diff --git a/.vscode/launch.json b/.vscode/launch.json new file mode 100644 index 00000000..2cd92153 --- /dev/null +++ b/.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" + } + ] +} diff --git a/course/README.md b/course/README.md new file mode 100644 index 00000000..446d3c74 --- /dev/null +++ b/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] \ No newline at end of file diff --git a/nanochat/tokenizer.py b/nanochat/tokenizer.py index a2146c2e..edc36bc3 100644 --- a/nanochat/tokenizer.py +++ b/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 diff --git a/scripts/tok_compare.py b/scripts/tok_compare.py new file mode 100644 index 00000000..dc62d2a8 --- /dev/null +++ b/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=") + 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=") + 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() diff --git a/scripts/tok_eval.py b/scripts/tok_eval.py index 9233d717..bf912c89 100644 --- a/scripts/tok_eval.py +++ b/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] = {} diff --git a/scripts/tok_train.py b/scripts/tok_train.py index 480e0e16..af9f0834 100644 --- a/scripts/tok_train.py +++ b/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")