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

model_name = "Bittar/outputs"
model = AutoModelForSequenceClassification.from_pretrained(model_name)
tokenizer = AutoTokenizer.from_pretrained(model_name)

mapping = {
    0: 'negative',
    1: 'positive'
}

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

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

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

iface = gr.Interface(
    fn=predict,
    inputs="text",
    outputs="text",
    layout="vertical",
    title="Movie feelings classifier",
    description="do u feel positive or negative about a movie? Write your review down and find out! (only works in english)"
)

iface.launch(share=True)