Update pages/19_RNN_LSTM_Shakespeare.py

pages/19_RNN_LSTM_Shakespeare.py

@@ -22,7 +22,7 @@ class LSTMModel(nn.Module):
 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.long).unsqueeze(0)
+    input_seq = torch.tensor(input_seq, dtype=torch.long).unsqueeze(0).unsqueeze(-1)
     h = (torch.zeros(num_layers, 1, hidden_size), torch.zeros(num_layers, 1, hidden_size))
     generated_text = start_str
 
@@ -31,7 +31,7 @@ def generate_text(model, start_str, length, char_to_int, int_to_char, num_layers
         _, 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.long)
+        input_seq = torch.tensor([[char_to_int[predicted_char]]], dtype=torch.long).unsqueeze(0).unsqueeze(-1)
 
     return generated_text
 
@@ -39,7 +39,7 @@ def generate_text(model, start_str, length, char_to_int, int_to_char, num_layers
 st.title("RNN/LSTM Text Generation")
 
 # Inputs
-text_data = st.text_area("Enter your text data for training:")
+text_data = st.text_area("Enter your text data for training:", "To be, or not to be, that is the question:\nWhether 'tis nobler in the mind to suffer\nThe slings and arrows of outrageous fortune,\nOr to take arms against a sea of troubles\nAnd by opposing end them. To die: to sleep;\nNo more; and by a sleep to say we end\nThe heart-ache and the thousand natural shocks\nThat flesh is heir to, 'tis a consummation\nDevoutly to be wish'd. To die, to sleep;\nTo sleep: perchance to dream: ay, there's the rub;\nFor in that sleep of death what dreams may come\nWhen we have shuffled off this mortal coil,\nMust give us pause: there's the respect\nThat makes calamity of so long life;")
 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)
@@ -70,7 +70,7 @@ if st.button("Train and Generate"):
         if len(dataX) == 0:
             st.error("Not enough data to create input-output pairs. Please provide more text data.")
         else:
-            X = np.reshape(dataX, (len(dataX), seq_length))
+            X = np.reshape(dataX, (len(dataX), seq_length, 1))
             X = X / float(len(chars))
             Y = np.array(dataY)
 
@@ -87,15 +87,15 @@ if st.button("Train and Generate"):
 
             # Training the model
             for epoch in range(num_epochs):
-                h = (torch.zeros(num_layers, 1, hidden_size), torch.zeros(num_layers, 1, hidden_size))
+                h = (torch.zeros(num_layers, X_tensor.size(0), hidden_size), torch.zeros(num_layers, X_tensor.size(0), hidden_size))
                 epoch_loss = 0
                 for i in range(len(dataX)):
-                    inputs = X_tensor[i].unsqueeze(0).unsqueeze(-1)
+                    inputs = X_tensor[i].unsqueeze(0)
                     targets = Y_tensor[i].unsqueeze(0)
 
                     # Forward pass
-                    outputs, h = model(inputs, (h[0].squeeze(1), h[1].squeeze(1)))
-                    h = (h[0].detach().unsqueeze(1), h[1].detach().unsqueeze(1))
+                    outputs, h = model(inputs, h)
+                    h = (h[0].detach(), h[1].detach())
                     loss = criterion(outputs, targets)
 
                     # Backward pass and optimization