Update pages/19_RNN_Shakespeare.py

pages/19_RNN_Shakespeare.py

@@ -51,58 +51,64 @@ generate_length = st.number_input("Generated text length:", min_value=50, value=
 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.array(dataY)
-
-    # Convert to PyTorch tensors
-    X_tensor = torch.tensor(X, dtype=torch.float32)
-    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)
-
-    # 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, 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)
-            h = (h[0].detach(), h[1].detach())
-            loss = criterion(outputs, targets)
-
-            # Backward pass and optimization
-            optimizer.zero_grad()
-            loss.backward()
-            optimizer.step()
-
-            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)
-    st.subheader("Generated Text")
-    st.write(generated_text)
+    if len(text) <= seq_length:
+        st.error("Text data is too short for the given sequence length. Please enter more text data.")
+    else:
+        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])
+
+        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, 1))
+            X = X / float(len(chars))
+            Y = np.array(dataY)
+
+            # Convert to PyTorch tensors
+            X_tensor = torch.tensor(X, dtype=torch.float32)
+            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)
+
+            # 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, 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)
+                    h = (h[0].detach(), h[1].detach())
+                    loss = criterion(outputs, targets)
+
+                    # Backward pass and optimization
+                    optimizer.zero_grad()
+                    loss.backward()
+                    optimizer.step()
+
+                    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)
+            st.subheader("Generated Text")
+            st.write(generated_text)