Use SmolLM tokenizer by default and validate vocab

train_2b.py

@@ -119,6 +119,21 @@ _LOCAL_SNAPSHOT_CACHE = {}
 _LOCAL_DATA_SUFFIXES = {".parquet", ".jsonl", ".json"}
 
 
+def _validate_tokenizer_compatibility(config, tokenizer_name: str) -> None:
+    from transformers import AutoTokenizer
+
+    tokenizer = AutoTokenizer.from_pretrained(tokenizer_name)
+    base_vocab_size = tokenizer.vocab_size
+    supported_base_vocab = config.vocab_size - 1
+    if base_vocab_size > supported_base_vocab:
+        raise SystemExit(
+            f"Tokenizer '{tokenizer_name}' has vocab_size={base_vocab_size}, "
+            f"but this preset only supports up to {supported_base_vocab} base tokens "
+            f"(config.vocab_size={config.vocab_size}, mask_token_id={config.mask_token_id}). "
+            "Use a compatible tokenizer such as 'HuggingFaceTB/SmolLM2-1.7B' or increase the model vocab."
+        )
+
+
 def _resolve_local_dataset_path(dataset_name):
     if os.path.isfile(dataset_name) or os.path.isdir(dataset_name):
         return dataset_name
@@ -788,7 +803,7 @@ def main():
     parser.add_argument("--batch_size", type=int, default=8)
     parser.add_argument("--device", type=str, default="cuda")
     parser.add_argument("--save_dir", type=str, default="results")
-    parser.add_argument("--tokenizer", type=str, default="gpt2")
+    parser.add_argument("--tokenizer", type=str, default="HuggingFaceTB/SmolLM2-1.7B")
     parser.add_argument("--data", type=str, default=None)
     parser.add_argument("--model-preset", choices=sorted(MODEL_PRESETS.keys()), default="dev",
                         help="Architecture preset to use")
@@ -819,6 +834,7 @@ def main():
     preset = MODEL_PRESETS[args.model_preset]
     config = PCSHOConfig(**preset["config"])
     msa_config = MSAConfig(**preset["msa"])
+    _validate_tokenizer_compatibility(config, args.tokenizer)
 
     save_root = args.save_dir if os.path.isabs(args.save_dir) else os.path.join(os.getcwd(), args.save_dir)
     free_bytes = shutil.disk_usage(os.path.dirname(save_root) or os.getcwd()).free