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

# Load model and tokenizer
model_name = "unsloth/Llama-3.2-1B-Instruct"  # Use the non-quantized version

tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(
    model_name,
    torch_dtype=torch.float32,
    low_cpu_mem_usage=True,
    device_map="cpu"
)

def generate_text(prompt, max_new_tokens, temperature):
    inputs = tokenizer(prompt, return_tensors="pt")
    
    with torch.no_grad():
        outputs = model.generate(
            **inputs,
            max_new_tokens=int(max_new_tokens),
            temperature=temperature,
            num_return_sequences=1,
            do_sample=True,
        )
    
    return tokenizer.decode(outputs[0], skip_special_tokens=True)

# Define the Gradio interface
iface = gr.Interface(
    fn=generate_text,
    inputs=[
        gr.Textbox(lines=5, label="Enter your prompt"),
        gr.Slider(50, 500, value=200, step=1, label="Maximum New Tokens"),
        gr.Slider(0.1, 2.0, value=0.7, step=0.1, label="Temperature")
    ],
    outputs=gr.Textbox(label="Generated Text"),
    title="Text Generation with Llama-3.2-1B-Instruct",
    description="Enter a prompt to generate text using the Llama-3.2-1B-Instruct model."
)

# Launch the interface
iface.launch()