import gradio as gr
from transformers import DistilBertTokenizer, DistilBertForSequenceClassification
from huggingface_hub import hf_hub_download
import numpy as np

tokenizer = DistilBertTokenizer.from_pretrained('distilbert-base-uncased')
model = DistilBertForSequenceClassification.from_pretrained('jdmartinev/imdbreviews_classification_distilbert_v02')
id2label = {0: "negative", 1: "positive"}

def classify_text(input_text,id2label):
    inputs = tokenizer(input_text, return_tensors="pt")
    outputs = model(**inputs)
    pred = np.argmax(outputs)
    return(pred)

demo = gr.Interface(classify_text, "text", "text")
demo.launch()