Update app.py

app.py

@@ -1,137 +1,131 @@
 import gradio as gr
-import torch
 from transformers import pipeline
 
-def initialize_model():
-    """Initialize the text generation pipeline with device detection"""
-    # Check if CUDA is available, otherwise fall back to CPU
-    device = "cuda" if torch.cuda.is_available() else "cpu"
-    print(f"Using device: {device}")
-    
-    try:
-        generator = pipeline(
-            "text-generation", 
-            model="akhaliq/MyGemmaGradioCoder", 
-            device=device
-        )
-        return generator
-    except Exception as e:
-        print(f"Error loading model: {e}")
-        # Fallback to CPU if CUDA fails
-        if device == "cuda":
-            print("Falling back to CPU...")
-            generator = pipeline(
-                "text-generation", 
-                model="akhaliq/gemma-3-270m-gradio-coder", 
-                device="cpu"
-            )
-            return generator
-        raise e
-
-# Initialize the model globally
-print("Loading model...")
-generator = initialize_model()
-print("Model loaded successfully!")
+# Initialize the text generation pipeline
+pipe = pipeline("text-generation", model="akhaliq/MyGemmaGradioCoder")
 
-def chat_response(message, history):
-    """Generate response for the chatbot"""
-    try:
-        # Format the message for the model
-        input_message = [{"role": "user", "content": message}]
-        
-        # Generate response
-        output = generator(
-            input_message, 
-            max_new_tokens=20000, 
-            return_full_text=False,
-            do_sample=True,
-            temperature=0.7,
-            pad_token_id=generator.tokenizer.eos_token_id
-        )[0]
-        
-        response = output["generated_text"]
-        return response
-        
-    except Exception as e:
-        return f"Sorry, I encountered an error: {str(e)}"
-
-# Create the Gradio interface
-def create_chatbot():
-    """Create and launch the Gradio chatbot interface"""
-    
-    # Custom CSS for better styling
-    css = """
-    .gradio-container {
-        max-width: 800px !important;
-        margin: auto !important;
-    }
-    .chat-message {
-        padding: 10px !important;
-        margin: 5px !important;
-        border-radius: 10px !important;
-    }
+def generate_code(user_input):
     """
+    Generate code based on user input using the Gemma model
+    """
+    messages = [
+        {"role": "user", "content": user_input},
+    ]
     
-    # Create the chatbot interface
-    with gr.Blocks(css=css, title="AI Chatbot") as demo:
-        gr.Markdown("# 🤖 AI Chatbot")
-        gr.Markdown("*Powered by Gemma-3-270m model via Transformers*")
-        
-        chatbot = gr.Chatbot(
-            height=500,
-            bubble_full_width=False,
-            show_label=False
-        )
+    # Generate response from the model
+    response = pipe(messages, max_new_tokens=512, temperature=0.7, do_sample=True)
+    
+    # Extract the generated text from the response
+    generated_text = response[0]['generated_text']
+    
+    # If the response contains the full conversation, extract just the assistant's response
+    if isinstance(generated_text, list):
+        # Handle conversation format
+        for msg in generated_text:
+            if msg.get('role') == 'assistant':
+                return msg.get('content', '')
+        # If no assistant message found, return the last message content
+        return generated_text[-1].get('content', '') if generated_text else ""
+    else:
+        # Handle string format - try to extract the code after the user input
+        if user_input in generated_text:
+            return generated_text.split(user_input)[-1].strip()
+        return generated_text
+
+# Create the Gradio interface
+with gr.Blocks(title="Text to Code Generator", theme=gr.themes.Soft()) as demo:
+    gr.Markdown(
+        """
+        # 🚀 Text to Code Generator
         
-        with gr.Row():
-            msg = gr.Textbox(
-                placeholder="Type your message here...",
-                show_label=False,
-                scale=4
+        Generate code from natural language descriptions using the Gemma Gradio Coder model.
+        Simply describe what you want to build, and the AI will generate the corresponding code!
+        """
+    )
+    
+    with gr.Row():
+        with gr.Column(scale=1):
+            # Input section
+            input_text = gr.Textbox(
+                label="Describe what you want to code",
+                placeholder="e.g., Create a Python function that calculates the factorial of a number",
+                lines=5,
+                max_lines=10
             )
-            send_btn = gr.Button("Send", scale=1, variant="primary")
-            clear_btn = gr.Button("Clear", scale=1)
-        
-        # Example questions
-        gr.Examples(
-            examples=[
-                "If you had a time machine, but could only go to the past or the future once and never return, which would you choose and why?",
-                "What's the most important lesson you've learned in life?",
-                "How do you think AI will change the world in the next 10 years?",
-                "What would you do if you had unlimited resources for one day?"
-            ],
-            inputs=msg
-        )
-        
-        def respond(message, chat_history):
-            if not message.strip():
-                return chat_history, ""
             
-            # Get bot response
-            bot_message = chat_response(message, chat_history)
+            with gr.Row():
+                generate_btn = gr.Button("Generate Code", variant="primary", scale=2)
+                clear_btn = gr.ClearButton([input_text], value="Clear", scale=1)
             
-            # Add to chat history
-            chat_history.append((message, bot_message))
-            return chat_history, ""
-        
-        def clear_chat():
-            return [], ""
+            # Examples section
+            gr.Examples(
+                examples=[
+                    ["Create a Python function to check if a number is prime"],
+                    ["Write a JavaScript function to reverse a string"],
+                    ["Create a React component for a todo list item"],
+                    ["Write a SQL query to find the top 5 customers by total purchase amount"],
+                    ["Create a Python class for a bank account with deposit and withdraw methods"],
+                ],
+                inputs=input_text,
+                label="Example Prompts"
+            )
         
-        # Event handlers
-        msg.submit(respond, [msg, chatbot], [chatbot, msg])
-        send_btn.click(respond, [msg, chatbot], [chatbot, msg])
-        clear_btn.click(clear_chat, None, [chatbot, msg])
+        with gr.Column(scale=1):
+            # Output section
+            output_code = gr.Code(
+                label="Generated Code",
+                language="python",
+                lines=20,
+                interactive=True,
+                show_line_numbers=True,
+                wrap_lines=True,
+                autocomplete=True
+            )
+            
+            with gr.Row():
+                copy_btn = gr.Button("📋 Copy Code", scale=1)
+            
+    # Add event handlers
+    generate_btn.click(
+        fn=generate_code,
+        inputs=input_text,
+        outputs=output_code,
+        api_name="generate"
+    )
     
-    return demo
-
-if __name__ == "__main__":
-    print("Creating Gradio interface...")
-    demo = create_chatbot()
+    input_text.submit(
+        fn=generate_code,
+        inputs=input_text,
+        outputs=output_code
+    )
+    
+    # Copy functionality (using JavaScript)
+    copy_btn.click(
+        fn=None,
+        inputs=output_code,
+        outputs=None,
+        js="""
+        (code) => {
+            navigator.clipboard.writeText(code);
+            alert('Code copied to clipboard!');
+            return null;
+        }
+        """
+    )
     
-    print("Starting Gradio server...")
-    demo.launch(
-        share=False,  # Set to True if you want a public link
-        server_name="0.0.0.0",  # Allow external connections
-        server_port=7860,
-        show_error=True
+    # Footer
+    gr.Markdown(
+        """
+        ---
+        💡 **Tips:**
+        - Be specific about the programming language you want
+        - Include details about inputs, outputs, and edge cases
+        - You can edit the generated code directly in the output box
+        
+        Powered by [akhaliq/MyGemmaGradioCoder](https://huggingface.co/akhaliq/MyGemmaGradioCoder)
+        """
     )
+
+# Launch the app
+if __name__ == "__main__":
+    demo.launch()