Create RNN.py

pages/RNN.py

@@ -0,0 +1,165 @@
+import streamlit as st
+import torch
+import torch.nn as nn
+import torch.optim as optim
+from torchtext.legacy import data, datasets
+import matplotlib.pyplot as plt
+import seaborn as sns
+import pandas as pd
+import numpy as np
+
+# Define the RNN model
+class RNN(nn.Module):
+    def __init__(self, vocab_size, embed_size, hidden_size, output_size, n_layers, dropout):
+        super(RNN, self).__init__()
+        self.embedding = nn.Embedding(vocab_size, embed_size)
+        self.rnn = nn.RNN(embed_size, hidden_size, n_layers, dropout=dropout, batch_first=True)
+        self.fc = nn.Linear(hidden_size, output_size)
+        self.dropout = nn.Dropout(dropout)
+
+    def forward(self, x):
+        x = self.dropout(self.embedding(x))
+        h0 = torch.zeros(n_layers, x.size(0), hidden_size).to(device)
+        out, _ = self.rnn(x, h0)
+        out = self.fc(out[:, -1, :])
+        return out
+
+# Function to load the data
+@st.cache_data
+def load_data():
+    TEXT = data.Field(tokenize='spacy', tokenizer_language='en_core_web_sm')
+    LABEL = data.LabelField(dtype=torch.float)
+    train_data, test_data = datasets.IMDB.splits(TEXT, LABEL)
+    train_data, valid_data = train_data.split(split_ratio=0.8)
+
+    MAX_VOCAB_SIZE = 25_000
+    TEXT.build_vocab(train_data, max_size=MAX_VOCAB_SIZE, vectors="glove.6B.100d", unk_init=torch.Tensor.normal_)
+    LABEL.build_vocab(train_data)
+
+    BATCH_SIZE = 64
+    train_iterator, valid_iterator, test_iterator = data.BucketIterator.splits(
+        (train_data, valid_data, test_data), 
+        batch_size=BATCH_SIZE,
+        device=device)
+    
+    return TEXT, LABEL, train_iterator, valid_iterator, test_iterator
+
+# Function to train the network
+def train_network(net, iterator, optimizer, criterion, epochs):
+    loss_values = []
+    for epoch in range(epochs):
+        epoch_loss = 0
+        net.train()
+        for batch in iterator:
+            optimizer.zero_grad()
+            predictions = net(batch.text).squeeze(1)
+            loss = criterion(predictions, batch.label)
+            loss.backward()
+            optimizer.step()
+            epoch_loss += loss.item()
+        epoch_loss /= len(iterator)
+        loss_values.append(epoch_loss)
+        st.write(f'Epoch {epoch + 1}: loss {epoch_loss:.3f}')
+    st.write('Finished Training')
+    return loss_values
+
+# Function to evaluate the network
+def evaluate_network(net, iterator, criterion):
+    epoch_loss = 0
+    correct = 0
+    total = 0
+    all_labels = []
+    all_predictions = []
+    net.eval()
+    with torch.no_grad():
+        for batch in iterator:
+            predictions = net(batch.text).squeeze(1)
+            loss = criterion(predictions, batch.label)
+            epoch_loss += loss.item()
+            rounded_preds = torch.round(torch.sigmoid(predictions))
+            correct += (rounded_preds == batch.label).sum().item()
+            total += len(batch.label)
+            all_labels.extend(batch.label.cpu().numpy())
+            all_predictions.extend(rounded_preds.cpu().numpy())
+    accuracy = 100 * correct / total
+    st.write(f'Loss: {epoch_loss / len(iterator):.4f}, Accuracy: {accuracy:.2f}%')
+    return accuracy, all_labels, all_predictions
+
+# Load the data
+device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+TEXT, LABEL, train_iterator, valid_iterator, test_iterator = load_data()
+
+# Streamlit interface
+st.title("RNN for Text Classification on IMDb Dataset")
+
+st.write("""
+This application demonstrates how to build and train a Recurrent Neural Network (RNN) for text classification using the IMDb dataset. You can adjust hyperparameters, visualize sample data, and see the model's performance.
+""")
+
+# Sidebar for input parameters
+st.sidebar.header('Model Hyperparameters')
+embed_size = st.sidebar.slider('Embedding Size', 50, 300, 100)
+hidden_size = st.sidebar.slider('Hidden Size', 50, 300, 256)
+n_layers = st.sidebar.slider('Number of RNN Layers', 1, 3, 2)
+dropout = st.sidebar.slider('Dropout', 0.0, 0.5, 0.2, step=0.1)
+learning_rate = st.sidebar.slider('Learning Rate', 0.001, 0.1, 0.01, step=0.001)
+epochs = st.sidebar.slider('Epochs', 1, 20, 5)
+
+# Create the network
+vocab_size = len(TEXT.vocab)
+output_size = 1
+net = RNN(vocab_size, embed_size, hidden_size, output_size, n_layers, dropout).to(device)
+criterion = nn.BCEWithLogitsLoss()
+optimizer = optim.Adam(net.parameters(), lr=learning_rate)
+
+# Add vertical space
+st.write('\n' * 10)
+
+# Train the network
+if st.sidebar.button('Train Network'):
+    loss_values = train_network(net, train_iterator, optimizer, criterion, epochs)
+    
+    # Plot the loss values
+    plt.figure(figsize=(10, 5))
+    plt.plot(range(1, epochs + 1), loss_values, marker='o')
+    plt.title('Training Loss Over Epochs')
+    plt.xlabel('Epoch')
+    plt.ylabel('Loss')
+    plt.grid(True)
+    st.pyplot(plt)
+    
+    # Store the trained model in the session state
+    st.session_state['trained_model'] = net
+
+# Test the network
+if 'trained_model' in st.session_state and st.sidebar.button('Test Network'):
+    accuracy, all_labels, all_predictions = evaluate_network(st.session_state['trained_model'], test_iterator, criterion)
+    st.write(f'Test Accuracy: {accuracy:.2f}%')
+    
+    # Display results in a table
+    st.write('Ground Truth vs Predicted')
+    results = pd.DataFrame({
+        'Ground Truth': all_labels,
+        'Predicted': all_predictions
+    })
+    st.table(results.head(50))  # Display first 50 results for brevity
+
+# Visualize some test results
+def visualize_text_predictions(iterator, net):
+    net.eval()
+    samples = []
+    with torch.no_grad():
+        for batch in iterator:
+            predictions = torch.round(torch.sigmoid(net(batch.text).squeeze(1)))
+            samples.extend(zip(batch.text.cpu(), batch.label.cpu(), predictions.cpu()))
+            if len(samples) >= 10:
+                break
+    return samples[:10]
+
+if 'trained_model' in st.session_state and st.sidebar.button('Show Test Results'):
+    samples = visualize_text_predictions(test_iterator, st.session_state['trained_model'])
+    st.write('Ground Truth vs Predicted for Sample Texts')
+    for i, (text, true_label, predicted) in enumerate(samples):
+        st.write(f'Sample {i+1}')
+        st.text(' '.join([TEXT.vocab.itos[token] for token in text]))
+        st.write(f'Ground Truth: {true_label.item()}, Predicted: {predicted.item()}')