Update app.py

app.py

@@ -2,13 +2,14 @@ import warnings
 warnings.filterwarnings("ignore")
 
 import torch
-import gradio as gr
+from fastapi import FastAPI
+from pydantic import BaseModel
 from transformers import AutoTokenizer, AutoModelForCausalLM
 
-# Reduce CPU pressure
 torch.set_num_threads(1)
 
-# ✅ Use lightweight model (IMPORTANT)
+app = FastAPI()
+
 BASE_MODEL = "distilgpt2"
 
 print("Loading model...")
@@ -20,6 +21,12 @@ model.eval()
 
 print("Model ready")
 
+# ─────────────────────────
+# Request schema
+# ─────────────────────────
+class Query(BaseModel):
+    question: str
+
 # ─────────────────────────
 # SQL FILTER
 # ─────────────────────────
@@ -34,28 +41,26 @@ def is_sql_related(text):
     return any(k in text for k in SQL_KEYWORDS)
 
 # ─────────────────────────
-# PROMPT
+# Endpoint
 # ─────────────────────────
-SYSTEM_PROMPT = """
-You are an expert SQL generator.
-Rules:
-- Only respond to SQL or database related questions.
-- Output ONLY SQL query.
-- No explanation.
-"""
+@app.post("/generate-sql")
+def generate_sql(data: Query):
 
-# ─────────────────────────
-# GENERATION
-# ─────────────────────────
-def generate_sql(user_input):
+    user_input = data.question
 
     if not user_input.strip():
-        return "Enter SQL question."
+        return {"error": "Empty input"}
 
     if not is_sql_related(user_input):
-        return "Only SQL/database questions are allowed."
+        return {"error": "Only SQL-related queries allowed"}
+
+    prompt = f"""
+You are an expert SQL generator.
+Only output SQL query.
 
-    prompt = f"{SYSTEM_PROMPT}\nUser: {user_input}\nSQL:"
+User: {user_input}
+SQL:
+"""
 
     inputs = tokenizer(prompt, return_tensors="pt")
 
@@ -70,33 +75,6 @@ def generate_sql(user_input):
 
     text = tokenizer.decode(output[0], skip_special_tokens=True)
 
-    result = text.split("SQL:")[-1].strip()
-    result = result.split("\n")[0]
+    result = text.split("SQL:")[-1].strip().split("\n")[0]
 
-    return result
-
-# ─────────────────────────
-# UI
-# ─────────────────────────
-demo = gr.Interface(
-    fn=generate_sql,
-    inputs=gr.Textbox(
-        lines=3,
-        label="SQL Question",
-        placeholder="Find duplicate emails in users table"
-    ),
-    outputs=gr.Textbox(
-        lines=6,
-        label="Generated SQL"
-    ),
-    title="AI SQL Generator (Portfolio Project)",
-    description="Only SQL/database queries are supported.",
-    examples=[
-        ["Find duplicate emails in users table"],
-        ["Top 5 highest paid employees"],
-        ["Count orders per customer last month"],
-        ["Write a joke about cats"]
-    ],
-)
-
-demo.launch(server_name="0.0.0.0", server_port=7860)
+    return {"sql": result}