#!/usr/bin/env python3
"""
Patch llama.cpp tools/perplexity/perplexity.cpp so common_tokenize() parses special
tokens (<|im_start|>, <|im_end|>, <think>, </think>, <tool_call>, </tool_call>) on chat/agent traces.
"""
import sys
from pathlib import Path

PPL_CPP = Path(sys.argv[1] if len(sys.argv) > 1 else "/content/llama.cpp/tools/perplexity/perplexity.cpp")
txt = PPL_CPP.read_text(encoding="utf-8")
if "common_tokenize(ctx, params.prompt, true, true)" in txt:
    print("perplexity.cpp already patched with parse_special=true")
else:
    patched = txt.replace(
        "common_tokenize(ctx, params.prompt, true)",
        "common_tokenize(ctx, params.prompt, true, true)",
    )
    assert patched != txt, "Failed to find common_tokenize(ctx, params.prompt, true) in perplexity.cpp"
    PPL_CPP.write_text(patched, encoding="utf-8")
    print(f"Patched {PPL_CPP} with parse_special=true")