Spaces:
Sleeping
Sleeping
File size: 2,656 Bytes
46c07fd |
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 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 |
import gradio as gr
import torch
from transformers import (
BertTokenizer,
XLNetTokenizer,
GPT2Tokenizer,
AutoModelForSequenceClassification
)
# Model repositories on Hugging Face Hub
model_repos = {
"BERT": "sk23aib/emotion-bert",
"XLNet": "sk23aib/emotion-xlnet",
"GPT-2": "sk23aib/emotion-gpt2"
}
# Emotion labels (must match model training order)
emotion_labels = [
"anger", "boredom", "empty", "enthusiasm", "fun", "happiness", "hate",
"love", "neutral", "relief", "sadness", "surprise", "worry"
]
# Load models and tokenizers
loaded_models = {}
# BERT
bert_tokenizer = BertTokenizer.from_pretrained(model_repos["BERT"])
bert_model = AutoModelForSequenceClassification.from_pretrained(model_repos["BERT"])
bert_model.eval()
loaded_models["BERT"] = {"tokenizer": bert_tokenizer, "model": bert_model}
# XLNet
xlnet_tokenizer = XLNetTokenizer.from_pretrained(model_repos["XLNet"])
xlnet_model = AutoModelForSequenceClassification.from_pretrained(model_repos["XLNet"])
xlnet_model.eval()
loaded_models["XLNet"] = {"tokenizer": xlnet_tokenizer, "model": xlnet_model}
# GPT-2
gpt2_tokenizer = GPT2Tokenizer.from_pretrained(model_repos["GPT-2"], padding_side="left")
gpt2_tokenizer.pad_token = gpt2_tokenizer.eos_token # Required for GPT-2
gpt2_model = AutoModelForSequenceClassification.from_pretrained(model_repos["GPT-2"])
gpt2_model.config.pad_token_id = gpt2_tokenizer.pad_token_id
gpt2_model.eval()
loaded_models["GPT-2"] = {"tokenizer": gpt2_tokenizer, "model": gpt2_model}
# Inference function – return top emotion + probability
def predict_emotions(text):
output_lines = []
with torch.no_grad():
for model_name, components in loaded_models.items():
tokenizer = components["tokenizer"]
model = components["model"]
inputs = tokenizer(text, return_tensors="pt", truncation=True, padding=True)
logits = model(**inputs).logits
probs = torch.nn.functional.softmax(logits, dim=1)[0]
top_idx = torch.argmax(probs).item()
top_emotion = emotion_labels[top_idx]
top_confidence = round(float(probs[top_idx]), 4)
output_lines.append(f"{model_name}: {top_emotion} ({top_confidence})")
return "\n".join(output_lines)
# Gradio Interface
interface = gr.Interface(
fn=predict_emotions,
inputs=gr.Textbox(lines=3, placeholder="Type a sentence to analyze..."),
outputs=gr.Textbox(label="Top Emotion by Model"),
title="Multi-Model Emotion Classifier",
description="See which emotion is predicted by BERT, XLNet, and GPT-2, along with their confidence."
)
interface.launch()
|