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

MODEL_NAME = "MhoOmm/News_Classifier_Model"

tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
model = AutoModelForSequenceClassification.from_pretrained(MODEL_NAME)

def classify(text):

    inputs = tokenizer(
        text,
        return_tensors="pt",
        truncation=True,
        padding=True
    )

    with torch.no_grad():
        outputs = model(**inputs)

    probs = torch.softmax(outputs.logits, dim=-1)[0]

    return {
        model.config.id2label[i]: float(probs[i])
        for i in range(len(probs))
    }

demo = gr.Interface(
    fn=classify,
    inputs=gr.Textbox(lines=4),
    outputs="label",
    title="News Classification"
)

demo.launch()