Course stuff regarding tokenization

.vscode/launch.json

@@ -0,0 +1,138 @@
+{
+  "version": "0.2.0",
+  "configurations": [
+    {
+      "name": "Debug: Base Train (minimal)",
+      "type": "debugpy",
+      "request": "launch",
+      "module": "scripts.base_train",
+      "args": [
+        "--depth=4",
+        "--max-seq-len=512",
+        "--device-batch-size=1",
+        "--total-batch-size=512",
+        "--num-iterations=20",
+        "--eval-every=-1",
+        "--core-metric-every=-1",
+        "--save-every=-1"
+      ],
+      "cwd": "${workspaceFolder}",
+      "env": {
+        "NANOCHAT_BASE_DIR": "${workspaceFolder}/.nanochat_debug"
+      },
+      "justMyCode": false,
+      "console": "integratedTerminal"
+    },
+    {
+      "name": "Debug: Mid Train (minimal)",
+      "type": "debugpy",
+      "request": "launch",
+      "module": "scripts.mid_train",
+      "args": [
+        "--device-batch-size=2",
+        "--max-seq-len=512",
+        "--total-batch-size=1024",
+        "--num-iterations=10",
+        "--eval-every=-1"
+      ],
+      "cwd": "${workspaceFolder}",
+      "env": {
+        "NANOCHAT_BASE_DIR": "${workspaceFolder}/.nanochat_debug"
+      },
+      "justMyCode": false,
+      "console": "integratedTerminal"
+    },
+    {
+      "name": "Debug: Chat SFT (minimal)",
+      "type": "debugpy",
+      "request": "launch",
+      "module": "scripts.chat_sft",
+      "args": [
+        "--source=mid",
+        "--device-batch-size=2",
+        "--num-iterations=5",
+        "--eval-every=-1",
+        "--eval-metrics-every=-1"
+      ],
+      "cwd": "${workspaceFolder}",
+      "env": {
+        "NANOCHAT_BASE_DIR": "${workspaceFolder}/.nanochat_debug"
+      },
+      "justMyCode": false,
+      "console": "integratedTerminal"
+    },
+    {
+      "name": "Debug: Tokenizer Train (minimal)",
+      "type": "debugpy",
+      "request": "launch",
+      "module": "scripts.tok_train",
+      "args": [
+        "--max-chars=500000",
+        "--vocab-size=4096",
+        "--doc-cap=5000"
+      ],
+      "cwd": "${workspaceFolder}",
+      "env": {
+        "NANOCHAT_BASE_DIR": "${workspaceFolder}/.nanochat_debug"
+      },
+      "justMyCode": false,
+      "console": "integratedTerminal"
+    },
+    {
+      "name": "Debug: Base Eval",
+      "type": "debugpy",
+      "request": "launch",
+      "module": "scripts.base_eval",
+      "args": [
+        "--max-per-task=10"
+      ],
+      "cwd": "${workspaceFolder}",
+      "env": {
+        "NANOCHAT_BASE_DIR": "${workspaceFolder}/.nanochat_debug"
+      },
+      "justMyCode": false,
+      "console": "integratedTerminal"
+    },
+    {
+      "name": "Debug: Report Generate",
+      "type": "debugpy",
+      "request": "launch",
+      "module": "nanochat.report",
+      "args": ["generate"],
+      "cwd": "${workspaceFolder}",
+      "env": {
+        "NANOCHAT_BASE_DIR": "${workspaceFolder}/.nanochat_debug"
+      },
+      "justMyCode": false,
+      "console": "integratedTerminal"
+    },
+    {
+      "name": "Debug: Pytest (tests)",
+      "type": "debugpy",
+      "request": "launch",
+      "module": "pytest",
+      "args": [
+        "tests/",
+        "-v",
+        "-s",
+        "--tb=short"
+      ],
+      "cwd": "${workspaceFolder}",
+      "justMyCode": false,
+      "console": "integratedTerminal"
+    },
+    {
+      "name": "Debug: Current File",
+      "type": "debugpy",
+      "request": "launch",
+      "program": "${file}",
+      "cwd": "${workspaceFolder}",
+      "env": {
+        "PYTHONPATH": "${workspaceFolder}",
+        "NANOCHAT_BASE_DIR": "${workspaceFolder}/.nanochat_debug"
+      },
+      "justMyCode": false,
+      "console": "integratedTerminal"
+    }
+  ]
+}

course/README.md

@@ -0,0 +1,25 @@
+# vast.ai config
+[build-system]
+requires = ["setuptools>=80.0", "wheel"]
+build-backend = "setuptools.build_meta"
+
+[tool.setuptools.packages.find]
+include = ["nanochat*"]
+exclude = ["dev*", "runs*", "tests*"]
+
+# Getting started
+ssh -i /path/to/key.pem username@server_ip
+git clone https://github.com/karpathy/nanochat
+pip install .
+sudo apt install tmux
+
+# Set up debugger
+create a debug config file so I can run through the debugger
+
+# Weights and biases
+Set up wandb and use it
+--run="[WANDB-NAME]"
+
+#  Stop and Resume
+--model-tag=[TAG-NAME]
+--resume-from-step=[CHECKPOINT-STEP,e.g.5000]

nanochat/tokenizer.py

@@ -391,8 +391,7 @@ def get_tokenizer():
     from nanochat.common import get_base_dir
     base_dir = get_base_dir()
     tokenizer_dir = os.path.join(base_dir, "tokenizer")
-    # return HuggingFaceTokenizer.from_directory(tokenizer_dir)
-    return RustBPETokenizer.from_directory(tokenizer_dir)
+    return HuggingFaceTokenizer.from_directory(tokenizer_dir)
 
 def get_token_bytes(device="cpu"):
     import torch

scripts/tok_compare.py

@@ -0,0 +1,111 @@
+"""
+Compare tokenization outputs between two tokenizer backends/directories.
+"""
+
+import argparse
+import os
+
+from nanochat.common import get_base_dir
+from nanochat.tokenizer import HuggingFaceTokenizer, RustBPETokenizer
+
+
+def load_tokenizer(backend: str, tokenizer_dir: str):
+    if backend == "huggingface":
+        expected = os.path.join(tokenizer_dir, "tokenizer.json")
+        if not os.path.exists(expected):
+            raise FileNotFoundError(
+                f"Missing HuggingFace tokenizer artifact: {expected}\n"
+                "Train one with: uv run python -m scripts.tok_train --tokenizer-backend huggingface\n"
+                "Or pass --tokenizer-dir-a/--tokenizer-dir-b to a directory that contains tokenizer.json."
+            )
+        return HuggingFaceTokenizer.from_directory(tokenizer_dir)
+    if backend == "rustbpe":
+        expected = os.path.join(tokenizer_dir, "tokenizer.pkl")
+        if not os.path.exists(expected):
+            raise FileNotFoundError(
+                f"Missing rustbpe tokenizer artifact: {expected}\n"
+                "Train one with: uv run python -m scripts.tok_train --tokenizer-backend rustbpe\n"
+                "Or pass --tokenizer-dir-a/--tokenizer-dir-b to a directory that contains tokenizer.pkl."
+            )
+        return RustBPETokenizer.from_directory(tokenizer_dir)
+    raise ValueError(f"Unknown backend: {backend}")
+
+
+def first_diff_index(a, b):
+    n = min(len(a), len(b))
+    for i in range(n):
+        if a[i] != b[i]:
+            return i
+    if len(a) != len(b):
+        return n
+    return -1
+
+
+def token_preview(tokenizer, token_ids, max_items=10):
+    preview_ids = token_ids[:max_items]
+    preview_tokens = [repr(tokenizer.decode([tid])) for tid in preview_ids]
+    return ", ".join(preview_tokens)
+
+
+def main():
+    parser = argparse.ArgumentParser(description="Compare token IDs between two tokenizers")
+    parser.add_argument("--backend-a", choices=["huggingface", "rustbpe"], required=True)
+    parser.add_argument("--backend-b", choices=["huggingface", "rustbpe"], required=True)
+    parser.add_argument("--tokenizer-dir-a", type=str, default=None, help="Directory for tokenizer A")
+    parser.add_argument("--tokenizer-dir-b", type=str, default=None, help="Directory for tokenizer B")
+    parser.add_argument("--text", action="append", default=[], help="Text to compare (repeatable)")
+    args = parser.parse_args()
+
+    base_dir = get_base_dir()
+    tok_dir_a = args.tokenizer_dir_a or os.path.join(base_dir, "tokenizer")
+    tok_dir_b = args.tokenizer_dir_b or os.path.join(base_dir, "tokenizer")
+
+    tokenizer_a = load_tokenizer(args.backend_a, tok_dir_a)
+    tokenizer_b = load_tokenizer(args.backend_b, tok_dir_b)
+
+    texts = args.text or [
+        "Hello world! 12345",
+        "Numbers: 1 23 456 7890",
+        "Unicode: 你好世界 🌍",
+        "A quick brown fox jumps over the lazy dog.",
+    ]
+
+    all_equal = True
+    print(f"A: {args.backend_a} @ {tok_dir_a}")
+    print(f"B: {args.backend_b} @ {tok_dir_b}")
+    print(f"Comparing {len(texts)} text sample(s)...")
+    print("-" * 80)
+
+    for i, text in enumerate(texts, start=1):
+        ids_a = tokenizer_a.encode(text)
+        ids_b = tokenizer_b.encode(text)
+        diff_i = first_diff_index(ids_a, ids_b)
+        equal = diff_i == -1
+        all_equal = all_equal and equal
+
+        print(f"[{i}] equal={equal} len_a={len(ids_a)} len_b={len(ids_b)}")
+        if equal:
+            continue
+
+        print(f"    first_diff_index={diff_i}")
+        if diff_i < len(ids_a):
+            print(f"    a_id={ids_a[diff_i]} a_tok={repr(tokenizer_a.decode([ids_a[diff_i]]))}")
+        else:
+            print("    a_id=<eos>")
+        if diff_i < len(ids_b):
+            print(f"    b_id={ids_b[diff_i]} b_tok={repr(tokenizer_b.decode([ids_b[diff_i]]))}")
+        else:
+            print("    b_id=<eos>")
+        print(f"    a_preview={token_preview(tokenizer_a, ids_a)}")
+        print(f"    b_preview={token_preview(tokenizer_b, ids_b)}")
+
+    print("-" * 80)
+    if all_equal:
+        print("All compared token ID sequences are identical.")
+    else:
+        print("Differences found.")
+        raise SystemExit(1)
+
+
+if __name__ == "__main__":
+    main()

scripts/tok_eval.py

@@ -2,9 +2,25 @@
 Evaluate compression ratio of the tokenizer.
 """
 
-from nanochat.tokenizer import get_tokenizer, RustBPETokenizer
+import argparse
+import os
+
+from nanochat.tokenizer import HuggingFaceTokenizer, RustBPETokenizer
+from nanochat.common import get_base_dir
 from nanochat.dataset import parquets_iter_batched
 
+parser = argparse.ArgumentParser(description='Evaluate tokenizer compression')
+parser.add_argument('--tokenizer-backend', type=str, default='huggingface', choices=['huggingface', 'rustbpe'],
+                    help='Tokenizer backend used for "ours" tokenizer')
+args = parser.parse_args()
+
+def load_ours_tokenizer():
+    base_dir = get_base_dir()
+    tokenizer_dir = os.path.join(base_dir, "tokenizer")
+    if args.tokenizer_backend == "huggingface":
+        return HuggingFaceTokenizer.from_directory(tokenizer_dir)
+    return RustBPETokenizer.from_directory(tokenizer_dir)
+
 # Random text I got from a random website this morning
 news_text = r"""
 (Washington, D.C., July 9, 2025)- Yesterday, Mexico’s National Service of Agro-Alimentary Health, Safety, and Quality (SENASICA) reported a new case of New World Screwworm (NWS) in Ixhuatlan de Madero, Veracruz in Mexico, which is approximately 160 miles northward of the current sterile fly dispersal grid, on the eastern side of the country and 370 miles south of the U.S./Mexico border. This new northward detection comes approximately two months after northern detections were reported in Oaxaca and Veracruz, less than 700 miles away from the U.S. border, which triggered the closure of our ports to Mexican cattle, bison, and horses on May 11, 2025.
@@ -171,7 +187,7 @@ for tokenizer_name in ["gpt2", "gpt4", "ours"]:
     elif tokenizer_name == "gpt4":
         tokenizer = RustBPETokenizer.from_pretrained("cl100k_base") # gpt-4 base model tokenizer
     else:
-        tokenizer = get_tokenizer()
+        tokenizer = load_ours_tokenizer()
 
     vocab_sizes[tokenizer_name] = tokenizer.get_vocab_size()
     tokenizer_results[tokenizer_name] = {}

scripts/tok_train.py

@@ -6,7 +6,7 @@ import os
 import time
 import argparse
 import torch
-from nanochat.tokenizer import RustBPETokenizer
+from nanochat.tokenizer import HuggingFaceTokenizer, RustBPETokenizer
 from nanochat.common import get_base_dir
 from nanochat.dataset import parquets_iter_batched
 
@@ -17,10 +17,13 @@ parser = argparse.ArgumentParser(description='Train a BPE tokenizer')
 parser.add_argument('--max-chars', type=int, default=2_000_000_000, help='Maximum characters to train on (default: 10B)')
 parser.add_argument('--doc-cap', type=int, default=10_000, help='Maximum characters per document (default: 10,000)')
 parser.add_argument('--vocab-size', type=int, default=32768, help='Vocabulary size (default: 32768 = 2^15)')
+parser.add_argument('--tokenizer-backend', type=str, default='huggingface', choices=['huggingface', 'rustbpe'],
+                    help='Tokenizer backend to train and save')
 args = parser.parse_args()
 print(f"max_chars: {args.max_chars:,}")
 print(f"doc_cap: {args.doc_cap:,}")
 print(f"vocab_size: {args.vocab_size:,}")
+print(f"tokenizer_backend: {args.tokenizer_backend}")
 
 # -----------------------------------------------------------------------------
 # Text iterator
@@ -46,7 +49,10 @@ text_iter = text_iterator()
 # -----------------------------------------------------------------------------
 # Train the tokenizer
 t0 = time.time()
-tokenizer = RustBPETokenizer.train_from_iterator(text_iter, args.vocab_size)
+if args.tokenizer_backend == "huggingface":
+    tokenizer = HuggingFaceTokenizer.train_from_iterator(text_iter, args.vocab_size)
+else:
+    tokenizer = RustBPETokenizer.train_from_iterator(text_iter, args.vocab_size)
 t1 = time.time()
 train_time = t1 - t0
 print(f"Training time: {train_time:.2f}s")