Spaces:
Running
Running
File size: 3,979 Bytes
f62f5b8 |
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 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 |
import gradio as gr
import torch
import torch.nn.functional as F
from transformers import AutoTokenizer
import pickle
from models.rnn import RNNClassifier
from models.lstm import LSTMClassifier
from models.transformer import TransformerClassifier
from utility import simple_tokenizer
# =========================
# Load models and vocab
# =========================
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model_name = "prajjwal1/bert-tiny"
def load_vocab():
with open("pretrained_models/vocab.pkl", "rb") as f:
return pickle.load(f)
def load_models(vocab_size, output_dim=6, padding_idx=0):
rnn_model = RNNClassifier(vocab_size, 128, 128, output_dim, padding_idx)
rnn_model.load_state_dict(torch.load("pretrained_models/best_rnn.pt", map_location=device))
rnn_model = rnn_model.to(device)
rnn_model.eval()
lstm_model = LSTMClassifier(vocab_size, 128, 128, output_dim, padding_idx)
lstm_model.load_state_dict(torch.load("pretrained_models/best_lstm.pt", map_location=device))
lstm_model = lstm_model.to(device)
lstm_model.eval()
transformer_model = TransformerClassifier(model_name, output_dim)
transformer_model.load_state_dict(torch.load("pretrained_models/best_transformer.pt", map_location=device))
transformer_model = transformer_model.to(device)
transformer_model.eval()
return rnn_model, lstm_model, transformer_model
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
vocab = load_vocab()
tokenizer = AutoTokenizer.from_pretrained(model_name)
rnn_model, lstm_model, transformer_model = load_models(len(vocab))
emotions = ["anger", "fear", "joy", "love", "sadness", "surprise"]
def predict(model, text, model_type, vocab, tokenizer=None, max_length=32):
if model_type in ["rnn", "lstm"]:
# Match collate_fn_rnn but with no random truncation
tokens = simple_tokenizer(text)
ids = [vocab.get(token, vocab["<UNK>"]) for token in tokens]
if len(ids) < max_length:
ids += [vocab["<PAD>"]] * (max_length - len(ids))
else:
ids = ids[:max_length]
input_ids = torch.tensor([ids], dtype=torch.long).to(device)
outputs = model(input_ids)
else:
# Match collate_fn_transformer but with no partial_prob
encoding = tokenizer(
text,
padding="max_length",
truncation=True,
max_length=128,
return_tensors="pt"
)
input_ids = encoding["input_ids"].to(device)
attention_mask = encoding["attention_mask"].to(device)
outputs = model(input_ids=input_ids, attention_mask=attention_mask)
probs = F.softmax(outputs, dim=-1)
return probs.squeeze().detach().cpu().numpy()
# =========================
# Gradio App
# =========================
def emotion_typeahead(text):
if len(text.strip()) <= 2:
return {}, {}, {}
rnn_probs = predict(rnn_model, text.strip(), "rnn", vocab)
lstm_probs = predict(lstm_model, text.strip(), "lstm", vocab)
transformer_probs = predict(transformer_model, text.strip(), "transformer", vocab, tokenizer)
rnn_dict = {emo: float(prob) for emo, prob in zip(emotions, rnn_probs)}
lstm_dict = {emo: float(prob) for emo, prob in zip(emotions, lstm_probs)}
transformer_dict = {emo: float(prob) for emo, prob in zip(emotions, transformer_probs)}
return rnn_dict, lstm_dict, transformer_dict
with gr.Blocks() as demo:
gr.Markdown("## 🎯 Emotion Typeahead Predictor (RNN, LSTM, Transformer)")
text_input = gr.Textbox(label="Type your sentence here...")
with gr.Row():
rnn_output = gr.Label(label="🧠 RNN Prediction")
lstm_output = gr.Label(label="🧠 LSTM Prediction")
transformer_output = gr.Label(label="🧠 Transformer Prediction")
text_input.change(emotion_typeahead, inputs=text_input, outputs=[rnn_output, lstm_output, transformer_output])
demo.launch()
|