import gradio as gr
from transformers import AutoModelForSequenceClassification, AutoTokenizer
import torch

model_name = "Daniel-Sousa/outputs"
model = AutoModelForSequenceClassification.from_pretrained(model_name)
tokenizer = AutoTokenizer.from_pretrained(model_name)

mapping = {
    0: 'It is not a disaster',
    1: 'Disaster',
}

def predict(text):
    inputs = tokenizer(text, return_tensors="pt")

    outputs = model(**inputs)
    predictions = outputs.logits

    return mapping[round(predictions.item())]

iface = gr.Interface(
    fn=predict,
    inputs="text",
    outputs="text",
    layout="vertical",
    title="Classificador de Tweets",
    description="Esse modelo é capaz de realizar uma análise de frases em Inglês do Tweet e diser se ela está refenciado um desastre ou não"
)

iface.launch()