Update app.py

app.py

@@ -6,21 +6,19 @@ from transformers import AutoTokenizer, AutoModelForCausalLM
 import spaces  # Mandatory library for Hugging Face ZeroGPU [1]
 
 MODEL_ID = "DevStudio-AI/Devstudio-Coder-1.5B"
-# Fetch your secure token from the Space Secrets environment
 HF_TOKEN = os.environ.get("HF_TOKEN")
 
 print("Loading tokenizer and base model...")
 
-# 1. Load the tokenizer from the official, guaranteed-clean Qwen repository
-# This bypasses any local repository file corruption or cache issues
+# We load the tokenizer from the official Qwen repository for safety
 tokenizer = AutoTokenizer.from_pretrained(
     "Qwen/Qwen2.5-Coder-1.5B-Instruct"
 )
 
-# 2. Load the model, pointing directly to the nested directory inside your repo
+# Load your custom fine-tuned model from the subfolder
 model = AutoModelForCausalLM.from_pretrained(
     MODEL_ID,
-    subfolder="models/final_merged",  # Tells the loader to look inside this folder [1]
+    subfolder="models/final_merged",
     torch_dtype=torch.float16,
     device_map="cpu",
     token=HF_TOKEN
@@ -34,8 +32,25 @@ def generate_code(prompt, temperature, max_tokens):
         # Move model to CUDA dynamically inside the GPU context [1]
         model.to("cuda")
         
-        inputs = tokenizer(prompt, return_tensors="pt").to("cuda")
+        # --- FIXED: Apply the identical System Prompt and ChatML formatting ---
+        system_prompt = (
+            "You are DevStudio-1.5B, an in-editor coding assistant developed by DevStudio AI. "
+            "You are a highly specialized master of modern single-file HTML and Tailwind CSS designs. "
+        )
+        
+        # Build conversational structure matching your training data
+        messages = [
+            {"role": "system", "content": system_prompt},
+            {"role": "user", "content": prompt}
+        ]
+        
+        # Convert structure into Qwen's ChatML template string
+        formatted_prompt = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
+        
+        # Tokenize the formatted ChatML prompt
+        inputs = tokenizer(formatted_prompt, return_tensors="pt").to("cuda")
         
+        # Generate completion
         outputs = model.generate(
             **inputs,
             max_new_tokens=int(max_tokens),