Update app.py

app.py

@@ -5,19 +5,17 @@ import torch
 import gradio as gr
 from transformers import AutoTokenizer, AutoModelForCausalLM
 
+# Reduce CPU pressure
 torch.set_num_threads(1)
 
-# BASE_MODEL = "TinyLlama/TinyLlama-1.1B-Chat-v1.0"
-BASE_MODEL = "tiiuae/falcon-rw-1b"
+# ✅ Use lightweight model (IMPORTANT)
+BASE_MODEL = "distilgpt2"
 
 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()
 
 print("Model ready")
@@ -36,52 +34,44 @@ def is_sql_related(text):
     return any(k in text for k in SQL_KEYWORDS)
 
 # ─────────────────────────
-# GENERATION
+# PROMPT
 # ─────────────────────────
 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.
+- No explanation.
 """
 
+# ─────────────────────────
+# GENERATION
+# ─────────────────────────
 def generate_sql(user_input):
 
     if not user_input.strip():
         return "Enter SQL question."
 
-    # HARD GUARD
     if not is_sql_related(user_input):
-        return "I only respond to SQL and database related questions. If you want, I can craft helpful database queries for you."
+        return "Only SQL/database questions are allowed."
 
-    prompt = f"""
-{SYSTEM_PROMPT}
-
-User request: {user_input}
-SQL:
-"""
+    prompt = f"{SYSTEM_PROMPT}\nUser: {user_input}\nSQL:"
 
     inputs = tokenizer(prompt, return_tensors="pt")
 
     with torch.no_grad():
         output = model.generate(
             **inputs,
-            max_new_tokens=120,
-            temperature=0.1,
-            do_sample=False,
+            max_new_tokens=80,
+            temperature=0.2,
+            do_sample=True,
             pad_token_id=tokenizer.eos_token_id,
         )
 
     text = tokenizer.decode(output[0], skip_special_tokens=True)
 
-    # return only SQL part
     result = text.split("SQL:")[-1].strip()
-
-    # extra safety: remove explanations
-    result = result.split("\n\n")[0]
+    result = result.split("\n")[0]
 
     return result
 
@@ -96,17 +86,17 @@ demo = gr.Interface(
         placeholder="Find duplicate emails in users table"
     ),
     outputs=gr.Textbox(
-        lines=8,
+        lines=6,
         label="Generated SQL"
     ),
     title="AI SQL Generator (Portfolio Project)",
-    description="This model ONLY responds to SQL/database queries.",
+    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"]  # will be blocked
+        ["Write a joke about cats"]
     ],
 )
 
-demo.launch(server_name="0.0.0.0")
+demo.launch(server_name="0.0.0.0", server_port=7860)