"""Apply the Olive Qwen3.5 tokenizer fix for onnxruntime-genai compatibility.

Qwen3.5's tokenizer.json uses Unicode property escapes (\\p{L}, \\p{N}) in its
Split pre-tokenizer, which onnxruntime-genai's C++ tokenizer (std::regex)
cannot parse. This keeps only the ByteLevel pre-tokenizer (with use_regex=True)
and retargets tokenizer_class to Qwen2Tokenizer, exactly as the official
microsoft/olive-recipes Qwen-Qwen3.5-4B example does in optimize.py.

Usage:
    python fix_tokenizer.py <model_dir> [<model_dir> ...]

This only affects the tokenizer used by onnxruntime-genai for chat/generation.
KnowLine scoring (score_knowline.py) tokenizes with the original HF tokenizer,
so prompts stay byte-identical to the PyTorch serving.
"""

import json
import sys
from pathlib import Path


def fix_tokenizer(model_dir: str) -> None:
    model_dir = Path(model_dir)
    tk_path = model_dir / "tokenizer.json"
    if tk_path.exists():
        tk = json.loads(tk_path.read_text(encoding="utf-8"))
        pt = tk.get("pre_tokenizer", {})
        if pt.get("type") == "Sequence":
            pt["pretokenizers"] = [s for s in pt["pretokenizers"] if s.get("type") == "ByteLevel"]
            for s in pt["pretokenizers"]:
                s["use_regex"] = True
        tk_path.write_text(json.dumps(tk, ensure_ascii=False), encoding="utf-8")
        print(f"  fixed {tk_path}")

    tc_path = model_dir / "tokenizer_config.json"
    if tc_path.exists():
        tc = json.loads(tc_path.read_text(encoding="utf-8"))
        tc["tokenizer_class"] = "Qwen2Tokenizer"
        tc_path.write_text(json.dumps(tc, indent=2, ensure_ascii=False), encoding="utf-8")
        print(f"  fixed {tc_path}")


if __name__ == "__main__":
    if len(sys.argv) < 2:
        print(__doc__)
        sys.exit(1)
    for d in sys.argv[1:]:
        print(f"Fixing tokenizer in {d}")
        fix_tokenizer(d)
    print("Done.")