Update app.py

app.py

@@ -66,9 +66,16 @@ def generate_code(data: RequestData):
     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()
-
+    # Extract only the JavaScript function that ends with return row;
+    match = re.search(r"function\s+transform\s*\([^)]*\)\s*{[^}]*return row;\s*}", generated_text, re.DOTALL)
+    if match:
+        clean_output = match.group(0).strip()
+    else:
+        # fallback: try to grab only up to "return row;"
+        fallback = generated_text.split("return row;")[0] + "return row;"
+        clean_output = fallback.strip()
+    
     from fastapi.responses import PlainTextResponse
     return PlainTextResponse(clean_output)
 
+