File size: 637 Bytes
103a097
4d54af4
103a097
4d54af4
 
103a097
4d54af4
 
 
 
 
 
 
103a097
4d54af4
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
import gradio as gr
from transformers import BertTokenizer, BertForSequenceClassification

tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')
model = BertForSequenceClassification.from_pretrained('bert-base-uncased', num_labels=8)

def classify(text):
    inputs = tokenizer(text, return_tensors='pt')
    outputs = model(inputs['input_ids'], attention_mask=inputs['attention_mask'])
    logits = outputs[0].logits
    probabilities = torch.softmax(logits, dim=-1).tolist()
    predicted_label = np.argmax(probabilities)
    return {'Label': predicted_label}

gr.Interface(fn=classify, inputs="text", outputs="json").launch()