Upload src/tokenizer.py with huggingface_hub

src/tokenizer.py

@@ -0,0 +1,79 @@
+"""
+Milestone 1: Character-level tokenizer and data loading.
+
+Reads tiny Shakespeare, builds a vocab from all unique characters,
+provides encode/decode, and a get_batch() function for training.
+"""
+
+import torch
+
+# ── Config ────────────────────────────────────────────────────────────────────
+DATA_PATH  = "data/input.txt"
+BLOCK_SIZE = 256   # context length (tokens per sample)
+BATCH_SIZE = 64    # samples per batch
+if torch.cuda.is_available():
+    DEVICE = "cuda"
+elif torch.backends.mps.is_available():
+    DEVICE = "mps"
+else:
+    DEVICE = "cpu"
+
+# ── Load data ─────────────────────────────────────────────────────────────────
+with open(DATA_PATH, "r") as f:
+    text = f.read()
+
+# ── Build vocab ───────────────────────────────────────────────────────────────
+chars  = sorted(set(text))
+VOCAB_SIZE = len(chars)
+
+stoi = {ch: i for i, ch in enumerate(chars)}   # char -> int
+itos = {i: ch for i, ch in enumerate(chars)}   # int  -> char
+
+def encode(s: str) -> list[int]:
+    return [stoi[c] for c in s]
+
+def decode(ids: list[int]) -> str:
+    return "".join(itos[i] for i in ids)
+
+# ── Train / val split ─────────────────────────────────────────────────────────
+data   = torch.tensor(encode(text), dtype=torch.long)
+n      = int(0.9 * len(data))
+train_data = data[:n]
+val_data   = data[n:]
+
+# ── Batch sampler ─────────────────────────────────────────────────────────────
+def get_batch(split: str):
+    """Return a random batch of (x, y) pairs.
+
+    x: (BATCH_SIZE, BLOCK_SIZE) input token ids
+    y: (BATCH_SIZE, BLOCK_SIZE) target token ids (x shifted right by 1)
+    """
+    src = train_data if split == "train" else val_data
+    ix  = torch.randint(len(src) - BLOCK_SIZE, (BATCH_SIZE,))
+    x   = torch.stack([src[i     : i + BLOCK_SIZE    ] for i in ix])
+    y   = torch.stack([src[i + 1 : i + BLOCK_SIZE + 1] for i in ix])
+    return x.to(DEVICE), y.to(DEVICE)
+
+
+# ── Quick sanity check ────────────────────────────────────────────────────────
+if __name__ == "__main__":
+    print(f"Dataset length : {len(text):,} characters")
+    print(f"Vocab size     : {VOCAB_SIZE} unique chars")
+    print(f"Train tokens   : {len(train_data):,}")
+    print(f"Val tokens     : {len(val_data):,}")
+    print(f"Device         : {DEVICE}")
+    print()
+
+    sample = text[:100]
+    encoded = encode(sample)
+    decoded = decode(encoded)
+    print(f"Sample text    : {repr(sample[:40])}")
+    print(f"Encoded[:10]   : {encoded[:10]}")
+    print(f"Round-trip OK  : {decoded == sample}")
+    print()
+
+    x, y = get_batch("train")
+    print(f"Batch x shape  : {x.shape}  (on {x.device})")
+    print(f"Batch y shape  : {y.shape}  (on {y.device})")
+    print(f"x[0,:8]        : {x[0,:8].tolist()}")
+    print(f"y[0,:8]        : {y[0,:8].tolist()}  (x shifted by 1)")