import gradio as gr
import spaces
import torch
from transformers import pipeline

MODELS = {
    "Qwen2.5-0.5B": "Qwen/Qwen2.5-0.5B-Instruct",
    "Qwen2.5-1B": "Qwen/Qwen2.5-1B-Instruct",
    "Llama-3.2-1B": "meta-llama/Llama-3.2-1B-Instruct",
    "Gemma-3-1B": "google/gemma-3-1b-it",
}

loaded = {}

def get_model(name):
    if name not in loaded:
        loaded[name] = pipeline(
            "text-generation",
            model=MODELS[name],
            torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
            device_map="auto" if torch.cuda.is_available() else None,
        )
    return loaded[name]


@spaces.GPU
def generate(model, prompt, max_tokens, temperature):
    pipe = get_model(model)

    result = pipe(
        prompt,
        max_new_tokens=max_tokens,
        temperature=temperature,
        do_sample=temperature > 0,
        return_full_text=False,
    )

    return result[0]["generated_text"]


demo = gr.Interface(
    fn=generate,
    inputs=[
        gr.Dropdown(
            choices=list(MODELS.keys()),
            value="Qwen2.5-0.5B",
            label="Model",
        ),
        gr.Textbox(lines=8, label="Prompt"),
        gr.Slider(1, 512, value=128, step=1, label="Max Tokens"),
        gr.Slider(0.0, 2.0, value=0.7, step=0.1, label="Temperature"),
    ],
    outputs=gr.Textbox(lines=12, label="Response"),
    title="Micro Tier LLM API",
)

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