Upload 18 files

app_gradio.py

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

models/lstm.py

@@ -0,0 +1,20 @@
+import torch
+import torch.nn as nn
+
+class LSTMClassifier(nn.Module):
+    def __init__(self, vocab_size, embed_dim, hidden_dim, output_dim, padding_idx):
+        super(LSTMClassifier, self).__init__()
+        self.embedding = nn.Embedding(vocab_size, embed_dim, padding_idx=padding_idx)
+        self.lstm = nn.LSTM(embed_dim, hidden_dim, num_layers=1, dropout=0.3, batch_first=True, bidirectional=True)
+        self.fc1 = nn.Linear(hidden_dim * 2, hidden_dim)
+        self.relu = nn.ReLU()
+        self.fc2 = nn.Linear(hidden_dim, output_dim)
+
+    def forward(self, x):
+        embedded = self.embedding(x)
+        output, (hidden, _) = self.lstm(embedded)
+        hidden_cat = torch.cat((hidden[-2,:,:], hidden[-1,:,:]), dim=1)  # concatenate last hidden states
+        x = self.fc1(hidden_cat)
+        x = self.relu(x)
+        out = self.fc2(x)
+        return out

models/rnn.py

@@ -0,0 +1,19 @@
+import torch
+import torch.nn as nn
+
+class RNNClassifier(nn.Module):
+    def __init__(self, vocab_size, embed_dim, hidden_dim, output_dim, padding_idx):
+        super(RNNClassifier, self).__init__()
+        self.embedding = nn.Embedding(vocab_size, embed_dim, padding_idx=padding_idx)
+        self.rnn = nn.RNN(embed_dim, hidden_dim, batch_first=True)
+        self.fc1 = nn.Linear(hidden_dim, hidden_dim // 2)  # New hidden layer
+        self.relu = nn.ReLU()
+        self.fc2 = nn.Linear(hidden_dim // 2, output_dim)
+
+    def forward(self, x):
+        embedded = self.embedding(x)  # [batch_size, seq_len, embed_dim]
+        output, hidden = self.rnn(embedded)  # hidden: [1, batch_size, hidden_dim]
+        x = self.fc1(hidden.squeeze(0))  # [batch_size, hidden_dim//2]
+        x = self.relu(x)
+        out = self.fc2(x)  # [batch_size, output_dim]
+        return out

models/transformer.py

@@ -0,0 +1,20 @@
+import torch
+import torch.nn as nn
+from transformers import AutoModel
+
+class TransformerClassifier(nn.Module):
+    def __init__(self, model_name, output_dim):
+        super(TransformerClassifier, self).__init__()
+        self.transformer = AutoModel.from_pretrained(model_name)
+        # Freeze bottom 3 layers, unfreeze top layers
+        for name, param in self.transformer.named_parameters():
+            if "layer.0" in name or "layer.1" in name or "layer.2" in name:
+                param.requires_grad = False
+        self.fc = nn.Linear(self.transformer.config.hidden_size, output_dim)
+
+    def forward(self, input_ids, attention_mask):
+        outputs = self.transformer(input_ids=input_ids, attention_mask=attention_mask)
+        hidden_state = outputs.last_hidden_state  # [batch_size, seq_len, hidden_dim]
+        pooled_output = hidden_state[:, 0]  # Use CLS token output
+        out = self.fc(pooled_output)
+        return out

pretrained_models/best_lstm.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:2ab822a8bc9c588f9e4b8dc7a2a85385528657b206b02d62d184217bb379ae76
+size 4984493

pretrained_models/best_rnn.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:0ef0a95f7b1655f46b26691395bbc94558c5346c09f7f90177dbf5c8cb578da0
+size 3959329

pretrained_models/best_transformer.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:7ae3d6c4db4de20f3e2251c67760fa91f6fe3f475135e45c3492114dfd11da0b
+size 17564687

pretrained_models/vocab.pkl

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:276d70ee4c7d337f6d07514ee81043fa7c7475e815fb153a667400a45df99f2d
+size 93532

requirements.txt

@@ -0,0 +1,10 @@
+torch
+transformers
+datasets
+scikit-learn
+gradio
+numpy==1.26.4
+scipy
+tqdm
+streamlit
+gradio

train.py

@@ -0,0 +1,256 @@
+import os
+import time
+import torch
+import matplotlib.pyplot as plt
+import seaborn as sns
+from collections import Counter
+from sklearn.metrics import confusion_matrix, ConfusionMatrixDisplay
+from torch.utils.data import DataLoader
+from transformers import AutoTokenizer
+from utility import (
+    load_emotion_dataset,
+    encode_labels,
+    build_vocab,
+    collate_fn_rnn,
+    collate_fn_transformer
+)
+from models.rnn import RNNClassifier
+from models.lstm import LSTMClassifier
+from models.transformer import TransformerClassifier
+from tqdm import tqdm
+
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+def summarize_class_distribution(dataset, label_encoder):
+    labels = [example["label"] for example in dataset]
+    counter = Counter(labels)
+    print("\n🔍 Class distribution:")
+    for label_idx, count in sorted(counter.items()):
+        label_name = label_encoder.inverse_transform([label_idx])[0]
+        print(f"{label_name:>10}: {count}")
+
+def plot_class_countplot(dataset, label_encoder):
+    labels = [example["label"] for example in dataset]
+    counts = Counter(labels)
+    label_display = [label_encoder.inverse_transform([i])[0] for i in sorted(counts.keys())]
+    values = [counts[i] for i in sorted(counts.keys())]
+
+    plt.figure(figsize=(8, 5))
+    sns.barplot(x=label_display, y=values)
+    plt.title("Emotion Class Distribution (Training Set)")
+    plt.xlabel("Emotion")
+    plt.ylabel("Count")
+    plt.tight_layout()
+    os.makedirs("plots", exist_ok=True)
+    plt.savefig("plots/class_distribution.png")
+    plt.close()
+
+def plot_loss_curve(train_losses, test_losses, model_name):
+    plt.figure(figsize=(8, 4))
+    plt.plot(train_losses, label="Train Loss")
+    plt.plot(test_losses, label="Test Loss")
+    plt.xlabel("Epoch")
+    plt.ylabel("Loss")
+    plt.title(f"{model_name} Train vs Test Loss")
+    plt.legend()
+    os.makedirs("plots", exist_ok=True)
+    plt.savefig(f"plots/{model_name.lower()}_loss_curve.png")
+    plt.close()
+
+def compute_test_loss(model, dataloader, criterion, model_type):
+    total_loss = 0
+    with torch.no_grad():
+        model.eval()
+        for batch in dataloader:
+            if isinstance(batch, tuple):
+                input_ids, labels = batch
+                attention_mask = None
+            else:
+                input_ids = batch["input_ids"]
+                attention_mask = batch.get("attention_mask", None)
+                labels = batch["labels"]
+
+            input_ids = input_ids.to(device)
+            labels = labels.to(device)
+            if attention_mask is not None:
+                attention_mask = attention_mask.to(device)
+
+            if model_type == "transformer":
+                outputs = model(input_ids=input_ids, attention_mask=attention_mask)
+            else:
+                outputs = model(input_ids)
+
+            loss = criterion(outputs, labels)
+            total_loss += loss.item()
+    return total_loss / len(dataloader)
+
+def train_model(model, train_loader, test_loader, optimizer, criterion, epochs, model_type="rnn"):
+    train_losses = []
+    test_losses = []
+
+    for epoch in range(epochs):
+        model.train()
+        start_time = time.time()
+        total_loss = 0
+        progress_bar = tqdm(train_loader, desc=f"Epoch {epoch + 1}", ncols=100)
+
+        for batch in progress_bar:
+            optimizer.zero_grad()
+
+            if isinstance(batch, tuple):
+                input_ids, labels = batch
+                attention_mask = None
+            else:
+                input_ids = batch["input_ids"]
+                attention_mask = batch.get("attention_mask", None)
+                labels = batch["labels"]
+
+            input_ids = input_ids.to(device)
+            labels = labels.to(device)
+            if attention_mask is not None:
+                attention_mask = attention_mask.to(device)
+
+            if model_type == "transformer":
+                outputs = model(input_ids=input_ids, attention_mask=attention_mask)
+            else:
+                outputs = model(input_ids)
+
+            loss = criterion(outputs, labels)
+            loss.backward()
+            optimizer.step()
+
+            total_loss += loss.item()
+            avg_loss = total_loss / len(train_loader)
+            progress_bar.set_postfix({"Avg Loss": f"{avg_loss:.4f}"})
+
+        test_loss = compute_test_loss(model, test_loader, criterion, model_type)
+        train_losses.append(avg_loss)
+        test_losses.append(test_loss)
+
+        print(f"✅ Epoch {epoch + 1} | Train: {avg_loss:.4f} | Test: {test_loss:.4f} | Time: {time.time() - start_time:.2f}s")
+
+    torch.cuda.empty_cache()
+    del model
+    return train_losses, test_losses
+
+def evaluate_preds(model, dataloader, model_type="rnn"):
+    model.eval()
+    all_preds = []
+    all_labels = []
+    with torch.no_grad():
+        for batch in dataloader:
+            if isinstance(batch, tuple):
+                input_ids, labels = batch
+                attention_mask = None
+            else:
+                input_ids = batch["input_ids"]
+                attention_mask = batch.get("attention_mask", None)
+                labels = batch["labels"]
+
+            input_ids = input_ids.to(device)
+            labels = labels.to(device)
+            if attention_mask is not None:
+                attention_mask = attention_mask.to(device)
+
+            if model_type == "transformer":
+                outputs = model(input_ids=input_ids, attention_mask=attention_mask)
+            else:
+                outputs = model(input_ids)
+
+            preds = torch.argmax(outputs, dim=1)
+            all_preds.extend(preds.cpu().tolist())
+            all_labels.extend(labels.cpu().tolist())
+    return all_labels, all_preds
+
+def plot_confusion_matrices(y_true_train, y_pred_train, y_true_test, y_pred_test, labels, title, filename):
+    fig, axes = plt.subplots(1, 2, figsize=(14, 6))
+    cm_train = confusion_matrix(y_true_train, y_pred_train)
+    cm_test = confusion_matrix(y_true_test, y_pred_test)
+
+    ConfusionMatrixDisplay(cm_train, display_labels=labels).plot(ax=axes[0], cmap='Blues', colorbar=False)
+    axes[0].set_title(f"{title} - Train")
+
+    ConfusionMatrixDisplay(cm_test, display_labels=labels).plot(ax=axes[1], cmap='Oranges', colorbar=False)
+    axes[1].set_title(f"{title} - Test")
+
+    plt.tight_layout()
+    os.makedirs("plots", exist_ok=True)
+    plt.savefig(f"plots/{filename}")
+    plt.close()
+
+# Load and encode data
+data = load_emotion_dataset("train")
+train_data, label_encoder = encode_labels(data)
+test_data, _ = encode_labels(load_emotion_dataset("test"))
+labels = label_encoder.classes_
+output_dim = len(labels)
+padding_idx = 0
+
+summarize_class_distribution(train_data, label_encoder)
+plot_class_countplot(train_data, label_encoder)
+
+# Build vocab
+vocab = build_vocab(train_data)
+
+model_name = "prajjwal1/bert-tiny"
+tokenizer = AutoTokenizer.from_pretrained(model_name)
+
+# DataLoaders (no augmentation)
+train_loader_rnn = DataLoader(train_data, batch_size=64, shuffle=True, collate_fn=lambda b: collate_fn_rnn(b, vocab, partial_prob=0.0))
+test_loader_rnn = DataLoader(test_data, batch_size=64, shuffle=False, collate_fn=lambda b: collate_fn_rnn(b, vocab, partial_prob=0.0))
+
+train_loader_tf = DataLoader(train_data, batch_size=64, shuffle=True, collate_fn=lambda b: collate_fn_transformer(b, tokenizer, partial_prob=0.0))
+test_loader_tf = DataLoader(test_data, batch_size=64, shuffle=False, collate_fn=lambda b: collate_fn_transformer(b, tokenizer, partial_prob=0.0))
+
+# Initialize and train models
+rnn = RNNClassifier(len(vocab), 128, 128, output_dim, padding_idx).to(device)
+lstm = LSTMClassifier(len(vocab), 128, 128, output_dim, padding_idx).to(device)
+transformer = TransformerClassifier(model_name, output_dim).to(device)
+
+criterion = torch.nn.CrossEntropyLoss()
+
+# rnn_train_losses, rnn_test_losses = train_model(rnn, train_loader_rnn, test_loader_rnn, torch.optim.Adam(rnn.parameters(), lr=1e-4), criterion, epochs=50, model_type="rnn")
+# torch.save(rnn.state_dict(), "pretrained_models/best_rnn.pt")
+# plot_loss_curve(rnn_train_losses, rnn_test_losses, "RNN")
+#
+# lstm_train_losses, lstm_test_losses = train_model(lstm, train_loader_rnn, test_loader_rnn, torch.optim.Adam(lstm.parameters(), lr=1e-4), criterion, epochs=50, model_type="lstm")
+# torch.save(lstm.state_dict(), "pretrained_models/best_lstm.pt")
+# plot_loss_curve(lstm_train_losses, lstm_test_losses, "LSTM")
+
+tf_train_losses, tf_test_losses = train_model(transformer, train_loader_tf, test_loader_tf, torch.optim.Adam(transformer.parameters(), lr=2e-5), criterion, epochs=50, model_type="transformer")
+torch.save(transformer.state_dict(), "pretrained_models/best_transformer.pt")
+plot_loss_curve(tf_train_losses, tf_test_losses, "Transformer")
+
+# Evaluate and plot
+model_paths = {
+    "RNN": "pretrained_models/best_rnn.pt",
+    "LSTM": "pretrained_models/best_lstm.pt",
+    "Transformer": "pretrained_models/best_transformer.pt"
+}
+
+for name in ["RNN", "LSTM", "Transformer"]:
+    if name == "RNN":
+        model = RNNClassifier(len(vocab), 128, 128, output_dim, padding_idx).to(device)
+        loader = train_loader_rnn
+        test_loader = test_loader_rnn
+    elif name == "LSTM":
+        model = LSTMClassifier(len(vocab), 128, 128, output_dim, padding_idx).to(device)
+        loader = train_loader_rnn
+        test_loader = test_loader_rnn
+    else:
+        model = TransformerClassifier(model_name, output_dim).to(device)
+        loader = train_loader_tf
+        test_loader = test_loader_tf
+
+    model.load_state_dict(torch.load(model_paths[name]))
+    model.eval()
+
+    y_train_true, y_train_pred = evaluate_preds(model, loader, model_type=name.lower())
+    y_test_true, y_test_pred = evaluate_preds(model, test_loader, model_type=name.lower())
+
+    plot_confusion_matrices(
+        y_train_true, y_train_pred, y_test_true, y_test_pred,
+        labels=labels,
+        title=name,
+        filename=f"{name.lower()}_confusion_matrices.png"
+    )

utility.py

@@ -0,0 +1,94 @@
+import re
+import torch
+from collections import Counter
+from datasets import load_dataset
+from sklearn.preprocessing import LabelEncoder
+from transformers import AutoTokenizer
+import random
+
+# ====== Dataset Loading ======
+
+def load_emotion_dataset(split="train"):
+    return load_dataset("dair-ai/emotion", split=split)
+
+def encode_labels(dataset):
+    le = LabelEncoder()
+    all_labels = [example["label"] for example in dataset]
+    le.fit(all_labels)
+    dataset = dataset.map(lambda x: {"label": le.transform([x["label"]])[0]})
+    return dataset, le
+
+# ====== Tokenizer for RNN/LSTM ======
+
+def simple_tokenizer(text):
+    text = text.lower()
+    text = re.sub(r"[^a-z0-9\s]", "", text)  # Remove special characters
+    return text.split()
+
+# ====== Vocab Builder for RNN/LSTM ======
+
+def build_vocab(dataset, min_freq=2):
+    counter = Counter()
+    for example in dataset:
+        tokens = simple_tokenizer(example["text"])
+        counter.update(tokens)
+
+    vocab = {"<PAD>": 0, "<UNK>": 1}
+    idx = 2
+    for word, freq in counter.items():
+        if freq >= min_freq:
+            vocab[word] = idx
+            idx += 1
+    return vocab
+
+# ====== Collate Function for RNN/LSTM ======
+
+def collate_fn_rnn(batch, vocab, max_length=32, partial_prob=0.0):
+    texts = [item["text"] for item in batch]
+    labels = [item["label"] for item in batch]
+
+    all_input_ids = []
+    for text in texts:
+        tokens = simple_tokenizer(text)
+
+        # 🔥 Randomly truncate tokens with some probability
+        if random.random() < partial_prob and len(tokens) > 5:
+            # Keep between 30% to 70% of the tokens
+            cutoff = random.randint(int(len(tokens)*0.3), int(len(tokens)*0.7))
+            tokens = tokens[:cutoff]
+
+        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]
+        all_input_ids.append(ids)
+
+    input_ids = torch.tensor(all_input_ids)
+    labels = torch.tensor(labels)
+    return input_ids, labels
+
+# ====== Collate Function for Transformer ======
+
+def collate_fn_transformer(batch, tokenizer, max_length=128, partial_prob=0.5):
+    import random
+    texts = []
+    labels = []
+
+    for item in batch:
+        text = item["text"]
+        tokens = text.split()
+
+        # 🔥 Random truncation
+        if random.random() < partial_prob and len(tokens) > 5:
+            cutoff = random.randint(int(len(tokens)*0.3), int(len(tokens)*0.7))
+            tokens = tokens[:cutoff]
+            text = " ".join(tokens)
+
+        texts.append(text)
+        labels.append(item["label"])
+
+    encoding = tokenizer(texts, padding="max_length", truncation=True, max_length=max_length, return_tensors="pt")
+    encoding["labels"] = torch.tensor(labels)
+    return encoding
+