Update app.py

app.py

@@ -1,24 +1,33 @@
-# UI Generator Hugging Face Space - MVP
-# Requirements: gradio, transformers, torch
+# Updated UI Generator with Better Open-Source Model
 
 import gradio as gr
-from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
+from transformers import AutoTokenizer, AutoModelForCausalLM, TextGenerationPipeline
+import torch
 
-# Load an open-source code generation model (you can swap with DeepSeek-Coder, CodeLlama, etc.)
-model_id = "Salesforce/codegen-350M-mono"
+# Use a better model for code generation, such as DeepSeek-Coder or Codestral
+model_id = "deepseek-ai/deepseek-coder-6.7b-instruct"  # Change to another if desired
+
+# Ensure torch uses CPU if no GPU
+device = 0 if torch.cuda.is_available() else -1
+
+# Load tokenizer and model
+print("Loading model...")
 tokenizer = AutoTokenizer.from_pretrained(model_id)
-model = AutoModelForCausalLM.from_pretrained(model_id)
-generator = pipeline("text-generation", model=model, tokenizer=tokenizer)
+model = AutoModelForCausalLM.from_pretrained(model_id, torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32)
+
+# Build the text generation pipeline
+generator = TextGenerationPipeline(model=model, tokenizer=tokenizer, device=device)
 
 def generate_ui(platform, framework, ui_prompt):
     prompt = f"""
-    You are an expert {framework} developer.
-    Generate full code for a {platform} UI screen based on the following description:
-    "{ui_prompt}"
-    Please include necessary imports and use best practices.
-    """
-    output = generator(prompt, max_length=512, do_sample=True, temperature=0.7)[0]["generated_text"]
-    return output.strip()
+You are an expert mobile app developer.
+Generate a complete {framework} UI code snippet for a {platform} app based on the description:
+"{ui_prompt}"
+Include all required imports, a main method, and best practices for UI structure.
+"""
+    response = generator(prompt, max_new_tokens=512, do_sample=True, temperature=0.7)[0]['generated_text']
+    # Trim the echoed prompt and just return the generated code
+    return response.split(""""""")[-1].strip()
 
 interface = gr.Interface(
     fn=generate_ui,