Update app.py

app.py

@@ -3,19 +3,19 @@ import os
 import gradio as gr
 import torch
 from transformers import AutoTokenizer, AutoModelForCausalLM
-import spaces  # Mandatory library for Hugging Face ZeroGPU [1]
+import spaces  # Mandatory library for Hugging Face ZeroGPU
 
 MODEL_ID = "DevStudio-AI/Devstudio-Coder-1.5B"
 HF_TOKEN = os.environ.get("HF_TOKEN")
 
 print("Loading tokenizer and base model...")
 
-# We load the tokenizer from the official Qwen repository for safety
+# Load tokenizer
 tokenizer = AutoTokenizer.from_pretrained(
     "Qwen/Qwen2.5-Coder-1.5B-Instruct"
 )
 
-# Load your custom fine-tuned model from the subfolder
+# Load model
 model = AutoModelForCausalLM.from_pretrained(
     MODEL_ID,
     subfolder="models/final_merged",
@@ -23,57 +23,103 @@ model = AutoModelForCausalLM.from_pretrained(
     device_map="cpu",
     token=HF_TOKEN
 )
+
 print("Model successfully loaded on CPU. Awaiting ZeroGPU allocation...")
 
-# The @spaces.GPU decorator dynamically requests Nvidia A100 resources for this call [1]
+
 @spaces.GPU
 def generate_code(prompt, temperature, max_tokens):
     try:
-        # Move model to CUDA dynamically inside the GPU context [1]
+        # Move model to GPU inside the ZeroGPU context
         model.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. "
+            "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
+
+        # Apply ChatML formatting
+        formatted_prompt = tokenizer.apply_chat_template(
+            messages,
+            tokenize=False,
+            add_generation_prompt=True
+        )
+
+        inputs = tokenizer(
+            formatted_prompt,
+            return_tensors="pt"
+        ).to("cuda")
+
         outputs = model.generate(
             **inputs,
             max_new_tokens=int(max_tokens),
             temperature=float(temperature),
             do_sample=True,
-            eos_token_id=tokenizer.eos_token_id
+            eos_token_id=tokenizer.eos_token_id,
+            pad_token_id=tokenizer.eos_token_id
         )
-        
-        # Isolate and decode newly generated tokens
+
+        # Decode only the generated part
         generated_ids = outputs[0][inputs["input_ids"].shape[1]:]
-        return tokenizer.decode(generated_ids, skip_special_tokens=True)
+        result = tokenizer.decode(
+            generated_ids,
+            skip_special_tokens=True
+        )
+
+        # ---------------------------
+        # Post-process escaped output
+        # ---------------------------
+        result = result.replace("\\n", "\n")
+        result = result.replace("\\t", "\t")
+        result = result.replace('\\"', '"')
+        result = result.replace("\\'", "'")
+
+        # Remove escaped markdown fences if present
+        result = result.replace("\\`", "`")
+
+        return result
+
     except Exception as e:
         return f"Error during generation: {str(e)}"
 
-# Define the Gradio web interface
+    finally:
+        # Free GPU memory after each request
+        if torch.cuda.is_available():
+            model.to("cpu")
+            torch.cuda.empty_cache()
+
+
 demo = gr.Interface(
     fn=generate_code,
     inputs=[
-        gr.Textbox(label="Prompt", placeholder="Enter your prompt here..."),
-        gr.Slider(minimum=0.1, maximum=1.0, value=0.3, label="Temperature"),
-        gr.Slider(minimum=64, maximum=2048, value=1024, step=64, label="Max Tokens")
+        gr.Textbox(
+            label="Prompt",
+            placeholder="Enter your prompt here...",
+            lines=8
+        ),
+        gr.Slider(
+            minimum=0.1,
+            maximum=1.0,
+            value=0.3,
+            label="Temperature"
+        ),
+        gr.Slider(
+            minimum=64,
+            maximum=2048,
+            value=1024,
+            step=64,
+            label="Max Tokens"
+        )
     ],
-    outputs=gr.Textbox(label="Generated Code"),
+    outputs=gr.Textbox(
+        label="Generated Code",
+        lines=30
+    ),
     title="DevStudio-1.5B API Engine",
     description="Static HTML + Tailwind CSS specialized completion endpoint running on ZeroGPU."
 )