Update app.py

app.py

@@ -6,9 +6,10 @@ from fastapi import FastAPI
 from pydantic import BaseModel
 from transformers import AutoTokenizer, AutoModelForCausalLM
 
+# Optimize CPU
 torch.set_num_threads(1)
 
-app = FastAPI()
+app = FastAPI(title="SQL Generator API")
 
 BASE_MODEL = "distilgpt2"
 
@@ -22,10 +23,10 @@ model.eval()
 print("Model ready")
 
 # ─────────────────────────
-# Request schema
+# Request Schema
 # ─────────────────────────
 class Query(BaseModel):
-    question: str
+    text: str
 
 # ─────────────────────────
 # SQL FILTER
@@ -37,30 +38,24 @@ SQL_KEYWORDS = [
 ]
 
 def is_sql_related(text):
-    text = text.lower()
-    return any(k in text for k in SQL_KEYWORDS)
+    return any(k in text.lower() for k in SQL_KEYWORDS)
 
 # ─────────────────────────
-# Endpoint
+# Generator
 # ─────────────────────────
-@app.post("/generate-sql")
-def generate_sql(data: Query):
-
-    user_input = data.question
+SYSTEM_PROMPT = """
+You are an expert SQL generator.
+Only output SQL query.
+"""
 
+def generate_sql(user_input: str):
     if not user_input.strip():
-        return {"error": "Empty input"}
+        return "Empty input."
 
     if not is_sql_related(user_input):
-        return {"error": "Only SQL-related queries allowed"}
-
-    prompt = f"""
-You are an expert SQL generator.
-Only output SQL query.
+        return "Only SQL-related queries allowed."
 
-User: {user_input}
-SQL:
-"""
+    prompt = f"{SYSTEM_PROMPT}\nUser: {user_input}\nSQL:"
 
     inputs = tokenizer(prompt, return_tensors="pt")
 
@@ -74,7 +69,18 @@ SQL:
         )
 
     text = tokenizer.decode(output[0], skip_special_tokens=True)
-
     result = text.split("SQL:")[-1].strip().split("\n")[0]
 
-    return {"sql": result}
+    return result
+
+# ─────────────────────────
+# Routes
+# ─────────────────────────
+@app.get("/")
+def root():
+    return {"status": "API is running"}
+
+@app.post("/generate")
+def generate(query: Query):
+    result = generate_sql(query.text)
+    return {"result": result}