from tokenizer import (
    build_xonelm_tokenizer,
    MultiTurnConversationFormatter,
    SpecialTokenConfig,
)

def run_tokenizer_demo():
    sample_corpus = [
        "Once upon a time, Lily found a golden key in the garden.",
        "Timmy and his dog Max played with a red ball.",
        "def solve_quadratic(a, b, c): return (-b + (b**2 - 4*a*c)**0.5) / (2*a)",
        "\\int_{0}^{\\infty} e^{-x^2} dx = \\frac{\\sqrt{\\pi}}{2}",
        "The system latency is <= 10ms with async/await workers.",
    ]

    tokenizer = build_xonelm_tokenizer(corpus=sample_corpus, vocab_size=1000)
    print("Tokenizer Vocab Size:", len(tokenizer))

    text_to_encode = "Lily solved \\alpha + \\beta == 42 async await."
    encoded = tokenizer.encode(text_to_encode)
    token_ids = encoded.ids if hasattr(encoded, "ids") else encoded["input_ids"]
    decoded = tokenizer.decode(token_ids)

    print("Single Text Tokenization")
    print("Input Text :", text_to_encode)
    print("Token IDs  :", token_ids)
    print("Decoded    :", decoded)

    conversation = [
        {"role": "system", "content": "You are a helpful and wise AI assistant."},
        {"role": "user", "content": "Can you explain how it's work?"},
        {"role": "assistant", "content": "No! I can't. hehe"},
    ]

    cfg = SpecialTokenConfig(
        pad_token_id=0,
        bos_token_id=1,
        eos_token_id=2,
        unk_token_id=3,
        eod_token_id=4,
    )
    formatter = MultiTurnConversationFormatter(tokenizer, cfg)
    formatted = formatter.format_conversation(conversation)

    print("Multi-Turn ChatML Formatting")
    print("Input IDs Length :", len(formatted["input_ids"]))
    print("Labels Length    :", len(formatted["labels"]))
    print("Formatted Text   :\n" + tokenizer.decode(formatted["input_ids"]))

if __name__ == "__main__":
    run_tokenizer_demo()