Upload app.py with huggingface_hub

app.py

@@ -11,6 +11,22 @@ import os
 # Configuration
 MODEL_ID = "amkyawdev/mm-coder-agent-v1-combined"
 
+# System prompt for better responses
+SYSTEM_PROMPT = """You are MM Coder Agent v1, a professional AI coding assistant.
+
+Guidelines:
+1. Write clean, efficient code with proper imports
+2. Be concise and practical
+3. Use idiomatic code for the target language
+4. Include error handling when needed
+5. Prioritize security and best practices
+
+Output format:
+- Start with code directly
+- Use code blocks with language tags
+- Add brief explanation only if complex
+- Always verify your code works"""
+
 # Global variables to store model and tokenizer
 model = None
 tokenizer = None
@@ -42,7 +58,10 @@ def generate_code(prompt, max_tokens=512, temperature=0.7):
     if model is None:
         load_model()
     
-    inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
+    # Format prompt with system instruction
+    formatted_prompt = f"{SYSTEM_PROMPT}\n\nUser: {prompt}\nAssistant:"
+    
+    inputs = tokenizer(formatted_prompt, return_tensors="pt").to(model.device)
     
     with torch.no_grad():
         outputs = model.generate(
@@ -56,8 +75,11 @@ def generate_code(prompt, max_tokens=512, temperature=0.7):
     
     response = tokenizer.decode(outputs[0], skip_special_tokens=True)
     
-    # Remove the prompt from response if it's included
-    if response.startswith(prompt):
+    # Remove the prompt and system prompt from response
+    prefix_to_remove = f"{SYSTEM_PROMPT}\n\nUser: {prompt}\nAssistant:"
+    if response.startswith(prefix_to_remove):
+        response = response[len(prefix_to_remove):].strip()
+    elif response.startswith(prompt):
         response = response[len(prompt):].strip()
     
     return response