import gradio as gr
import torch
from transformers import DistilBertTokenizer, DistilBertForSequenceClassification

# Carregar o modelo e o tokenizer treinados
model_name = "./fine_tuned_bert"
tokenizer = DistilBertTokenizer.from_pretrained(model_name)
model = DistilBertForSequenceClassification.from_pretrained(model_name)

# Mover o modelo para a GPU se disponível
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model.to(device)

# Função para prever o sentimento
def predict(text):
    inputs = tokenizer(text, return_tensors="pt", padding=True, truncation=True).to(device)
    with torch.no_grad():
        outputs = model(**inputs)
    prediction = torch.argmax(outputs.logits, dim=-1)
    return "Positivo" if prediction.item() == 1 else "Negativo"

# Criar interface Gradio
iface = gr.Interface(
    fn=predict,
    inputs=gr.Textbox(label="Digite um comentário"),
    outputs=gr.Textbox(label="Classificação"),
    title="Análise de Sentimento IMDb com BERT",
    description="Modelo de análise de sentimentos treinado com BERT para classificar comentários como positivos ou negativos."
)

# Iniciar a aplicação
iface.launch()