import gradio as gr
from transformers import AutoModelForSequenceClassification, AutoTokenizer

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


def classify_text(text):
    inputs = tokenizer(text, return_tensors="pt")
    outputs = model(**inputs)
    predictions = outputs.logits.argmax(dim=1)

    return f"Predicted class: {predictions.item()}"


iface = gr.Interface(fn=classify_text, inputs="text", outputs="text")

iface.launch()