import torch
from transformers import pipeline
import gradio as gr

# Carregar o modelo
model_id = "meta-llama/Llama-3.2-1B-Instruct"

pipe = pipeline(
    "text-classification", 
    model=model_id, 
    torch_dtype=torch.bfloat16, 
    device_map="auto"
)

# Definir os rótulos das classes
pipe.model.config.id2label = {0: 'greeting', 1: 'farewell', 2: 'other'}

# Função para usar o modelo e retornar o resultado
def classify_text(text):
    result = pipe(text)
    return result[0]['label']

# Interface Gradio
iface = gr.Interface(
    fn=classify_text,                  # Função que será chamada
    inputs=gr.Textbox(label="Texto"),   # Caixa de texto para o usuário inserir a entrada
    outputs=gr.Label(label="Classificação"),  # Exibição da classificação
    title="Classificador de Texto",     # Título do app
    description="Este modelo classifica o texto em três categorias: 'greeting', 'farewell', e 'other'."
)

# Iniciar o app Gradio
iface.launch()