Update app.py

app.py

@@ -1,95 +1,77 @@
 import gradio as gr
 from transformers import pipeline, AutoTokenizer
 
-# Load the StarChat beta model
-MODEL_NAME = "HuggingFaceH4/starchat-beta"
+# Load a smaller model that fits within free tier memory limits
+MODEL_NAME = "Salesforce/codegen-350M-mono"
 tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
 generator = pipeline(
     "text-generation",
     model=MODEL_NAME,
     tokenizer=tokenizer,
-    device="cpu",  # Free tier uses CPU
+    device="cpu",  # Force CPU usage for free tier
 )
 
 def generate_code(natural_language_input):
     try:
-        # Create a proper prompt for the model
-        chat_prompt = [
-            {"role": "system", "content": "You are a helpful AI assistant that generates Python code based on natural language descriptions."},
-            {"role": "user", "content": f"Write Python code that: {natural_language_input}"}
-        ]
+        # Create a clear prompt for code generation
+        prompt = f"# Python code that {natural_language_input}\n#"
         
-        # Generate the prompt for the model
-        prompt = tokenizer.apply_chat_template(chat_prompt, tokenize=False, add_generation_prompt=True)
-        
-        # Generate code with appropriate parameters for code generation
+        # Generate with conservative parameters to save memory
         generated = generator(
             prompt,
-            max_new_tokens=256,
-            temperature=0.2,
-            top_p=0.95,
-            do_sample=True,
+            max_new_tokens=128,  # Reduced from 256 to save memory
+            temperature=0.7,
+            top_p=0.9,
             pad_token_id=tokenizer.eos_token_id,
-            eos_token_id=tokenizer.eos_token_id,
+            do_sample=True,
         )
         
         # Extract and clean the generated code
         full_response = generated[0]['generated_text']
         code_response = full_response.replace(prompt, "").strip()
         
-        # Sometimes the model adds non-code text - we'll try to extract just the code blocks
-        if "```python" in code_response:
-            code_response = code_response.split("```python")[1].split("```")[0]
-        elif "```" in code_response:
-            code_response = code_response.split("```")[1].split("```")[0]
+        # Remove any non-code text after generation
+        if "```" in code_response:
+            code_response = code_response.split("```")[0]
         
         return code_response
     except Exception as e:
-        return f"Error generating code: {str(e)}"
+        return f"Error generating code: {str(e)}\n\n(Tip: Try a simpler description or break your request into smaller parts)"
 
-# Gradio interface
-with gr.Blocks(theme=gr.themes.Soft()) as demo:
+# Streamlined Gradio interface to reduce memory footprint
+with gr.Blocks(title="Text to Python Code Generator") as demo:
     gr.Markdown("""
-    # 🚀 Text to Python Code Generator
-    ### Powered by StarChat Beta
-    Describe what you want the code to do in natural language, and get Python code!
+    # 🐍 Python Code Generator
+    *Free tier version using Salesforce/codegen-350M-mono*
     """)
     
     with gr.Row():
-        with gr.Column(scale=3):
-            input_text = gr.Textbox(
-                label="Describe your code task",
-                placeholder="e.g., 'create a Flask web server with two routes'",
-                lines=3,
-                max_lines=6
-            )
-            with gr.Row():
-                generate_btn = gr.Button("Generate Code", variant="primary")
-                clear_btn = gr.Button("Clear")
-            
-            gr.Markdown("**Tips:** Be specific in your description for better results.")
-            
-        with gr.Column(scale=4):
-            output_code = gr.Code(
-                label="Generated Python Code",
-                language="python",
-                interactive=True,
-                lines=10
-            )
-            with gr.Row():
-                copy_btn = gr.Button("Copy to Clipboard")
-                refresh_btn = gr.Button("Try Again")
+        input_text = gr.Textbox(
+            label="Describe what you want the code to do",
+            placeholder="e.g., 'calculate fibonacci sequence up to N numbers'",
+            lines=3
+        )
+    
+    with gr.Row():
+        generate_btn = gr.Button("Generate", variant="primary")
+        clear_btn = gr.Button("Clear")
+    
+    output_code = gr.Code(
+        label="Generated Python Code",
+        language="python",
+        interactive=True
+    )
     
-    # Examples
+    # Memory-friendly examples
     examples = gr.Examples(
         examples=[
-            ["create a function to calculate factorial recursively"],
-            ["write code to download a file from URL and save it locally"],
-            ["create a pandas DataFrame from a dictionary and plot a bar chart"],
-            ["implement a simple REST API using FastAPI with GET and POST endpoints"]
+            ["print hello world"],
+            ["function to calculate square of a number"],
+            ["read a CSV file using pandas"],
+            ["simple Flask app with one route"]
         ],
         inputs=input_text,
-        label="Click any example to try it"
+        label="Try these examples (keep requests simple)"
     )
     
     # Button actions
@@ -104,18 +86,6 @@ with gr.Blocks(theme=gr.themes.Soft()) as demo:
         inputs=None,
         outputs=[input_text, output_code]
     )
-    
-    refresh_btn.click(
-        fn=generate_code,
-        inputs=input_text,
-        outputs=output_code
-    )
-    
-    copy_btn.click(
-        fn=lambda code: gr.Clipboard().copy(code),
-        inputs=output_code,
-        outputs=None,
-        api_name="copy_code"
-    )
 
-demo.launch(debug=True)
+# Launch with minimal resources
+demo.launch(debug=False, show_error=True)