Update app.py

app.py

@@ -1,5 +1,5 @@
 import os
-os.environ["HF_HOME"] = "/tmp/hf"
+os.environ["HF_HOME"] = "/tmp/hf"  # Prevents write errors on Hugging Face Spaces
 
 from fastapi import FastAPI
 from pydantic import BaseModel
@@ -7,17 +7,23 @@ from transformers import AutoTokenizer, AutoModelForCausalLM
 import torch
 import re
 
-app = FastAPI()
+app = FastAPI(
+    title="JavaScript Code Generator API",
+    description="Generate simple JavaScript functions from natural language instructions",
+    version="1.0"
+)
 
-# Load model and tokenizer from Hugging Face
+# 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()
 
-# Your system prompt
+# 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) {
@@ -41,22 +47,21 @@ Don't add any markdown block markers either.
 Every function must end with return row;
 """
 
-# Define the expected request body
+# Input schema
 class RequestData(BaseModel):
     instruction: str
 
-# POST endpoint
-@app.post("/generate")
+# Main route
+@app.post("/generate", summary="Generate JavaScript code", tags=["Code Generation"])
 def generate_code(data: RequestData):
     instruction = data.instruction
-    full_prompt = system_prompt + f"\n### Instruction:\n{instruction}\n\n### Response:\n"
-    
-    # Tokenize input
-    input_ids = tokenizer(full_prompt, return_tensors="pt").input_ids.to(device)
-    
+    full_prompt = f"{system_prompt}\n### Instruction:\n{instruction}\n\n### Response:\n"
+
+    inputs = tokenizer(full_prompt, return_tensors="pt").to(device)
+
     with torch.no_grad():
-        output_ids = model.generate(
-            input_ids,
+        outputs = model.generate(
+            **inputs,
             max_new_tokens=200,
             temperature=0.3,
             top_k=50,
@@ -65,10 +70,6 @@ def generate_code(data: RequestData):
             pad_token_id=tokenizer.eos_token_id
         )
 
-    generated_text = tokenizer.decode(output_ids[0][input_ids.shape[-1]:], skip_special_tokens=True)
-
-    # Extract clean JS function
-    match = re.search(r"function\s*\(.*?\)\s*{.*?return row;\s*}", generated_text, re.DOTALL)
-    clean_output = match.group(0).strip() if match else generated_text.strip()
-
-    return {"result": clean_output}
+    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()}