Support Laoyao BPE tokenizer in SFT heldout builder

dataset_building/README.md

@@ -56,8 +56,10 @@ dataset_building/generated/training_mix_v4_train1m_test2p8k_noupsample_nobbh_202
 ```
 
 Use `OUT=/path/to/output` and `TOKENIZER=/path/to/tokenizer` to place the
-validation set outside the repo or build it with a specific tokenizer. Generated
-outputs stay ignored by git.
+validation set outside the repo or build it with a specific tokenizer. For
+Laoyao custom BPE tokenizers, set `LAOYAO_MODEL_ROOT=/path/to/laoyao/repo` so
+the builder can import `laoyaomodel.tokenization.bpe`. Generated outputs stay
+ignored by git.
 
 ## SFT Builders
 

dataset_building/build_training_and_test_mix_v3.py

@@ -7,6 +7,7 @@ import json
 import os
 import random
 import re
+import sys
 from collections import Counter, defaultdict
 from pathlib import Path
 
@@ -145,15 +146,30 @@ def prompt_hash(user):
 
 
 def load_tokenizer(path):
+    path_obj = Path(path)
     try:
         return AutoTokenizer.from_pretrained(path, trust_remote_code=True)
     except Exception:
-        return Tokenizer.from_file(str(Path(path) / "tokenizer.json"))
+        tokenizer_path = path_obj if path_obj.is_file() else path_obj / "tokenizer.json"
+        try:
+            return Tokenizer.from_file(str(tokenizer_path))
+        except Exception:
+            payload = json.loads(tokenizer_path.read_text(encoding="utf-8"))
+            if payload.get("implementation") != "laoyaomodel.bytes_bpe":
+                raise
+            model_root = os.environ.get("LAOYAO_MODEL_ROOT")
+            if model_root:
+                sys.path.insert(0, model_root)
+            from laoyaomodel.tokenization.bpe import BPETokenizer
+
+            return BPETokenizer.from_tokenizer_json(tokenizer_path)
 
 
 def token_len(tok, text):
     if isinstance(tok, Tokenizer):
         return len(tok.encode(text, add_special_tokens=False).ids)
+    if tok.__class__.__name__ == "BPETokenizer":
+        return len(tok.encode(text))
     return len(tok(text, add_special_tokens=False)["input_ids"])