Create 19_RNN_Shakespeare.py

pages/19_RNN_Shakespeare.py

@@ -0,0 +1,104 @@
+import streamlit as st
+import torch
+import torch.nn as nn
+import torch.optim as optim
+import numpy as np
+
+# Define the RNN or LSTM Model
+class LSTMModel(nn.Module):
+    def __init__(self, input_size, hidden_size, output_size, num_layers):
+        super(LSTMModel, self).__init__()
+        self.hidden_size = hidden_size
+        self.num_layers = num_layers
+        self.lstm = nn.LSTM(input_size, hidden_size, num_layers, batch_first=True)
+        self.fc = nn.Linear(hidden_size, output_size)
+
+    def forward(self, x, h):
+        out, h = self.lstm(x, h)
+        out = self.fc(out[:, -1, :])
+        return out, h
+
+# Text generation function
+def generate_text(model, start_str, length, char_to_int, int_to_char, num_layers, hidden_size):
+    model.eval()
+    input_seq = [char_to_int[c] for c in start_str]
+    input_seq = torch.tensor(input_seq, dtype=torch.float32).unsqueeze(0).unsqueeze(-1)
+    h = (torch.zeros(num_layers, 1, hidden_size), torch.zeros(num_layers, 1, hidden_size))
+    generated_text = start_str
+
+    for _ in range(length):
+        output, h = model(input_seq, h)
+        _, predicted = torch.max(output, 1)
+        predicted_char = int_to_char[predicted.item()]
+        generated_text += predicted_char
+        input_seq = torch.tensor([[char_to_int[predicted_char]]], dtype=torch.float32).unsqueeze(0).unsqueeze(-1)
+
+    return generated_text
+
+# Streamlit interface
+st.title("RNN/LSTM Text Generation")
+
+# Inputs
+text_data = st.text_area("Enter your text data for training:")
+start_string = st.text_input("Enter the start string for text generation:")
+seq_length = st.number_input("Sequence length:", min_value=10, value=100)
+hidden_size = st.number_input("Hidden size:", min_value=50, value=256)
+num_layers = st.number_input("Number of layers:", min_value=1, value=2)
+learning_rate = st.number_input("Learning rate:", min_value=0.0001, value=0.003, format="%.4f")
+num_epochs = st.number_input("Number of epochs:", min_value=1, value=20)
+generate_length = st.number_input("Generated text length:", min_value=50, value=500)
+
+if st.button("Train and Generate"):
+    # Data Preparation
+    text = text_data
+    chars = sorted(list(set(text)))
+    char_to_int = {c: i for i, c in enumerate(chars)}
+    int_to_char = {i: c for i, c in enumerate(chars)}
+
+    # Prepare input-output pairs
+    dataX = []
+    dataY = []
+    for i in range(0, len(text) - seq_length):
+        seq_in = text[i:i + seq_length]
+        seq_out = text[i + seq_length]
+        dataX.append([char_to_int[char] for char in seq_in])
+        dataY.append(char_to_int[seq_out])
+
+    X = np.reshape(dataX, (len(dataX), seq_length, 1))
+    X = X / float(len(chars))
+    Y = np.eye(len(chars))[dataY]
+
+    # Convert to PyTorch tensors
+    X_tensor = torch.tensor(X, dtype=torch.float32)
+    Y_tensor = torch.tensor(dataY, dtype=torch.long)
+
+    # Model initialization
+    model = LSTMModel(input_size=1, hidden_size=hidden_size, output_size=len(chars), num_layers=num_layers)
+
+    # Loss and optimizer
+    criterion = nn.CrossEntropyLoss()
+    optimizer = optim.Adam(model.parameters(), lr=learning_rate)
+
+    # Training the model
+    for epoch in range(num_epochs):
+        h = (torch.zeros(num_layers, X_tensor.size(0), hidden_size), torch.zeros(num_layers, X_tensor.size(0), hidden_size))
+        for i in range(0, X_tensor.size(0), seq_length):
+            inputs = X_tensor[i:i + seq_length]
+            targets = Y_tensor[i:i + seq_length]
+
+            # Forward pass
+            outputs, h = model(inputs, h)
+            h = (h[0].detach(), h[1].detach())
+            loss = criterion(outputs, targets)
+
+            # Backward pass and optimization
+            optimizer.zero_grad()
+            loss.backward()
+            optimizer.step()
+
+        st.write(f'Epoch [{epoch + 1}/{num_epochs}], Loss: {loss.item():.4f}')
+
+    # Text generation
+    generated_text = generate_text(model, start_string, generate_length, char_to_int, int_to_char, num_layers, hidden_size)
+    st.subheader("Generated Text")
+    st.write(generated_text)