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()}')