Spaces:
Sleeping
Sleeping
import gradio as gr | |
import numpy as np | |
import pandas as pd | |
from transformers import AutoTokenizer, AutoModelForSequenceClassification | |
import torch | |
target_list = ["Playful", "Infuriating", "Sentimental", "Cynical", "Depressing", "Awe-inspiring", "Patriotic", "Begrudging", "Educational", "Hopeful", | |
"Sarcastic", "Disrespectful", "Disparaging"] | |
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') | |
model_name = 'valurank/finetuned-distilbert-multi-label-emotion' | |
model = AutoModelForSequenceClassification.from_pretrained(model_name) | |
tokenizer = AutoTokenizer.from_pretrained(model_name) | |
def get_probs(logits, threshold=0.5): | |
sigm = 1 / (1 + np.exp(-logits)) | |
return sigm | |
def multi_label_emotions(text): | |
inputs = tokenizer(text, return_tensors="pt", truncation=True) | |
model.to(device) | |
with torch.no_grad(): | |
logits = model(**inputs).logits | |
#probs = logits.int().numpy()[0] | |
log_probs = get_probs(logits) | |
final_log_probs = [] | |
for log in log_probs: | |
final_log_probs.append(log.numpy()) | |
final_output = [] | |
for i in zip(final_log_probs[0], target_list): | |
final_output.append(i) | |
final_output.sort(reverse=True) | |
final_dict = {} | |
for k,v in final_output: | |
final_dict[v] = float(k) | |
return final_dict | |
demo = gr.Interface(multi_label_emotions, inputs=gr.inputs.Textbox(), | |
outputs = gr.Label(num_top_classes=16), | |
title='Multi-label-emotion-classification') | |
if __name__ == '__main__': | |
demo.launch(debug=True) |