Upload app.py

app.py

@@ -1,38 +1,39 @@
-import os
-import json
-import time
-import sqlite3
-import textwrap
-import requests
+import os, time, json, sqlite3, textwrap, requests
 import gradio as gr
 
-# -----------------------------
-# Config
-# -----------------------------
+# -------------------------------------------------
+# 1. CONFIGURATION
+# -------------------------------------------------
+MODEL_ID = "defog/sqlcoder-7b-nl2sql-beta"        # working public model
+API_URL  = f"https://api-inference.huggingface.co/models/{MODEL_ID}"
+
+HF_TOKEN = os.getenv("HF_TOKEN")                  # set in Space → Settings → Secrets
+if not HF_TOKEN:
+    raise RuntimeError("HF_TOKEN secret not found. "
+                       "Add it in Space Settings → Secrets.")
+
+HEADERS = {"Authorization": f"Bearer {HF_TOKEN}"}
+
 DB_PATH     = "company.db"
 SCHEMA_FILE = "schema.sql"
-MODEL_ID = "tscholak/t5-base-spider"
-API_URL  = f"https://api-inference.huggingface.co/models/{MODEL_ID}"
-HF_TOKEN    = os.getenv("HF_TOKEN")           # set in Space > Settings > Secrets
-HEADERS     = {"Authorization": f"Bearer {HF_TOKEN}"} if HF_TOKEN else {}
 
-# -----------------------------
-# Helper: Build DB once
-# -----------------------------
+# -------------------------------------------------
+# 2. UTIL: BUILD SQLITE DB IF NEEDED
+# -------------------------------------------------
 def create_db_if_needed():
+    """Create SQLite DB from schema.sql the first time the app runs."""
     if os.path.exists(DB_PATH):
         return
     if not os.path.isfile(SCHEMA_FILE):
-        raise FileNotFoundError("schema.sql not found in Space.")
-    with open(SCHEMA_FILE) as f:
-        schema_sql = f.read()
-    with sqlite3.connect(DB_PATH) as conn:
-        conn.executescript(schema_sql)
-
-# -----------------------------
-# Helper: Call HF Inference API
-# -----------------------------
-def nlp_to_sql(nl_query: str, schema_ddl: str) -> str:
+        raise FileNotFoundError("schema.sql file is missing in the Space.")
+    with open(SCHEMA_FILE) as f, sqlite3.connect(DB_PATH) as conn:
+        conn.executescript(f.read())
+
+# -------------------------------------------------
+# 3. UTIL: CALL HUGGING FACE MODEL
+# -------------------------------------------------
+def nlp_to_sql(question: str, schema_ddl: str) -> str:
+    """Call HF model to convert NL question into SQL."""
     prompt = textwrap.dedent(f"""
         ### Task
         Translate the following natural language question into ONE valid SQLite SQL query.
@@ -41,93 +42,84 @@ def nlp_to_sql(nl_query: str, schema_ddl: str) -> str:
         {schema_ddl}
 
         ### Question
-        {nl_query}
+        {question}
 
         ### SQL
     """)
-    payload = {
-        "inputs": prompt,
-        "parameters": {"max_new_tokens": 256}
-    }
+    payload = {"inputs": prompt, "parameters": {"max_new_tokens": 256}}
 
     try:
-        resp = requests.post(API_URL, headers=HEADERS, json=payload, timeout=60)
+        r = requests.post(API_URL, headers=HEADERS, json=payload, timeout=60)
     except Exception as e:
-        return f"✖ Exception contacting API: {e}"
+        return f"✖ Connection error: {e}"
 
-    if resp.status_code != 200:
-        return f"✖ API error {resp.status_code}: {resp.text}"
+    if r.status_code != 200:
+        return f"✖ API error {r.status_code}: {r.text}"
 
+    # Parse JSON
     try:
-        output = resp.json()
-    except ValueError:
-        return "✖ Non-JSON response from API."
-
-    if not isinstance(output, list) or not output:
-        return "✖ Model returned empty output."
+        generated = r.json()[0]["generated_text"]
+    except Exception:
+        return "✖ Invalid JSON response."
 
-    generated = output[0].get("generated_text", "")
-    if "### SQL" in generated:
-        sql = generated.split("### SQL")[-1].strip()
-    else:
-        sql = generated.strip()
+    # Extract SQL
+    sql = generated.split("### SQL")[-1].strip()
+    return sql or "✖ Empty SQL returned."
 
-    return sql or "✖ Empty SQL string returned."
-
-# -----------------------------
-# Pipeline: NL → SQL → Execute
-# -----------------------------
+# -------------------------------------------------
+# 4. PIPELINE: NL → SQL → EXECUTE
+# -------------------------------------------------
 def run_pipeline(nl_query: str):
+    start = time.time()
     trace = []
-    t0 = time.time()
 
-    # build db if needed
+    # DB setup
     create_db_if_needed()
 
-    # load schema for prompt
+    # Load schema
     with open(SCHEMA_FILE) as f:
         schema_ddl = f.read()
-    trace.append(("Schema loaded", f"{len(schema_ddl.splitlines())} lines"))
+    trace.append(("Schema", "loaded"))
 
-    # convert NL → SQL
+    # Convert NL → SQL
     sql_query = nlp_to_sql(nl_query, schema_ddl)
-    trace.append(("LLM output", sql_query))
+    trace.append(("LLM", sql_query))
 
-    # execute SQL
+    # Execute SQL
     try:
         with sqlite3.connect(DB_PATH) as conn:
             cur = conn.execute(sql_query)
             rows = cur.fetchall()
             cols = [d[0] for d in cur.description] if cur.description else []
         result = {"columns": cols, "rows": rows}
-        trace.append(("Execution", f"{len(rows)} rows"))
+        trace.append(("Exec", f"{len(rows)} rows"))
     except Exception as e:
         result = {"error": str(e)}
-        trace.append(("Execution error", str(e)))
-
-    trace.append(("Total time", f"{time.time() - t0:.2f}s"))
-
-    # stringify outputs for UI
-    result_json = json.dumps(result, indent=2, ensure_ascii=False)
-    trace_log   = "\n".join(f"{s}: {m}" for s, m in trace)
-
-    return sql_query, result_json, trace_log
-
-# -----------------------------
-# Gradio UI
-# -----------------------------
-with gr.Blocks(title="NLP ➜ SQL (SQLite, HF Hub)") as demo:
-    gr.Markdown("## NLP → SQL Query Demo (auto-build SQLite, no terminal required)")
-    nl_in  = gr.Textbox(label="Natural-language question",
-                        placeholder="e.g. List all employees in Engineering hired after 2021")
+        trace.append(("Exec error", str(e)))
+
+    trace.append(("Time", f"{time.time() - start:.2f}s"))
+
+    return (
+        sql_query,
+        json.dumps(result, indent=2, ensure_ascii=False),
+        "\n".join(f"{s}: {m}" for s, m in trace),
+    )
+
+# -------------------------------------------------
+# 5. GRADIO UI
+# -------------------------------------------------
+with gr.Blocks(title="NLP → SQL (SQLite, HF Hub)") as demo:
+    gr.Markdown("### NLP → SQL demo • SQLite backend • Hugging Face Inference API")
+    query_in = gr.Textbox(
+        label="Natural-language question",
+        placeholder="e.g. List all employees in Engineering hired after 2021",
+    )
     with gr.Row():
-        sql_out    = gr.Code(label="Generated SQL", language="sql")
-        result_out = gr.Code(label="Query result (JSON)")
+        sql_out = gr.Code(label="Generated SQL", language="sql")
+        res_out = gr.Code(label="Query result (JSON)")
     trace_out = gr.Textbox(label="Trace", lines=6)
-
     run_btn = gr.Button("Run")
-    run_btn.click(fn=run_pipeline, inputs=nl_in,
-                  outputs=[sql_out, result_out, trace_out])
+    run_btn.click(run_pipeline, query_in, [sql_out, res_out, trace_out])
 
 if __name__ == "__main__":
     demo.launch()