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

MODEL = "meta-llama/Llama-3.2-1B-Instruct"

tokenizer = AutoTokenizer.from_pretrained(MODEL)

model = AutoModelForCausalLM.from_pretrained(
    MODEL,
    torch_dtype=torch.float16,
    low_cpu_mem_usage=True,
    device_map="auto"
)

model.config.use_cache = True

def gerar(prompt):
    prompt = "Responda direto e curto:\n" + prompt

    inputs = tokenizer(prompt, return_tensors="pt")

    with torch.no_grad():
        output = model.generate(
            **inputs,
            max_new_tokens=30,
            do_sample=False,
            temperature=0.7,
            top_p=0.9,
            repetition_penalty=1.1
        )

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

demo = gr.Interface(
    fn=gerar,
    inputs=gr.Textbox(lines=2),
    outputs="text",
    title="⚡ Llama 3.2 1B (Modo Bruto)"
)

demo.launch()