feat(dataset): add dataset and data loader for tokenized text processing

llm/dataset.py

@@ -0,0 +1,49 @@
+import tiktoken
+import torch
+from tiktoken import Encoding
+from torch.utils.data import DataLoader, Dataset
+
+
+def create_data_loader(
+    text: str,
+    max_length: int,
+    stride: int,
+    batch_size: int,
+    shuffle: bool,
+    drop_last: bool = True,
+    num_workers: int = 0,
+) -> DataLoader:
+    tokenizer = tiktoken.get_encoding("gpt2")
+    dataset = GptDatasetV1(
+        text=text,
+        tokenizer=tokenizer,
+        max_length=max_length,
+        stride=stride,
+    )
+    return DataLoader(
+        dataset=dataset,
+        batch_size=batch_size,
+        shuffle=shuffle,
+        drop_last=drop_last,
+        num_workers=num_workers,
+    )
+
+
+class GptDatasetV1(Dataset):
+    def __init__(self, text: str, tokenizer: Encoding, max_length: int, stride: int):
+        self.input_ids = []
+        self.target_ids = []
+
+        token_ids = tokenizer.encode(text)
+
+        for i in range(0, len(token_ids) - max_length, stride):
+            input_chunk = token_ids[i : i + max_length]
+            target_chunk = token_ids[i + 1 : i + max_length + 1]
+            self.input_ids.append(torch.tensor(input_chunk))
+            self.target_ids.append(torch.tensor(target_chunk))
+
+    def __len__(self) -> int:
+        return len(self.input_ids)
+
+    def __getitem__(self, index: int) -> tuple[int, int]:
+        return self.input_ids[index], self.target_ids[index]

pyproject.toml

@@ -3,6 +3,10 @@ name = "largelanguagemodel"
 version = "0.1.0"
 description = "Experimenting with LLMs"
 requires-python = ">=3.13"
+dependencies = [
+    "tiktoken>=0.9.0",
+    "torch>=2.7.0",
+]
 
 [dependency-groups]
 dev = [

tests/test_data_loader.py

@@ -0,0 +1,24 @@
+from pathlib import Path
+
+import torch
+
+from llm.dataset import create_data_loader
+
+
+def test_data_loader(resource: Path):
+    path = resource / "the-verdict.txt"
+    text = path.read_text(encoding="utf-8")
+
+    data_loader = create_data_loader(
+        text=text,
+        batch_size=1,
+        max_length=4,
+        stride=1,
+        shuffle=False,
+    )
+
+    data_iter = iter(data_loader)
+    inputs, targets = next(data_iter)
+
+    assert torch.equal(inputs, torch.tensor([[40, 367, 2885, 1464]]))
+    assert torch.equal(targets, torch.tensor([[367, 2885, 1464, 1807]]))