"""Fix tokenizer v3: remove 32 BPE merges that produce tokens not in vocab (surgery orphans)."""
import json
from pathlib import Path

src = Path("tokenizer.json")
with open(src) as f:
    t = json.load(f)

v = t["model"]["vocab"]
merges = t["model"]["merges"]
print(f"Before: vocab={len(v)} merges={len(merges)}")

clean_merges = []
removed = []
for m in merges:
    if isinstance(m, list) and len(m) == 2:
        comb = m[0] + m[1]
        if comb in v:
            clean_merges.append(m)
        else:
            removed.append((m, comb))
    else:
        # Keep non-standard entries verbatim
        clean_merges.append(m)

t["model"]["merges"] = clean_merges
print(f"After: vocab={len(v)} merges={len(clean_merges)} removed={len(removed)}")
print("Removed examples:")
for m, comb in removed[:15]:
    print(f"  {m} -> {comb!r}")

with open(src, "w", encoding="utf-8") as f:
    json.dump(t, f, ensure_ascii=False)
print(f"Saved {src}")

# Test load
from tokenizers import Tokenizer
try:
    tk = Tokenizer.from_file(str(src))
    print(f"LOAD OK vocab_size={tk.get_vocab_size()}")
    # Test FR + ChatML encoding
    test = "Bonjour, je suis ARCHON. <|im_start|>system test<|im_end|>"
    ids = tk.encode(test).ids
    print(f"encode test: ids[:15]={ids[:15]} len={len(ids)}")
    print(f"decode: {tk.decode(ids)!r}")
except Exception as e:
    print(f"LOAD FAIL: {e}")