Update pages/17_RNN_News.py

pages/17_RNN_News.py

@@ -2,7 +2,10 @@ import streamlit as st
 import torch
 import torch.nn as nn
 import torch.optim as optim
-from torchtext.legacy import data, datasets
+from torchtext.data.utils import get_tokenizer
+from torchtext.vocab import build_vocab_from_iterator
+from torchtext.datasets import AG_NEWS
+from torch.utils.data import DataLoader, random_split
 import matplotlib.pyplot as plt
 import pandas as pd
 import numpy as np
@@ -23,38 +26,62 @@ class RNN(nn.Module):
         out = self.fc(out[:, -1, :])
         return out
 
-# Load the data
-@st.cache(allow_output_mutation=True)
+# Create a custom collate function to pad sequences
+def collate_batch(batch):
+    label_list, text_list, lengths = [], [], []
+    for _label, _text in batch:
+        label_list.append(label_pipeline(_label))
+        processed_text = torch.tensor(text_pipeline(_text), dtype=torch.int64)
+        text_list.append(processed_text)
+        lengths.append(processed_text.size(0))
+    labels = torch.tensor(label_list, dtype=torch.int64)
+    texts = pad_sequence(text_list, batch_first=True, padding_value=vocab["<pad>"])
+    return texts, labels
+
+# Function to load the data
+@st.cache_data
 def load_data():
-    TEXT = data.Field(tokenize='spacy', tokenizer_language='en_core_web_sm', include_lengths=True)
-    LABEL = data.LabelField(dtype=torch.long)
-    train_data, test_data = datasets.AG_NEWS.splits(TEXT, LABEL)
-    train_data, valid_data = train_data.split(split_ratio=0.8)
+    tokenizer = get_tokenizer("basic_english")
+    train_iter = AG_NEWS(split='train')
+    test_iter = AG_NEWS(split='test')
 
-    TEXT.build_vocab(train_data, max_size=25000, vectors="glove.6B.100d", unk_init=torch.Tensor.normal_)
-    LABEL.build_vocab(train_data)
+    def yield_tokens(data_iter):
+        for _, text in data_iter:
+            yield tokenizer(text)
 
-    BATCH_SIZE = 64
+    vocab = build_vocab_from_iterator(yield_tokens(train_iter), specials=["<unk>", "<pad>"])
+    vocab.set_default_index(vocab["<unk>"])
 
-    train_iterator, valid_iterator, test_iterator = data.BucketIterator.splits(
-        (train_data, valid_data, test_data),
-        batch_size=BATCH_SIZE,
-        sort_within_batch=True,
-        device=device)
+    global text_pipeline, label_pipeline
+    text_pipeline = lambda x: vocab(tokenizer(x))
+    label_pipeline = lambda x: int(x) - 1
 
-    return TEXT, LABEL, train_iterator, valid_iterator, test_iterator
+    # Create DataLoaders
+    train_dataset = list(train_iter)
+    test_dataset = list(test_iter)
 
-# Train the network
+    train_size = int(0.8 * len(train_dataset))
+    valid_size = len(train_dataset) - train_size
+    train_dataset, valid_dataset = random_split(train_dataset, [train_size, valid_size])
+
+    BATCH_SIZE = 64
+    train_loader = DataLoader(train_dataset, batch_size=BATCH_SIZE, shuffle=True, collate_fn=collate_batch)
+    valid_loader = DataLoader(valid_dataset, batch_size=BATCH_SIZE, shuffle=True, collate_fn=collate_batch)
+    test_loader = DataLoader(test_dataset, batch_size=BATCH_SIZE, shuffle=True, collate_fn=collate_batch)
+
+    return vocab, train_loader, valid_loader, test_loader
+
+# 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:
+        for texts, labels in iterator:
+            texts, labels = texts.to(device), labels.to(device)
             optimizer.zero_grad()
-            text, text_lengths = batch.text
-            predictions = net(text).squeeze(1)
-            loss = criterion(predictions, batch.label)
+            predictions = net(texts).squeeze(1)
+            loss = criterion(predictions, labels)
             loss.backward()
             optimizer.step()
             epoch_loss += loss.item()
@@ -64,7 +91,7 @@ def train_network(net, iterator, optimizer, criterion, epochs):
     st.write('Finished Training')
     return loss_values
 
-# Evaluate the network
+# Function to evaluate the network
 def evaluate_network(net, iterator, criterion):
     epoch_loss = 0
     correct = 0
@@ -73,15 +100,15 @@ def evaluate_network(net, iterator, criterion):
     all_predictions = []
     net.eval()
     with torch.no_grad():
-        for batch in iterator:
-            text, text_lengths = batch.text
-            predictions = net(text).squeeze(1)
-            loss = criterion(predictions, batch.label)
+        for texts, labels in iterator:
+            texts, labels = texts.to(device), labels.to(device)
+            predictions = net(texts).squeeze(1)
+            loss = criterion(predictions, labels)
             epoch_loss += loss.item()
             _, predicted = torch.max(predictions, 1)
-            correct += (predicted == batch.label).sum().item()
-            total += len(batch.label)
-            all_labels.extend(batch.label.cpu().numpy())
+            correct += (predicted == labels).sum().item()
+            total += len(labels)
+            all_labels.extend(labels.cpu().numpy())
             all_predictions.extend(predicted.cpu().numpy())
     accuracy = 100 * correct / total
     st.write(f'Loss: {epoch_loss / len(iterator):.4f}, Accuracy: {accuracy:.2f}%')
@@ -89,7 +116,7 @@ def evaluate_network(net, iterator, criterion):
 
 # Load data
 device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
-TEXT, LABEL, train_iterator, valid_iterator, test_iterator = load_data()
+vocab, train_loader, valid_loader, test_loader = load_data()
 
 # Streamlit interface
 st.title("RNN for Text Classification on AG News Dataset")
@@ -108,8 +135,8 @@ 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 = len(LABEL.vocab)
+vocab_size = len(vocab)
+output_size = 4  # Number of classes in AG_NEWS
 net = RNN(vocab_size, embed_size, hidden_size, output_size, n_layers, dropout).to(device)
 criterion = nn.CrossEntropyLoss()
 optimizer = optim.Adam(net.parameters(), lr=learning_rate)
@@ -119,7 +146,7 @@ st.write('\n' * 10)
 
 # Train the network
 if st.sidebar.button('Train Network'):
-    loss_values = train_network(net, train_iterator, optimizer, criterion, epochs)
+    loss_values = train_network(net, train_loader, optimizer, criterion, epochs)
     
     # Plot the loss values
     plt.figure(figsize=(10, 5))
@@ -135,7 +162,7 @@ if st.sidebar.button('Train Network'):
 
 # 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)
+    accuracy, all_labels, all_predictions = evaluate_network(st.session_state['trained_model'], test_loader, criterion)
     st.write(f'Test Accuracy: {accuracy:.2f}%')
     
     # Display results in a table
@@ -151,18 +178,17 @@ def visualize_text_predictions(iterator, net):
     net.eval()
     samples = []
     with torch.no_grad():
-        for batch in iterator:
-            text, text_lengths = batch.text
-            predictions = torch.max(net(text), 1)[1]
-            samples.extend(zip(text.cpu(), batch.label.cpu(), predictions.cpu()))
+        for texts, labels in iterator:
+            predictions = torch.max(net(texts), 1)[1]
+            samples.extend(zip(texts.cpu(), labels.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'])
+    samples = visualize_text_predictions(test_loader, 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.text(' '.join([vocab.get_itos()[token] for token in text]))
         st.write(f'Ground Truth: {LABEL.vocab.itos[true_label.item()]}, Predicted: {LABEL.vocab.itos[predicted.item()]}')