Retrieve dataset from dataset instead of text file

bigram.py

@@ -3,6 +3,7 @@ import torch
 import torch.nn as nn
 from torch.nn import functional as F
 import numpy as np
+from datasets import load_dataset
 
 from encoder import encode, decode
 from self_attention import Head, MultiHead
@@ -13,13 +14,12 @@ class Batcher():
         self.device = device
         self.batch_size = batch_size
         self.block_size = block_size
-        with open('input.txt', 'r', encoding='utf-8') as f:
-            text = f.read()
-            my_tensors = torch.tensor(encode(text), dtype=torch.long)
-            n = int(0.9*len(my_tensors))
-            self.train_data = my_tensors[:n]
-            self.val_data = my_tensors[n:]
-            self.vocab = set(text)
+        from dataset import make_dataset
+        train_data = make_dataset('train')
+        val_data = make_dataset('validation')
+        self.train_data = torch.tensor(encode(train_data), dtype=torch.long)
+        self.val_data = torch.tensor(encode(val_data), dtype=torch.long)
+        self.vocab = set(train_data + val_data)
 
     def get_batch(self, split: str = 'val'):
         data = self.train_data if split == 'train' else self.val_data

dataset.py

@@ -0,0 +1,19 @@
+from typing import Literal
+from datasets import load_dataset, Dataset
+
+_datasets = {
+    '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: Dataset = load_dataset(
+            "karpathy/tiny_shakespeare", split=split, trust_remote_code=True)
+        out = str(list(ds)[0]['text'])
+        _datasets[split] = out
+    return _datasets[split]

encoder.py

@@ -1,5 +1,6 @@
-with open('input.txt', 'r', encoding='utf-8') as f:
-    text = f.read()
+from dataset import make_dataset
+
+text = make_dataset('train') + make_dataset('validation')
 
 chars = sorted(list(set(text)))
 stoi = {ch: i for i, ch in enumerate(chars)}

model.pth

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:9036021f6ede2817a9b030d5dd605d38cf89ee627129bc791f5e9cc0b948aae1
+oid sha256:0f623313ac43af74c994754923b93641afef4c026c03a09b28d0e06640875675
 size 139095034

pyproject.toml

@@ -10,8 +10,12 @@ package-mode = false
 python = "^3.10"
 torch = "^2.3.0"
 numpy = "^1.26.4"
+datasets = "^2.19.0"
 
 
+[tool.poetry.group.dev.dependencies]
+ipykernel = "^6.29.4"
+
 [build-system]
 requires = ["poetry-core"]
 build-backend = "poetry.core.masonry.api"