Gilvan's picture
Update app.py
b3bcbc8 verified
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()