MhoOmm's picture
Update app.py
02c07ba verified
Raw
History Blame Contribute Delete
804 Bytes
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()