from typing import Literal, Union
from datasets import load_dataset, DatasetDict

DatasetType = Union[None, str]

_datasets: dict[str, DatasetType] = {
    'train': None,
    'validation': None,
    'test': None,
}

# Lazy load the dataset


def make_dataset(split: Literal['train', 'validation', 'test'] = 'train'):
    if _datasets[split] is None:
        ds = load_dataset(
            "karpathy/tiny_shakespeare", split=split, trust_remote_code=True)
        out = str(list(ds)[0]['text'])
        _datasets[split] = out
    return str(_datasets[split])