File size: 3,993 Bytes
d5ccd98 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 | import os, sys, time, json, sqlite3, textwrap, requests
import gradio as gr
# -------------------------------------------------
# 1. CONFIGURATION
# -------------------------------------------------
MODEL_ID = "gpt2" # always exists; later swap for sqlcoder
API_URL = f"https://api-inference.huggingface.co/models/{MODEL_ID}"
HF_TOKEN = os.getenv("HF_TOKEN")
if not HF_TOKEN:
raise RuntimeError(
"HF_TOKEN not found. Go to Space → Settings → Secrets and add it."
)
HEADERS = {"Authorization": f"Bearer {HF_TOKEN}"}
DB_PATH = "company.db"
SCHEMA_FILE = "schema.sql"
# -------------------------------------------------
# 2. UTIL: BUILD DB IF NEEDED
# -------------------------------------------------
def create_db_if_needed():
if os.path.exists(DB_PATH):
return
with open(SCHEMA_FILE) as f, sqlite3.connect(DB_PATH) as conn:
conn.executescript(f.read())
# -------------------------------------------------
# 3. UTIL: CALL HF MODEL (with token debug)
# -------------------------------------------------
def nlp_to_sql(question: str, schema_ddl: str) -> str:
prompt = textwrap.dedent(f"""
Translate the following natural language question into a single valid SQLite SQL query.
### Schema
{schema_ddl}
### Question
{question}
### SQL
""")
payload = {"inputs": prompt, "parameters": {"max_new_tokens": 64}}
# ---------- DEBUG ----------
print("=" * 60, file=sys.stderr)
print("DEBUG URL:", API_URL, file=sys.stderr)
print("DEBUG token starts with:", HF_TOKEN[:8], file=sys.stderr)
# ---------------------------
try:
r = requests.post(API_URL, headers=HEADERS, json=payload, timeout=60)
except Exception as e:
return f"[ConnErr] {e}"
# ---------- MORE DEBUG ----------
print("DEBUG status:", r.status_code, file=sys.stderr)
print("DEBUG first 200 bytes:", r.text[:200], file=sys.stderr)
print("=" * 60, file=sys.stderr)
# -------------------------------
if r.status_code != 200:
return f"[API {r.status_code}] {r.text[:100]}"
try:
generated = r.json()[0]["generated_text"]
except Exception as e:
return f"[JSONErr] {e}"
return generated.split("### SQL")[-1].strip() or "[Empty SQL]"
# -------------------------------------------------
# 4. PIPELINE
# -------------------------------------------------
def run_pipeline(query: str):
t0, trace = time.time(), []
create_db_if_needed()
with open(SCHEMA_FILE) as f:
schema = f.read()
trace.append(("Schema", "loaded"))
sql = nlp_to_sql(query, schema)
trace.append(("LLM", sql))
try:
with sqlite3.connect(DB_PATH) as conn:
cur = conn.execute(sql)
rows = cur.fetchall()
cols = [d[0] for d in cur.description] if cur.description else []
result = {"columns": cols, "rows": rows}
trace.append(("Exec", f"{len(rows)} rows"))
except Exception as e:
result = {"error": str(e)}
trace.append(("Exec error", str(e)))
trace.append(("Time", f"{time.time() - t0:.2f}s"))
return sql, json.dumps(result, indent=2), "\n".join(f"{s}: {m}" for s, m in trace)
# -------------------------------------------------
# 5. UI
# -------------------------------------------------
with gr.Blocks(title="Debug HF Token & API") as demo:
gr.Markdown("### Debugging HF TOKEN → API (uses GPT-2)")
q = gr.Textbox(label="Question", placeholder="e.g., How many employees?")
with gr.Row():
sql_box = gr.Code(label="SQL / debug output")
res_box = gr.Code(label="Result / error")
trace_box = gr.Textbox(label="Trace")
btn = gr.Button("Run")
btn.click(run_pipeline, q, [sql_box, res_box, trace_box])
if __name__ == "__main__":
demo.launch()
|