Spaces:
Running
Running
File size: 4,683 Bytes
c4fe4b1 | 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 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 | import torch
import torch.nn as nn
import re
import json
import os
# Define the 5 target emotions in alphabetical order for consistency
EMOTIONS = ["Bored", "Confident", "Confused", "Curious", "Frustrated"]
class BiLSTMClassifier(nn.Module):
def __init__(self, vocab_size, embedding_dim=64, hidden_dim=64, output_dim=5):
super().__init__()
self.embedding = nn.Embedding(vocab_size, embedding_dim, padding_idx=0)
self.lstm = nn.LSTM(embedding_dim, hidden_dim, batch_first=True, bidirectional=True)
self.fc = nn.Linear(hidden_dim * 2, output_dim)
def forward(self, x):
# x shape: [batch, seq_len]
embedded = self.embedding(x) # [batch, seq_len, embedding_dim]
lstm_out, _ = self.lstm(embedded) # [batch, seq_len, hidden_dim * 2]
# Mean pooling over sequence dimension
pooled = torch.mean(lstm_out, dim=1) # [batch, hidden_dim * 2]
logits = self.fc(pooled) # [batch, output_dim]
return logits
def clean_text(text):
"""
Cleans text by lowercasing, removing non-alphanumeric characters except spaces,
and removing extra whitespaces.
"""
if not text:
return ""
text = text.lower()
text = re.sub(r'[^a-z0-9\s]', '', text)
text = re.sub(r'\s+', ' ', text).strip()
return text
def tokenize(text):
"""
Tokenizes clean text into a list of words.
"""
return clean_text(text).split()
def text_to_indices(text, vocab, max_len=50):
"""
Converts a text string into a list of vocabulary indices with padding/truncation.
"""
words = tokenize(text)
indices = []
for w in words:
indices.append(vocab.get(w, 1)) # 1 is <unk>
# Pad or truncate
if len(indices) < max_len:
indices += [0] * (max_len - len(indices)) # 0 is <pad>
else:
indices = indices[:max_len]
return indices
class BiLSTMWrapper:
def __init__(self, model_path="bilstm_model.pth", vocab_path="vocab.json"):
self.model_path = model_path
self.vocab_path = vocab_path
self.model = None
self.vocab = None
self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
if os.path.exists(model_path) and os.path.exists(vocab_path):
self.load()
def load(self):
# Load vocabulary
with open(self.vocab_path, "r", encoding="utf-8") as f:
self.vocab = json.load(f)
# Initialize and load model weights
vocab_size = len(self.vocab)
self.model = BiLSTMClassifier(vocab_size)
self.model.load_state_dict(torch.load(self.model_path, map_location=self.device, weights_only=True))
self.model.to(self.device)
self.model.eval()
def predict(self, text):
"""
Runs inference on the input text and returns predictions in the unified schema.
"""
if self.model is None or self.vocab is None:
# If not loaded, return default flat predictions
flat_score = 1.0 / len(EMOTIONS)
all_emotions = {e: flat_score for e in EMOTIONS}
return {
"primary_emotion": EMOTIONS[0],
"primary_confidence": round(flat_score, 4),
"all_emotions": {e: round(flat_score, 4) for e in EMOTIONS},
"mixed_emotions": []
}
indices = text_to_indices(text, self.vocab)
tensor = torch.tensor([indices], dtype=torch.long).to(self.device)
with torch.no_grad():
logits = self.model(tensor)
probs = torch.softmax(logits, dim=1).squeeze().cpu().numpy()
all_emotions = {}
for emotion, prob in zip(EMOTIONS, probs):
all_emotions[emotion] = round(float(prob), 4)
# Find primary emotion
primary_emotion = max(all_emotions, key=all_emotions.get)
primary_confidence = all_emotions[primary_emotion]
# Mixed emotions: any emotion other than the top one with confidence >= 15%
mixed_emotions = []
for emotion, prob in all_emotions.items():
if emotion != primary_emotion and prob >= 0.15:
mixed_emotions.append([emotion, round(prob, 4)])
# Sort mixed emotions descending by score
mixed_emotions.sort(key=lambda x: x[1], reverse=True)
return {
"primary_emotion": primary_emotion,
"primary_confidence": primary_confidence,
"all_emotions": all_emotions,
"mixed_emotions": mixed_emotions
}
|