import torch
from transformers import pipeline
import gradio as gr

# Define the model ID
model_id = "unsloth/Llama-3.2-1B"

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

# Define custom labels for classification
pipe.model.config.id2label = {0: 'saudação', 1: 'fim de conversa', 2: 'outro'}
pipe.model.config.label2id = {'saudação': 0, 'fim de conversa': 1, 'outro': 2}

# Function to classify input text
def classify_text(text):
    result = pipe(text)
    return result[0]['label'] + ": " + str(result[0]['score'])

# Create Gradio interface
iface = gr.Interface(
    fn=classify_text,  # Function to be called
    inputs=gr.Textbox(label="Digite o texto"),  # Textbox input for user
    outputs=gr.Label(label="Classificação"),  # Output label showing classification
    title="Classificação Textual",  # Title of the app
    description="Classifique seu texto como 'saudação', 'fim de conversa', ou 'outro'."  # Description of the app
)

# Launch the Gradio app
iface.launch()