CODEX / app.py
Erinaldordso's picture
Update app.py
c0c0d8c verified
Raw History Blame Contribute Delete
935 Bytes
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()