Update pages/17_RNN.py

pages/17_RNN.py

@@ -1,205 +1,97 @@
 import streamlit as st
 import torch
 import torch.nn as nn
-import torch.optim as optim
-from torchtext.data.utils import get_tokenizer
-from torchtext.vocab import build_vocab_from_iterator
-from torchtext.datasets import IMDB
-from torch.utils.data import DataLoader, random_split
-import matplotlib.pyplot as plt
-import seaborn as sns
-import pandas as pd
 import numpy as np
-from collections import Counter
-from torch.nn.utils.rnn import pad_sequence
 
-# Define the RNN model
+# Define the dataset
+sequence = "hellohellohello"
+chars = list(set(sequence))
+data_size, vocab_size = len(sequence), len(chars)
+
+# Create mappings from characters to indices and vice versa
+char_to_idx = {ch: i for i, ch in enumerate(chars)}
+idx_to_char = {i: ch for i, ch in enumerate(chars)}
+
+# Convert the sequence to indices
+indices = np.array([char_to_idx[ch] for ch in sequence])
+
 class RNN(nn.Module):
-    def __init__(self, vocab_size, embed_size, hidden_size, output_size, n_layers, dropout):
+    def __init__(self, input_size, hidden_size, output_size):
         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
-
-# Create a custom collate function to pad sequences
-def collate_batch(batch):
-    texts, labels = zip(*batch)
-    text_lengths = [len(text) for text in texts]
-    texts_padded = pad_sequence(texts, batch_first=True, padding_value=vocab["<pad>"])
-    return texts_padded, torch.tensor(labels, dtype=torch.float), text_lengths
-
-# Function to load the data
-@st.cache_data
-def load_data():
-    tokenizer = get_tokenizer("basic_english")
-    train_iter, test_iter = IMDB(split=('train', 'test'))
-
-    def yield_tokens(data_iter):
-        for _, text in data_iter:
-            yield tokenizer(text)
-    
-    vocab = build_vocab_from_iterator(yield_tokens(train_iter), specials=["<unk>", "<pad>"])
-    vocab.set_default_index(vocab["<unk>"])
-
-    # Define the text and label processing pipelines
-    text_pipeline = lambda x: vocab(tokenizer(x))
-    label_pipeline = lambda x: 1 if x == 'pos' else 0
-
-    # Process the data into tensors
-    def process_data(data_iter):
-        texts, labels = [], []
-        for label, text in data_iter:
-            texts.append(torch.tensor(text_pipeline(text), dtype=torch.long))
-            labels.append(label_pipeline(label))
-        return texts, torch.tensor(labels, dtype=torch.float)
-
-    train_texts, train_labels = process_data(train_iter)
-    test_texts, test_labels = process_data(test_iter)
-
-    # Create DataLoaders
-    train_dataset = list(zip(train_texts, train_labels))
-    test_dataset = list(zip(test_texts, test_labels))
-
-    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 texts, labels, _ in iterator:
-            texts, labels = texts.to(device), labels.to(device)
-            optimizer.zero_grad()
-            predictions = net(texts).squeeze(1)
-            loss = criterion(predictions, labels)
-            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 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()
-            rounded_preds = torch.round(torch.sigmoid(predictions))
-            correct += (rounded_preds == labels).sum().item()
-            total += len(labels)
-            all_labels.extend(labels.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')
-
-# Display a loading message with some vertical space
-st.markdown("<div style='margin-top: 50px;'><b>Loading data...</b></div>", unsafe_allow_html=True)
-vocab, train_loader, valid_loader, test_loader = load_data()
+        self.hidden_size = hidden_size
+        self.i2h = nn.Linear(input_size + hidden_size, hidden_size)
+        self.i2o = nn.Linear(input_size + hidden_size, output_size)
+        self.softmax = nn.LogSoftmax(dim=1)
+
+    def forward(self, input, hidden):
+        combined = torch.cat((input, hidden), 1)
+        hidden = self.i2h(combined)
+        output = self.i2o(combined)
+        output = self.softmax(output)
+        return output, hidden
+
+    def init_hidden(self):
+        return torch.zeros(1, self.hidden_size)
+
+# Hyperparameters
+n_hidden = 128
+learning_rate = 0.005
+n_epochs = 500
+
+# Initialize the model, loss function, and optimizer
+rnn = RNN(vocab_size, n_hidden, vocab_size)
+criterion = nn.NLLLoss()
+optimizer = torch.optim.Adam(rnn.parameters(), lr=learning_rate)
+
+def char_tensor(char):
+    tensor = torch.zeros(1, vocab_size)
+    tensor[0][char_to_idx[char]] = 1
+    return tensor
+
+# Training loop
+for epoch in range(n_epochs):
+    hidden = rnn.init_hidden()
+    rnn.zero_grad()
+    loss = 0
+
+    for i in range(data_size - 1):
+        input_char = char_tensor(sequence[i])
+        target_char = torch.tensor([char_to_idx[sequence[i + 1]]], dtype=torch.long)
+        
+        output, hidden = rnn(input_char, hidden)
+        loss += criterion(output, target_char)
+
+    loss.backward()
+    optimizer.step()
+
+    if epoch % 10 == 0:
+        print(f'Epoch {epoch} loss: {loss.item() / (data_size - 1)}')
+
+print("Training complete.")
+
+def generate(start_char, predict_len=100):
+    hidden = rnn.init_hidden()
+    input_char = char_tensor(start_char)
+    predicted_str = start_char
+
+    for _ in range(predict_len):
+        output, hidden = rnn(input_char, hidden)
+        topv, topi = output.topk(1)
+        predicted_char_idx = topi[0][0].item()
+        predicted_char = idx_to_char[predicted_char_idx]
+        predicted_str += predicted_char
+        input_char = char_tensor(predicted_char)
+
+    return predicted_str
 
 # 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(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_loader, 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_loader, 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 texts, labels, _ in iterator:
-            predictions = torch.round(torch.sigmoid(net(texts).squeeze(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_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([vocab.get_itos()[token] for token in text]))
-        st.write(f'Ground Truth: {true_label.item()}, Predicted: {predicted.item()}')
+st.title('RNN Character Prediction')
+st.write('This app uses a Recurrent Neural Network (RNN) to predict the next character in a given string.')
+
+start_char = st.text_input('Enter a starting character:', 'h')
+predict_len = st.slider('Select the length of the generated text:', min_value=10, max_value=200, value=50)
+
+if st.button('Generate Text'):
+    generated_text = generate(start_char, predict_len)
+    st.write('Generated Text:')
+    st.text(generated_text)