Update app.py

app.py

@@ -2,23 +2,21 @@ import warnings
 warnings.filterwarnings("ignore")
 
 import torch
-from fastapi import FastAPI
-from pydantic import BaseModel
-from transformers import AutoTokenizer, AutoModelForCausalLM
-
 import gradio as gr
-import threading
+from transformers import AutoTokenizer, AutoModelForCausalLM
 
 torch.set_num_threads(1)
 
-app = FastAPI(title="SQL Generator API")
-
 BASE_MODEL = "TinyLlama/TinyLlama-1.1B-Chat-v1.0"
 
 print("Loading model...")
 
+model = AutoModelForCausalLM.from_pretrained(
+    BASE_MODEL,
+    torch_dtype=torch.float32
+)
+
 tokenizer = AutoTokenizer.from_pretrained(BASE_MODEL)
-model = AutoModelForCausalLM.from_pretrained(BASE_MODEL)
 
 model.eval()
 
@@ -34,21 +32,34 @@ SQL_KEYWORDS = [
 ]
 
 def is_sql_related(text):
-    return any(k in text.lower() for k in SQL_KEYWORDS)
+    text = text.lower()
+    return any(k in text for k in SQL_KEYWORDS)
 
+# ─────────────────────────
+# GENERATION
+# ─────────────────────────
 SYSTEM_PROMPT = """
 You are an expert SQL generator.
-Only output SQL query.
+Rules:
+- Only respond to SQL or database related questions.
+- If the question is not about SQL or databases, refuse.
+- Output ONLY SQL query.
+- Do not explain.
 """
 
-def generate_sql(user_input: str):
+def generate_sql(user_input):
+
     if not user_input.strip():
         return "Enter SQL question."
 
     if not is_sql_related(user_input):
-        return "Only SQL/database questions allowed."
+        return "I only respond to SQL and database related questions. If you want, I can craft helpful database queries for you."
 
-    prompt = f"{SYSTEM_PROMPT}\nUser: {user_input}\nSQL:"
+    prompt = f"""
+{SYSTEM_PROMPT}
+User request: {user_input}
+SQL:
+"""
 
     inputs = tokenizer(prompt, return_tensors="pt")
 
@@ -62,35 +73,34 @@ def generate_sql(user_input: str):
         )
 
     text = tokenizer.decode(output[0], skip_special_tokens=True)
-    result = text.split("SQL:")[-1].strip().split("\n")[0]
 
-    return result
+    result = text.split("SQL:")[-1].strip()
+    result = result.split("\n\n")[0]
 
-# ─────────────────────────
-# API
-# ─────────────────────────
-class Query(BaseModel):
-    text: str
-
-@app.get("/")
-def root():
-    return {"status": "API running"}
-
-@app.post("/generate")
-def generate(query: Query):
-    return {"result": generate_sql(query.text)}
+    return result
 
 # ─────────────────────────
-# Gradio UI (for testing)
+# UI
 # ─────────────────────────
-def run_gradio():
-    demo = gr.Interface(
-        fn=generate_sql,
-        inputs=gr.Textbox(lines=3, label="SQL Question"),
-        outputs=gr.Textbox(lines=6, label="Generated SQL"),
-        title="SQL Generator"
-    )
-    demo.launch(server_name="0.0.0.0", server_port=7860)
-
-# Run UI in separate thread
-threading.Thread(target=run_gradio).start()
+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=8,
+        label="Generated SQL"
+    ),
+    title="AI SQL Generator (Portfolio Project)",
+    description="This model ONLY responds to SQL/database queries.",
+    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")