Update app.py

app.py

@@ -7,9 +7,6 @@ from transformers import AutoTokenizer, AutoModelForCausalLM
 
 torch.set_num_threads(1)
 
-# ─────────────────────────
-# MODEL
-# ─────────────────────────
 BASE_MODEL = "TinyLlama/TinyLlama-1.1B-Chat-v1.0"
 
 print("Loading model...")
@@ -20,22 +17,49 @@ model = AutoModelForCausalLM.from_pretrained(
 )
 
 tokenizer = AutoTokenizer.from_pretrained(BASE_MODEL)
-
 model.eval()
+
 print("Model ready")
 
+# ─────────────────────────
+# SQL FILTER
+# ─────────────────────────
+SQL_KEYWORDS = [
+    "sql", "database", "table", "select", "insert",
+    "update", "delete", "join", "group by",
+    "postgres", "mysql", "sqlite", "query"
+]
+
+def is_sql_related(text):
+    text = text.lower()
+    return any(k in text for k in SQL_KEYWORDS)
+
 # ─────────────────────────
 # GENERATION
 # ─────────────────────────
-def generate_sql(question):
-    if not question.strip():
+SYSTEM_PROMPT = """
+You are an expert SQL generator.
+
+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):
+
+    if not user_input.strip():
         return "Enter SQL question."
 
+    # 🔴 HARD GUARD
+    if not is_sql_related(user_input):
+        return "❌ This demo only supports SQL and database related questions."
+
     prompt = f"""
-You are a SQL expert.
-Return ONLY SQL query.
+{SYSTEM_PROMPT}
 
-User: {question}
+User request: {user_input}
 SQL:
 """
 
@@ -45,28 +69,42 @@ SQL:
         output = model.generate(
             **inputs,
             max_new_tokens=120,
-            temperature=0.2,
+            temperature=0.1,
             do_sample=False,
             pad_token_id=tokenizer.eos_token_id,
         )
 
     text = tokenizer.decode(output[0], skip_special_tokens=True)
 
-    return text.split("SQL:")[-1].strip()
+    # return only SQL part
+    result = text.split("SQL:")[-1].strip()
+
+    # extra safety: remove explanations
+    result = result.split("\n\n")[0]
+
+    return result
 
 # ─────────────────────────
 # UI
 # ─────────────────────────
 demo = gr.Interface(
     fn=generate_sql,
-    inputs=gr.Textbox(lines=3, label="SQL Question"),
-    outputs=gr.Textbox(lines=8, label="Generated SQL"),
-    title="SQL Generator (Portfolio Demo)",
-    description="Fast CPU AI SQL generator.",
+    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"],
-        ["Orders per customer last month"],
+        ["Count orders per customer last month"],
+        ["Write a joke about cats"]  # will be blocked
     ],
 )