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)


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,
    ).to(device)

    with torch.no_grad():
        output_ids = model.generate(
            **inputs,
            max_new_tokens=int(max_new_tokens),
            do_sample=bool(do_sample),
            temperature=float(temperature) if do_sample else None,
            pad_token_id=tokenizer.eos_token_id,
        )

    return tokenizer.decode(output_ids[0], skip_special_tokens=True)


examples = [
    ["12 + 34 =", 4, 1.0, False],
    ["7 + 8 =", 3, 1.0, False],
    ["25 - 9 =", 4, 1.0, False],
    ["3 * 6 =", 4, 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 =`.
"""

demo = gr.Interface(
    fn=generate,
    inputs=[
        gr.Textbox(
            label="Prompt",
            value="12 + 34 =",
            lines=3,
        ),
        gr.Slider(
            minimum=1,
            maximum=64,
            value=8,
            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()