Update app.py

app.py

@@ -1,40 +1,119 @@
+import warnings
 
-import pandas as pd
-from sklearn.model_selection import train_test_split
-from sklearn.feature_extraction.text import TfidfVectorizer
-from sklearn.naive_bayes import MultinomialNB
-from sklearn.metrics import classification_report, accuracy_score
 import gradio as gr
+import matplotlib.pyplot as plt
+import seaborn as sns
+import torch
+from torch import cuda, nn
+from torch.nn import functional as F
 
-import datasets
+warnings.filterwarnings("ignore")
 
-df_true = pd.read_csv('training.1600000.processed.noemoticon.csv')
+device = 'cuda' if cuda.is_available() else 'cpu'
 
 
+class Attention(nn.Module):
+    def __init__(self, in_features, *args, bias=True, **kwargs):
+        super(Attention, self).__init__(*args, **kwargs)
+        self.bias = bias
+        self.W = nn.Parameter(torch.randn(in_features, in_features))
+        if self.bias:
+            self.b = nn.Parameter(torch.randn(in_features))
+        self.u = nn.Parameter(torch.randn(in_features))
+        self.tanh = nn.Tanh()
+        self.softmax = nn.Softmax(dim=-1)
 
-df['polarity'] = df['polarity'].replace({0: 'negative', 2: 'neutral', 4: 'positive'})
-df = df[['text', 'polarity']]
+    def forward(self, x):
+        uit = torch.matmul(x, self.W)
+        if self.bias:
+            uit += self.b
+        ait = torch.matmul(self.tanh(uit), self.u)
+        attention = self.softmax(ait)
+        return attention
 
 
-train_data, test_data = train_test_split(df, test_size=0.2, random_state=42)
+class GenderClassifierWithAttention(nn.Module):
+    def __init__(self, input_size, hidden_size, embed_size, *args, **kwargs):
+        super(GenderClassifierWithAttention, self).__init__(*args, **kwargs)
+        self.hid_dim = hidden_size
 
+        # embedding layer
+        self.embedding = nn.Embedding(
+            num_embeddings=input_size, embedding_dim=embed_size)
 
-vectorizer = TfidfVectorizer(max_features=5000)
-X_train = vectorizer.fit_transform(train_data['text'])
-y_train = train_data['polarity']
-classifier = MultinomialNB()
-classifier.fit(X_train, y_train)
+        # attention layer
+        self.attention = Attention(self.hid_dim * 2)
 
+        # LSTM layer
+        self.lstm = nn.LSTM(embed_size, hidden_size=self.hid_dim,
+                            num_layers=2, bidirectional=True, batch_first=True)
 
-X_test = vectorizer.transform(test_data['text'])
-y_pred = classifier.predict(X_test)
-print("Accuracy:", accuracy_score(test_data['polarity'], y_pred))
-print("Classification Report:\n", classification_report(test_data['polarity'], y_pred))
+        # dropout layers
+        self.dropout1 = nn.Dropout(p=0.4)
+        self.dropout2 = nn.Dropout(p=0.4)
 
+        # normalization layers
+        self.batch_norm1 = nn.BatchNorm1d(num_features=15)
+        self.batch_norm2 = nn.BatchNorm1d(num_features=7)
+        self.layer_norm = nn.LayerNorm(40)
 
-def predict_sentiment(text):
-    text_vectorized = vectorizer.transform([text])
-    prediction = classifier.predict(text_vectorized)[0]
-    return prediction
-iface = gr.Interface(fn=predict_sentiment, inputs="text", outputs="text")
-iface.launch()
+        # linear layers
+        self.fc1 = nn.Linear(in_features=self.hid_dim * 2, out_features=15)
+        self.fc2 = nn.Linear(in_features=15, out_features=7)
+        self.fc3 = nn.Linear(in_features=7, out_features=2)
+
+        # activation functions
+        self.tanh = nn.Tanh()
+        self.relu = nn.ReLU()
+
+    def forward(self, x):
+        x = self.embedding(x).float()
+
+        out, (h, c) = self.lstm(x)
+        out = self.layer_norm(out)
+        attention = self.attention(self.tanh(out))
+        x = torch.einsum('ijk,ij->ik', out, attention)
+
+        x = self.tanh(self.fc1(x))
+        x = self.batch_norm1(x)
+        x = self.dropout1(x)
+
+        x = self.tanh(self.fc2(x))
+        x = self.batch_norm2(x)
+        x = self.dropout2(x)
+
+        x = self.fc3(x)
+        return x, attention
+
+
+def get_attention(model):
+    def inner(name):
+        import string
+        char_stoi = {val: key for key, val in dict(
+            enumerate(["<PAD>"] + list(string.ascii_lowercase))).items()}
+
+        name_mapped = [char_stoi[char] for char in name.lower()]
+
+        probs, attention = model(torch.tensor([name_mapped]).to(device))
+        probs = F.softmax(probs, dim=-1)
+        probs_dict = dict(
+            zip(["female", "male"], probs.cpu().detach().numpy().flatten().tolist()))
+
+        fig, ax = plt.subplots(nrows=1, ncols=1)
+        _ = sns.barplot(x=list(range(attention.shape[1])),
+                        y=attention.squeeze().cpu().detach().numpy(),
+                        color="#1FCECB", ax=ax)
+        _ = ax.set_xticklabels(list(name))
+
+        return probs_dict, fig
+
+    return inner
+
+
+if __name__ == "__main__":
+    model = torch.load("lstm_attention.pt", map_location=device)
+    interface = gr.Interface(get_attention(model),
+                             inputs="text", outputs=[gr.outputs.Label(label="gender"), gr.outputs.Image(type="plot", label="attention")],
+                             title="Visualizing Attention in Gender Classification",
+                             examples=["annabella", "ewart",
+                                       "blancha", "bronson"], allow_flagging="never").launch()