Update app.py

app.py

@@ -1,29 +1,22 @@
-import os
-os.environ["HF_HOME"] = "/tmp/hf"  # Prevents write errors on Hugging Face Spaces
-
 from fastapi import FastAPI
 from pydantic import BaseModel
 from transformers import AutoTokenizer, AutoModelForCausalLM
 import torch
 import re
+import os
 
-app = FastAPI(
-    title="JavaScript Code Generator API",
-    description="Generate simple JavaScript functions from natural language instructions",
-    version="1.0"
-)
+# Fix the permission issue by changing HF cache location
+os.environ["HF_HOME"] = "/tmp/hf"
+
+app = FastAPI()
 
-# Load model and tokenizer
 model_id = "misalsathsara/phi1.5-js-codegen"
 tokenizer = AutoTokenizer.from_pretrained(model_id)
 model = AutoModelForCausalLM.from_pretrained(model_id)
-
-# Device setup
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 model.to(device)
 model.eval()
 
-# Prompt Template
 system_prompt = """
 You are a smart javascript assistant that only generates only the best simple javascript functions without any comments like this:
 function transform(row) {
@@ -47,21 +40,19 @@ Don't add any markdown block markers either.
 Every function must end with return row;
 """
 
-# Input schema
 class RequestData(BaseModel):
     instruction: str
 
-# Main route
-@app.post("/generate", summary="Generate JavaScript code", tags=["Code Generation"])
+@app.post("/generate")
 def generate_code(data: RequestData):
     instruction = data.instruction
-    full_prompt = f"{system_prompt}\n### Instruction:\n{instruction}\n\n### Response:\n"
-
-    inputs = tokenizer(full_prompt, return_tensors="pt").to(device)
-
+    full_prompt = system_prompt + f"\n### Instruction:\n{instruction}\n\n### Response:\n"
+    
+    input_ids = tokenizer(full_prompt, return_tensors="pt").input_ids.to(device)
+    
     with torch.no_grad():
-        outputs = model.generate(
-            **inputs,
+        output_ids = model.generate(
+            input_ids,
             max_new_tokens=200,
             temperature=0.3,
             top_k=50,
@@ -70,6 +61,10 @@ def generate_code(data: RequestData):
             pad_token_id=tokenizer.eos_token_id
         )
 
-    result = tokenizer.decode(outputs[0][inputs["input_ids"].shape[-1]:], skip_special_tokens=True)
-    match = re.search(r"function\s*\(.*?\)\s*{.*?return row;\s*}", result, re.DOTALL)
-    return {"result": match.group(0).strip() if match else result.strip()}
+    generated_text = tokenizer.decode(output_ids[0][input_ids.shape[-1]:], skip_special_tokens=True)
+
+    # Only return JavaScript function — no extra text
+    match = re.search(r"function\s+transform\(.*?\)\s*{.*?return row;\s*}", generated_text, re.DOTALL)
+    clean_output = match.group(0).strip() if match else generated_text.strip()
+
+    return clean_output