Update app.py

app.py

@@ -1,108 +1,107 @@
-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"))
-    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"))
-    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()
+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()