Update app.py

app.py

@@ -264,11 +264,8 @@ class SpaceInferenceEngine:
                 generated_text += token_str
                 yield generated_text
 
-                # Optional: Stop generation if needed
-                # if new_token_id == self.tokenizer.token_to_id("<EOS>"): break
-
 # ============================================================================
-# 3. GRADIO INTERFACE
+# 3. GRADIO INTERFACE (Using Blocks for robustness)
 # ============================================================================
 
 # Initialize engine globally
@@ -288,29 +285,53 @@ def predict(prompt, max_tokens, temperature, top_k):
     ):
         yield current_text
 
-# Custom CSS for a cleaner look
+# Custom CSS
 custom_css = """
-#component-0 {max_width: 800px; margin: auto;}
+#container {max_width: 800px; margin: auto;}
+.output-text {min_height: 300px;}
 """
 
-with gr.Interface(
-    fn=predict,
-    inputs=[
-        gr.Textbox(lines=3, placeholder="Enter your prompt here...", label="Input Prompt"),
-        gr.Slider(minimum=10, maximum=512, value=150, step=10, label="Max New Tokens"),
-        gr.Slider(minimum=0.1, maximum=2.0, value=0.8, step=0.1, label="Temperature"),
-        gr.Slider(minimum=1, maximum=100, value=40, step=1, label="Top-K"),
-    ],
-    outputs=gr.Textbox(lines=10, label="Generated Output"),
-    title="i3-200m-v2 (RWKV-Hybrid)",
-    description="A 200M parameter hybrid model combining RWKV (RNN) and Standard Attention layers.",
-    css=custom_css,
-    examples=[
-        ["The history of science is"],
-        ["Once upon a time in a digital world,"],
-        ["The quick brown fox jumps over"]
-    ],
-    cache_examples=False 
-) as demo:
-    demo.queue() # Enable queuing for streaming
+with gr.Blocks(css=custom_css) as demo:
+    with gr.Column(elem_id="container"):
+        gr.Markdown("# i3-200m-v2 (RWKV-Hybrid)")
+        gr.Markdown("A 200M parameter hybrid model combining RWKV (RNN) and Standard Attention layers.")
+        
+        with gr.Group():
+            prompt_input = gr.Textbox(
+                lines=3, 
+                placeholder="Enter your prompt here...", 
+                label="Input Prompt"
+            )
+            with gr.Row():
+                max_tokens = gr.Slider(minimum=10, maximum=512, value=150, step=10, label="Max New Tokens")
+                temp = gr.Slider(minimum=0.1, maximum=2.0, value=0.8, step=0.1, label="Temperature")
+                topk = gr.Slider(minimum=1, maximum=100, value=40, step=1, label="Top-K")
+            
+            submit_btn = gr.Button("Generate", variant="primary")
+        
+        output_text = gr.Textbox(
+            lines=10, 
+            label="Generated Output", 
+            elem_classes="output-text"
+        )
+
+        # Set up event listeners
+        submit_btn.click(
+            fn=predict,
+            inputs=[prompt_input, max_tokens, temp, topk],
+            outputs=output_text
+        )
+        
+        # Example inputs to easily test
+        gr.Examples(
+            examples=[
+                ["The history of science is"],
+                ["Once upon a time in a digital world,"],
+                ["The quick brown fox jumps over"]
+            ],
+            inputs=prompt_input
+        )
+
+if __name__ == "__main__":
+    demo.queue()
     demo.launch()