Update app.py

app.py

@@ -265,7 +265,7 @@ class SpaceInferenceEngine:
                 yield generated_text
 
 # ============================================================================
-# 3. GRADIO INTERFACE (Using Blocks for robustness)
+# 3. GRADIO INTERFACE (UI Upgrade)
 # ============================================================================
 
 # Initialize engine globally
@@ -273,8 +273,9 @@ print("Starting Engine...")
 engine = SpaceInferenceEngine()
 
 def predict(prompt, max_tokens, temperature, top_k):
-    if not prompt:
-        return "Please enter a prompt."
+    if not prompt.strip():
+        yield "⚠️ Please enter a prompt to generate text."
+        return
     
     # Use the generator for streaming
     for current_text in engine.generate_stream(
@@ -287,50 +288,130 @@ def predict(prompt, max_tokens, temperature, top_k):
 
 # Custom CSS
 custom_css = """
-#container {max_width: 800px; margin: auto;}
-.output-text {min_height: 300px;}
+.gradio-container {
+    max-width: 1200px !important;
+}
+.main-header {
+    text-align: center;
+    margin-bottom: 2rem;
+}
 """
 
-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():
+with gr.Blocks(css=custom_css, theme=gr.themes.Soft()) as demo:
+    # Header
+    with gr.Row():
+        gr.Markdown(
+            """
+            # 🚀 i3-200M Text Generation
+            ### Powered by RWKV-Hybrid Architecture
+            Generate creative text using the i3-200M language model combining RNN efficiency with Attention precision.
+            """,
+            elem_classes="main-header"
+        )
+    
+    # Main Generation Area
+    with gr.Row():
+        # Left Column: Inputs
+        with gr.Column(scale=2):
             prompt_input = gr.Textbox(
-                lines=3, 
-                placeholder="Enter your prompt here...", 
-                label="Input Prompt"
+                label="✍️ Enter Your Prompt",
+                placeholder="Once upon a time in a distant galaxy...",
+                lines=4,
+                max_lines=8
             )
-            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
-        )
+            with gr.Accordion("⚙️ Generation Parameters", open=True):
+                with gr.Row():
+                    max_tokens_input = gr.Slider(
+                        minimum=10, 
+                        maximum=512, 
+                        value=150, 
+                        step=10, 
+                        label="Max Tokens",
+                        info="Maximum number of tokens to generate"
+                    )
+                    temp_input = gr.Slider(
+                        minimum=0.1, 
+                        maximum=2.0, 
+                        value=0.8, 
+                        step=0.1, 
+                        label="Temperature",
+                        info="Higher = more creative, Lower = more focused"
+                    )
+                
+                topk_input = gr.Slider(
+                    minimum=1, 
+                    maximum=100, 
+                    value=40, 
+                    step=1, 
+                    label="Top-k Sampling",
+                    info="Number of top tokens to consider"
+                )
+            
+            with gr.Row():
+                generate_btn = gr.Button("🎨 Generate Text", variant="primary", size="lg")
+                clear_btn = gr.ClearButton(components=[prompt_input], value="🗑️ Clear", size="lg")
         
-        # Example inputs to easily test
+        # Right Column: Output
+        with gr.Column(scale=2):
+            output_text = gr.Textbox(
+                label="📝 Generated Output",
+                lines=12,
+                max_lines=20,
+                show_copy_button=True
+            )
+    
+    # Examples Section
+    with gr.Row():
         gr.Examples(
             examples=[
-                ["The history of science is"],
-                ["Once upon a time in a digital world,"],
-                ["The quick brown fox jumps over"]
+                ["The history of science is", 150, 0.7, 50],
+                ["In a world where technology and nature coexist", 200, 0.9, 40],
+                ["The scientist discovered something remarkable", 120, 0.8, 45],
             ],
-            inputs=prompt_input
+            inputs=[prompt_input, max_tokens_input, temp_input, topk_input],
+            label="💡 Try These Examples"
         )
+    
+    # Developer Panel
+    with gr.Accordion("🔧 Developer Info", open=False):
+        total_params = sum(p.numel() for p in engine.model.parameters())
+        
+        with gr.Row():
+            with gr.Column():
+                gr.Markdown(f"""
+                **Model Architecture:**
+                - **Model:** i3-200M Hybrid
+                - **Device:** {engine.device}
+                - **Vocab Size:** {engine.config['vocab_size']:,}
+                - **Parameters:** {total_params:,} ({total_params/1e6:.2f}M)
+                """)
+            
+            with gr.Column():
+                gr.Markdown(f"""
+                **Configuration:**
+                - **d_model:** {engine.config['d_model']}
+                - **RWKV Layers:** {engine.config['rwkv_layers']}
+                - **Attention Layers:** {engine.config['attn_layers']}
+                - **Max Seq Len:** {engine.model.max_seq_len}
+                """)
+    
+    # Footer
+    gr.Markdown(
+        """
+        ---
+        <div style="text-align: center; color: #666;">
+        <p>Built with ❤️ using Gradio | Model: FlameF0X/i3-200m-v2</p>
+        </div>
+        """
+    )
+    
+    # Connect UI
+    generate_btn.click(
+        predict,
+        inputs=[prompt_input, max_tokens_input, temp_input, topk_input],
+        outputs=[output_text]
+    )
 
 if __name__ == "__main__":
     demo.queue()