Update pages/17_RNN.py

pages/17_RNN.py

@@ -10,7 +10,7 @@ data_size, vocab_size = len(sequence), len(chars)
 
 # Create mappings from characters to indices and vice versa
 char_to_idx = {ch: i for i, ch in enumerate(chars)}
-idx_to_char = {i: ch for i, ch in enumerate(chars)}
+idx_to_char = {i, ch for i, ch in enumerate(chars)}
 
 # Convert the sequence to indices
 indices = np.array([char_to_idx[ch] for ch in sequence])
@@ -44,6 +44,8 @@ criterion = nn.NLLLoss()
 optimizer = torch.optim.Adam(rnn.parameters(), lr=learning_rate)
 
 def char_tensor(char):
+    if char not in char_to_idx:
+        raise ValueError(f"Character '{char}' not in vocabulary.")
     tensor = torch.zeros(1, vocab_size)
     tensor[0][char_to_idx[char]] = 1
     return tensor
@@ -70,6 +72,8 @@ for epoch in range(n_epochs):
 print("Training complete.")
 
 def generate(start_char, predict_len=100):
+    if start_char not in char_to_idx:
+        raise ValueError(f"Start character '{start_char}' not in vocabulary.")
     hidden = rnn.init_hidden()
     input_char = char_tensor(start_char)
     predicted_str = start_char
@@ -92,6 +96,9 @@ 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)
 
 if st.button('Generate Text'):
-    generated_text = generate(start_char, predict_len)
-    st.write('Generated Text:')
-    st.text(generated_text)
+    try:
+        generated_text = generate(start_char, predict_len)
+        st.write('Generated Text:')
+        st.text(generated_text)
+    except ValueError as e:
+        st.error(str(e))