Update pages/17_RNN.py

pages/17_RNN.py

@@ -3,6 +3,14 @@ import torch
 import torch.nn as nn
 import numpy as np
 
+# Introduction
+st.title('RNN Character Prediction')
+st.write("""
+This app demonstrates how to train a Recurrent Neural Network (RNN) to predict the next character in a given string. 
+The RNN learns the sequence of characters from a provided text and generates text based on the learned patterns.
+You can choose different options for training the RNN and see how it affects the generated text.
+""")
+
 # Define the dataset
 sequence = "In the vast expanse of the digital realm."
 chars = list(set(sequence))
@@ -36,11 +44,20 @@ class RNN(nn.Module):
 # 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()
+
+# Define training options
+options = {
+    'Quick Train (100 epochs)': 100,
+    'Medium Train (500 epochs)': 500,
+    'Long Train (1000 epochs)': 1000
+}
+
+train_option = st.selectbox('Select training option:', list(options.keys()))
+n_epochs = options[train_option]
 optimizer = torch.optim.Adam(rnn.parameters(), lr=learning_rate)
 
 def char_tensor(char):
@@ -50,26 +67,31 @@ def char_tensor(char):
     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
+# Training function
+def train_model(n_epochs):
+    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)
+        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()
 
-    loss.backward()
-    optimizer.step()
+        if epoch % 10 == 0:
+            st.write(f'Epoch {epoch} loss: {loss.item() / (data_size - 1)}')
 
-    if epoch % 10 == 0:
-        print(f'Epoch {epoch} loss: {loss.item() / (data_size - 1)}')
+    st.write("Training complete.")
 
-print("Training complete.")
+# Train the model
+if st.button('Train Model'):
+    train_model(n_epochs)
 
 def generate(start_char, predict_len=100):
     if start_char not in char_to_idx:
@@ -88,10 +110,6 @@ def generate(start_char, predict_len=100):
 
     return predicted_str
 
-# Streamlit interface
-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)