File size: 1,001 Bytes
f5459a1 cf0a4fc f5459a1 |
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 |
import gradio as gr
from transformers import DistilBertTokenizer, DistilBertForSequenceClassification
import torch
description = "Sentiment Analysis :) && :("
title = "SentBERT"
examples = [["That ice cream was really bad"], ["Great to meet you!"], ["Hey, there's a snake there"]]
class2interpret = {
0: 'Positive/Neutral',
1: 'Negative'
}
def classify(example):
tokenizer = DistilBertTokenizer.from_pretrained("distilbert-base-uncased")
model = DistilBertForSequenceClassification.from_pretrained("distilbert-base-uncased")
inputs = tokenizer(example, return_tensors="pt")
with torch.no_grad():
logits = model(**inputs).logits
probs = torch.nn.Softmax(dim=1)(logits).tolist()[0]
return {class2interpret[0]: probs[0], class2interpret[1]: probs[1]}, {class2interpret[0]: probs[0], class2interpret[1]: probs[1]}
interface = gr.Interface(fn=classify, inputs='text', outputs=['label', 'json'], examples=examples, description=description, title=title)
interface.launch() |