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

MODEL_NAME = "TinyLlama/TinyLlama-1.1B-Chat-v1.0"

# Load tokenizer and model
tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
if tokenizer.pad_token is None:
    tokenizer.pad_token = tokenizer.eos_token

model = AutoModelForCausalLM.from_pretrained(
    MODEL_NAME,
    torch_dtype=torch.bfloat16 if torch.cuda.is_available() else torch.float32
)
if torch.cuda.is_available():
    model.to("cuda")
model.eval()

def generate_text(prompt, max_new_tokens=100, temperature=0.7, top_k=50):
    if not prompt:
        return "Please enter a prompt."

    messages = [{"role": "user", "content": prompt}]
    encoded = tokenizer.apply_chat_template(
        messages,
        add_generation_prompt=True,
        return_tensors="pt",
        padding=True,
        return_attention_mask=True,
    )

    input_ids = encoded["input_ids"]
    attention_mask = encoded["attention_mask"]

    if torch.cuda.is_available():
        input_ids = input_ids.to("cuda")
        attention_mask = attention_mask.to("cuda")

    output_ids = model.generate(
        input_ids,
        attention_mask=attention_mask,
        max_new_tokens=max_new_tokens,
        do_sample=True,
        temperature=temperature,
        top_k=top_k,
        pad_token_id=tokenizer.eos_token_id
    )

    response = tokenizer.decode(output_ids[0][input_ids.shape[-1]:], skip_special_tokens=True)
    return response

# Gradio interface
demo = gr.Interface(
    fn=generate_text,
    inputs=[
        gr.Textbox(label="Prompt"),
        gr.Slider(minimum=10, maximum=500, value=100, label="Max New Tokens"),
        gr.Slider(minimum=0.1, maximum=1.0, value=0.7, step=0.05, label="Temperature"),
        gr.Slider(minimum=0, maximum=100, value=50, step=1, label="Top K")
    ],
    outputs=gr.Textbox(label="Generated Text"),
    title="TinyLlama Gradio API",
    description="Use this via UI or API via `/run/predict`"
)

if __name__ == "__main__":
    demo.launch(server_name="0.0.0.0", server_port=7860)