Add extra typing to make pyright shut up

.gitignore

@@ -2,4 +2,5 @@
 __pycache__
 .git
 .wolf*
-.env
+.env
+pyrightconfig.json

dataset.py

@@ -1,7 +1,9 @@
-from typing import Literal
-from datasets import load_dataset, Dataset
+from typing import Literal, Union
+from datasets import load_dataset, DatasetDict
 
-_datasets = {
+DatasetType = Union[None, str]
+
+_datasets: dict[str, DatasetType] = {
     'train': None,
     'validation': None,
     'test': None,
@@ -12,8 +14,8 @@ _datasets = {
 
 def make_dataset(split: Literal['train', 'validation', 'test'] = 'train'):
     if _datasets[split] is None:
-        ds: Dataset = load_dataset(
+        ds = load_dataset(
             "karpathy/tiny_shakespeare", split=split, trust_remote_code=True)
         out = str(list(ds)[0]['text'])
         _datasets[split] = out
-    return _datasets[split]
+    return str(_datasets[split])