#!/usr/bin/env python3
"""
Patch the llama.cpp converter to handle Comma v0.1's tokenizer format.
This adds support for the specific BPE pre-tokenizer checksum.
"""

import sys
from pathlib import Path

converter_path = Path("llama.cpp/convert_hf_to_gguf.py")

if not converter_path.exists():
    print(f"ERROR: {converter_path} not found!")
    sys.exit(1)

print(f"Patching {converter_path}...")

# Read the file
with open(converter_path, 'r', encoding='utf-8', errors='ignore') as f:
    content = f.read()

# The checksum that's failing
failing_checksum = "bf66900d65fe80247e435184a4ac839c5c332657cf567e64b8ede5fbd63f5fd9"

# Find the section where checksums are checked and add our checksum
# We need to find the get_vocab_base_pre function and add a case for this checksum

patch_code = f'''
        # Patch for Comma v0.1 tokenizer (Llama 3 compatible)
        if chkhsh == "{failing_checksum}":
            # Comma v0.1 uses Llama 3 style BPE
            return gguf.MODEL_TENSOR.TOKEN_EMBD
'''

# Try to find a good insertion point
if 'def get_vocab_base_pre(' in content:
    print("Found get_vocab_base_pre function")

    # Look for where the NotImplementedError is raised
    error_line = 'raise NotImplementedError("BPE pre-tokenizer was not recognized'

    if error_line in content:
        # Insert our patch before the error
        content = content.replace(
            error_line,
            patch_code + "\n        " + error_line
        )
        print("[OK] Patch applied!")

        # Write back
        with open(converter_path, 'w', encoding='utf-8') as f:
            f.write(content)

        print(f"\n{converter_path} has been patched.")
        print("Now try running: python llama.cpp/convert_hf_to_gguf.py comma-v0.1-2t --outfile comma-v0.1-2t.gguf --outtype q4_K_M")
    else:
        print("ERROR: Could not find the error line to patch")
        sys.exit(1)
else:
    print("ERROR: Could not find get_vocab_base_pre function")
    sys.exit(1)