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

MODEL_ID = "SupraLabs/Supra-1.5-50M-Instruct-exp"

print(f"Loading model: {MODEL_ID}")
tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
model = AutoModelForCausalLM.from_pretrained(MODEL_ID, dtype=torch.float32)
model.eval()
print("Model loaded.")


def generate_text(
    message,
    history,
    max_new_tokens=200,
    temperature=0.7,
    top_p=0.9,
    repetition_penalty=1.2,
):
    max_new_tokens = int(max_new_tokens or 200)
    temperature = float(temperature or 0.7)
    top_p = float(top_p or 0.9)
    repetition_penalty = float(repetition_penalty or 1.2)

    prompt = f"User: {message}\nAssistant:"
    inputs = tokenizer(prompt, return_tensors="pt")
    input_ids = inputs["input_ids"]

    with torch.no_grad():
        output = model.generate(
            input_ids,
            max_new_tokens=max_new_tokens,
            temperature=temperature if temperature > 0 else 1.0,
            top_p=top_p,
            repetition_penalty=repetition_penalty,
            do_sample=temperature > 0,
            pad_token_id=tokenizer.eos_token_id,
            eos_token_id=tokenizer.eos_token_id,
        )

    generated = output[0][input_ids.shape[-1]:]
    response = tokenizer.decode(generated, skip_special_tokens=True).strip()

    if "User:" in response:
        response = response.split("User:")[0].strip()

    return response


with gr.Blocks(title="Supra-1.5-50M Instruct") as demo:
    gr.Markdown(
        """
        # 🧠 Supra-1.5-50M-Instruct-exp
        **SupraLabs** · Experimental 50M-parameter instruction-tuned SLM · [Model Card](https://huggingface.co/SupraLabs/Supra-1.5-50M-Instruct-exp)

        > ⚠️ Experimental model — responses may be limited or inconsistent.
        """
    )

    gr.ChatInterface(
        fn=generate_text,
        additional_inputs=[
            gr.Slider(32, 512, value=200, step=8, label="Max New Tokens"),
            gr.Slider(0.0, 1.5, value=0.7, step=0.05, label="Temperature"),
            gr.Slider(0.5, 1.0, value=0.9, step=0.05, label="Top-p"),
            gr.Slider(1.0, 2.0, value=1.2, step=0.05, label="Repetition Penalty"),
        ],
        additional_inputs_accordion=gr.Accordion("⚙️ Generation Settings", open=False),
        examples=[
            ["What is a language model?"],
            ["Write a short poem about the stars."],
            ["Explain gravity in simple terms."],
            ["What is the capital of France?"],
        ],
    )

    gr.Markdown("---\nBuilt with ❤️ by [SupraLabs](https://huggingface.co/SupraLabs) · Apache 2.0 License")

demo.launch(theme=gr.themes.Soft())