Update pages/19_RNN_Shakespeare.py

pages/19_RNN_Shakespeare.py

@@ -66,11 +66,11 @@ if st.button("Train and Generate"):
 
     X = np.reshape(dataX, (len(dataX), seq_length, 1))
     X = X / float(len(chars))
-    Y = np.eye(len(chars))[dataY]
+    Y = np.array(dataY)
 
     # Convert to PyTorch tensors
     X_tensor = torch.tensor(X, dtype=torch.float32)
-    Y_tensor = torch.tensor(dataY, dtype=torch.long)
+    Y_tensor = torch.tensor(Y, dtype=torch.long)
 
     # Model initialization
     model = LSTMModel(input_size=1, hidden_size=hidden_size, output_size=len(chars), num_layers=num_layers)
@@ -81,10 +81,11 @@ if st.button("Train and Generate"):
 
     # 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]
+        h = (torch.zeros(num_layers, 1, hidden_size), torch.zeros(num_layers, 1, hidden_size))
+        epoch_loss = 0
+        for i in range(len(dataX)):
+            inputs = X_tensor[i].unsqueeze(0)
+            targets = Y_tensor[i].unsqueeze(0)
 
             # Forward pass
             outputs, h = model(inputs, h)
@@ -96,7 +97,10 @@ if st.button("Train and Generate"):
             loss.backward()
             optimizer.step()
 
-        st.write(f'Epoch [{epoch + 1}/{num_epochs}], Loss: {loss.item():.4f}')
+            epoch_loss += loss.item()
+
+        avg_loss = epoch_loss / len(dataX)
+        st.write(f'Epoch [{epoch + 1}/{num_epochs}], Loss: {avg_loss:.4f}')
 
     # Text generation
     generated_text = generate_text(model, start_string, generate_length, char_to_int, int_to_char, num_layers, hidden_size)