File size: 804 Bytes
02c07ba
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
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()