Update app.py

app.py

@@ -1,49 +1,60 @@
 import gradio as gr
 
-def make_html(training_text, iterations, gen_length):
-    # Sanitize text for JS (escape quotes and backslashes)
+def make_html(training_text, iterations, gen_length, layer1_size, layer2_size):
+    # Sanitize training text for JS
     safe_text = training_text.replace("\\", "\\\\").replace("\"", "\\\"").replace("\n", "\\n")
-    
+
     html = f"""
     <div>
-      <h3>Output:</h3>
-      <pre id="output">Training...</pre>
+      <h3>Live LSTM Training Output (updates every 10s)</h3>
+      <pre id="output">Starting training...</pre>
       <script src="https://karpathy.github.io/recurrentjs/recurrent.js"></script>
       <script>
         const text = "{safe_text}".split("");
-        const iterations = {iterations};
+        const totalIterations = {iterations};
         const genLength = {gen_length};
 
-        const lstm = new RNN("lstm", {{ hiddenSizes: [128] }});
+        const lstm = new RNN("lstm", {{ hiddenSizes: [{layer1_size}, {layer2_size}] }});
         const trainer = new RNNTrainer(lstm, {{ learningRate: 0.01, momentum: 0.1, batchSize: 5 }});
 
-        for (let i = 0; i < iterations; i++) {{
-          const idx = Math.floor(Math.random() * (text.length - 10));
-          const input = text.slice(idx, idx + 5);
-          const output = text.slice(idx + 1, idx + 6);
-          trainer.train(input, output);
-        }}
+        let iteration = 0;
+        const interval = setInterval(() => {{
+          for (let i = 0; i < 5 && iteration < totalIterations; i++, iteration++) {{
+            const idx = Math.floor(Math.random() * (text.length - 10));
+            const input = text.slice(idx, idx + 5);
+            const output = text.slice(idx + 1, idx + 6);
+            trainer.train(input, output);
+          }}
 
-        const generated = lstm.sample(["H"], genLength).join("");
-        document.getElementById("output").innerText = generated;
+          const sample = lstm.sample(["H"], genLength).join("");
+          document.getElementById("output").innerText = `Epochs completed: ${iteration} / ${totalIterations}\\n\\n` + sample;
+
+          if (iteration >= totalIterations) {{
+            clearInterval(interval);
+            document.getElementById("output").innerText += "\\n\\nTraining complete!";
+          }}
+        }}, 10000); // every 10 seconds
       </script>
     </div>
     """
     return html
 
 with gr.Blocks() as demo:
-    gr.Markdown("## 🧠 LSTM Trainer with recurrent.js (Browser-Based)")
+    gr.Markdown("## 🧠 Live LSTM Trainer with recurrent.js (Browser-Based)")
     
+    training_text = gr.Textbox(label="Training Text", lines=6, placeholder="Paste your training data here...")
+
     with gr.Row():
-        training_text = gr.Textbox(label="Training Text", lines=8, placeholder="Paste your training text here...")
-    
+        iterations = gr.Slider(10, 200, value=50, step=10, label="Total Training Epochs")
+        gen_length = gr.Slider(20, 500, value=100, step=10, label="Characters to Generate")
+
     with gr.Row():
-        iterations = gr.Slider(1, 100, value=20, step=1, label="Training Iterations")
-        gen_length = gr.Slider(10, 500, value=100, step=10, label="Characters to Generate")
-    
-    run_button = gr.Button("Train and Generate")
+        layer1_size = gr.Slider(16, 256, value=128, step=16, label="Neurons in Layer 1")
+        layer2_size = gr.Slider(16, 256, value=128, step=16, label="Neurons in Layer 2")
+
+    run_button = gr.Button("Start Training")
     output_html = gr.HTML()
 
-    run_button.click(fn=make_html, inputs=[training_text, iterations, gen_length], outputs=output_html)
+    run_button.click(fn=make_html, inputs=[training_text, iterations, gen_length, layer1_size, layer2_size], outputs=output_html)
 
 demo.launch()