import re
import torch
import gradio as gr
from transformers import AutoModelForCausalLM, AutoTokenizer

MODEL_ID = "UniversalComputingResearch/Atom2.7m"

tokenizer = AutoTokenizer.from_pretrained(
    MODEL_ID,
    trust_remote_code=True,
)

model = AutoModelForCausalLM.from_pretrained(
    MODEL_ID,
    trust_remote_code=True,
).eval()

device = "cuda" if torch.cuda.is_available() else "cpu"
model.to(device)

print("Tokenizer class:", type(tokenizer))
print("Model class:", type(model))
print("Device:", device)


def decode_output(output_ids):
    """
    First try the official tokenizer.decode path.
    If the Space displays raw internal byte-level/LSD tokens, repair display only.
    """
    text = tokenizer.decode(output_ids, skip_special_tokens=True)

    # If normal decode worked, return it.
    if "Ġ" not in text and not re.search(r"\d\s+\d", text):
        return text

    # Fallback renderer for Atom's internal digit-token representation.
    # This does not change model behavior; it only fixes display if the Space
    # is not applying the custom decode path correctly.
    ids = output_ids.tolist() if hasattr(output_ids, "tolist") else list(output_ids)
    ids = [int(x) for x in ids]

    digit_id_to_char = {}
    for digit in "0123456789":
        token_id = tokenizer.convert_tokens_to_ids(digit)
        if token_id is not None and token_id != tokenizer.unk_token_id:
            digit_id_to_char[int(token_id)] = digit

    special_ids = set(tokenizer.all_special_ids)

    pieces = []
    text_buffer = []
    digit_buffer = []

    def flush_text():
        nonlocal text_buffer
        if text_buffer:
            part = tokenizer.backend_tokenizer.decode(
                text_buffer,
                skip_special_tokens=True,
            )
            part = part.replace("Ġ", " ")
            pieces.append(part)
            text_buffer = []

    def flush_digits():
        nonlocal digit_buffer
        if digit_buffer:
            pieces.extend(reversed(digit_buffer))
            digit_buffer = []

    for token_id in ids:
        if token_id in special_ids:
            continue

        if token_id in digit_id_to_char:
            flush_text()
            digit_buffer.append(digit_id_to_char[token_id])
        else:
            flush_digits()
            text_buffer.append(token_id)

    flush_text()
    flush_digits()

    return "".join(pieces).strip()


def generate(prompt, max_new_tokens, temperature, do_sample):
    if not prompt.strip():
        return "Enter a prompt first."

    inputs = tokenizer(
        prompt,
        return_tensors="pt",
        add_special_tokens=False,
    )

    inputs = {k: v.to(device) for k, v in inputs.items()}

    generation_kwargs = {
        **inputs,
        "max_new_tokens": int(max_new_tokens),
        "do_sample": bool(do_sample),
    }

    if do_sample:
        generation_kwargs["temperature"] = float(temperature)

    with torch.no_grad():
        output_ids = model.generate(**generation_kwargs)

    return decode_output(output_ids[0])


examples = [
    ["12 + 34 =", 3, 1.0, False],
    ["7 + 8 =", 3, 1.0, False],
    ["25 - 9 =", 3, 1.0, False],
    ["5 * 11 =", 3, 1.0, False],
    ["The capital of France is", 12, 0.8, True],
]

description = """
Atom2.7m is a tiny causal language model for text continuation, with arithmetic-aware handling for numeric spans.

It is not an instruction-tuned chatbot. It works best with short continuation prompts such as `12 + 34 =`.
For arithmetic, greedy decoding with a small number of new tokens usually works best.
"""

demo = gr.Interface(
    fn=generate,
    inputs=[
        gr.Textbox(
            label="Prompt",
            value="12 + 34 =",
            lines=3,
        ),
        gr.Slider(
            minimum=1,
            maximum=32,
            value=3,
            step=1,
            label="Max new tokens",
        ),
        gr.Slider(
            minimum=0.1,
            maximum=2.0,
            value=1.0,
            step=0.1,
            label="Temperature",
        ),
        gr.Checkbox(
            value=False,
            label="Sample instead of greedy decoding",
        ),
    ],
    outputs=gr.Textbox(
        label="Model output",
        lines=6,
    ),
    title="Atom2.7m Arithmetic Demo",
    description=description,
    examples=examples,
    allow_flagging="never",
)

if __name__ == "__main__":
    demo.launch()