Update app.py

app.py

@@ -5,43 +5,38 @@ from transformers import AutoTokenizer, AutoModelForCausalLM
 
 MODEL_ID = "saadkhi/SQL_Chat_finetuned_model"
 
-app = FastAPI(title="SQL Chatbot API")
+app = FastAPI()
 
-# Load model once (on startup)
+# ---- LOAD ONCE ONLY ----
 tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
+
 model = AutoModelForCausalLM.from_pretrained(
     MODEL_ID,
-    torch_dtype=torch.float16,
-    device_map="auto"
+    dtype=torch.float16,     # use dtype, not torch_dtype
+    device_map="auto",
+    low_cpu_mem_usage=True
 )
 
+model.eval()
+
+
 class QueryRequest(BaseModel):
     prompt: str
     max_new_tokens: int = 256
 
-class QueryResponse(BaseModel):
-    response: str
 
-
-@app.post("/generate", response_model=QueryResponse)
-def generate_answer(request: QueryRequest):
-    inputs = tokenizer(
-        request.prompt,
-        return_tensors="pt"
-    ).to(model.device)
+@app.post("/generate")
+def generate(req: QueryRequest):
+    inputs = tokenizer(req.prompt, return_tensors="pt").to(model.device)
 
     with torch.no_grad():
-        output_ids = model.generate(
+        outputs = model.generate(
             **inputs,
-            max_new_tokens=request.max_new_tokens,
+            max_new_tokens=req.max_new_tokens,
             do_sample=True,
             temperature=0.7,
             top_p=0.9
         )
 
-    output_text = tokenizer.decode(
-        output_ids[0],
-        skip_special_tokens=True
-    )
-
-    return {"response": output_text}
+    text = tokenizer.decode(outputs[0], skip_special_tokens=True)
+    return {"response": text}