Spaces:

stvnnnnnn
/

nl2sql-tapex-backend

Sleeping

App Files Files Community

stvnnnnnn commited on Oct 13

Commit

c1c647f

verified ·

1 Parent(s): 719fbce

Update app.py

Browse files

Files changed (1) hide show

app.py +128 -38

app.py CHANGED Viewed

@@ -1,75 +1,165 @@
 from fastapi import FastAPI, HTTPException
 from fastapi.middleware.cors import CORSMiddleware
 from pydantic import BaseModel
-from deep_translator import GoogleTranslator
 from datasets import load_dataset
 from transformers import TapexTokenizer, BartForConditionalGeneration
-import pandas as pd, torch, os
-# === Config ===
 HF_MODEL_ID = os.getenv("HF_MODEL_ID", "stvnnnnnn/tapex-wikisql-best")
 TABLE_SPLIT = os.getenv("TABLE_SPLIT", "validation")
 TABLE_INDEX = int(os.getenv("TABLE_INDEX", "10"))
-MAX_ROWS    = int(os.getenv("MAX_ROWS", "12"))
-torch.set_num_threads(1)
-app = FastAPI(title="NL→SQL – TAPEX + WikiSQL (HF Space)")
-# CORS: permite que Vercel (o cualquier origen) consuma la API
 app.add_middleware(
     CORSMiddleware,
-    allow_origins=["*"], allow_credentials=False,
-    allow_methods=["*"], allow_headers=["*"],
 )
-# Carga modelo/tokenizer con bajo pico de RAM (CPU)
-tok = TapexTokenizer.from_pretrained(HF_MODEL_ID)
-model = BartForConditionalGeneration.from_pretrained(
-    HF_MODEL_ID, low_cpu_mem_usage=True
-).to("cpu")
-class NLQuery(BaseModel):
-    nl_query: str
-def get_example(split, index):
-    # streaming para no cargar todo WikiSQL en RAM
-    ds = load_dataset("Salesforce/wikisql", split=split, streaming=True)
-    for i, ex in enumerate(ds):
-        if i == index:
-            return ex
-    raise IndexError("Index fuera de rango")
-def load_table(split=TABLE_SPLIT, index=TABLE_INDEX, max_rows=MAX_ROWS):
-    ex = get_example(split, index)
     header = [str(h) for h in ex["table"]["header"]]
     rows = ex["table"]["rows"][:max_rows]
-    return pd.DataFrame(rows, columns=header)
 @app.get("/api/health")
 def health():
     return {"ok": True, "model": HF_MODEL_ID, "split": TABLE_SPLIT, "index": TABLE_INDEX}
 @app.get("/api/preview")
 def preview():
-    df = load_table()
-    return {"columns": df.columns.tolist(), "rows": df.head(8).to_dict(orient="records")}
 @app.post("/api/nl2sql")
 def nl2sql(q: NLQuery):
     try:
-        text = (q.nl_query or "").strip()
-        if not text:
-            raise ValueError("Consulta vacía.")
-        is_ascii = all(ord(c) < 128 for c in text)
-        query_en = text if is_ascii else GoogleTranslator(source="auto", target="en").translate(text)
-        df = load_table()
-        enc = tok(table=df, query=query_en, return_tensors="pt", truncation=True)
-        out = model.generate(**enc, max_length=160, num_beams=1)
-        sql = tok.batch_decode(out, skip_special_tokens=True)[0]
-        return {"consulta_original": text, "consulta_traducida": query_en, "sql_generado": sql}
     except Exception as e:
         raise HTTPException(status_code=500, detail=str(e))

+# app.py  — NL→SQL (TAPEX + WikiSQL) backend for HF Spaces
 from fastapi import FastAPI, HTTPException
+from fastapi.responses import HTMLResponse, JSONResponse
 from fastapi.middleware.cors import CORSMiddleware
 from pydantic import BaseModel
+import os
+import torch
+import pandas as pd
+from functools import lru_cache
 from datasets import load_dataset
+from deep_translator import GoogleTranslator
 from transformers import TapexTokenizer, BartForConditionalGeneration
+# --------- Configuración y defaults ----------
 HF_MODEL_ID = os.getenv("HF_MODEL_ID", "stvnnnnnn/tapex-wikisql-best")
 TABLE_SPLIT = os.getenv("TABLE_SPLIT", "validation")
 TABLE_INDEX = int(os.getenv("TABLE_INDEX", "10"))
+MAX_ROWS    = int(os.getenv("MAX_ROWS", "100"))  # límite prudente en CPU
+# Asegura caché escribible en Space
+os.environ["HF_HOME"] = "/app/.cache/huggingface"
+os.makedirs(os.environ["HF_HOME"], exist_ok=True)
+# --------- App & CORS ----------
+app = FastAPI(title="NL→SQL – TAPEX + WikiSQL (HF Space)")
 app.add_middleware(
     CORSMiddleware,
+    allow_origins=["*"],          # cambia a tu dominio de Vercel cuando lo tengas
+    allow_credentials=True,
+    allow_methods=["*"],
+    allow_headers=["*"],
 )
+# --------- Carga perezosa (lazy) ---------
+@lru_cache(maxsize=1)
+def get_model_and_tokenizer():
+    tok = TapexTokenizer.from_pretrained(HF_MODEL_ID)
+    # En Spaces free es CPU; no uses device_map="auto" para evitar dependencia de accelerate
+    model = BartForConditionalGeneration.from_pretrained(HF_MODEL_ID)
+    model.eval()
+    return tok, model
+@lru_cache(maxsize=32)
+def get_table(split: str, index: int, max_rows: int) -> pd.DataFrame:
+    """
+    Carga una tabla de WikiSQL y devuelve un DataFrame.
+    Evitamos el revision parquet (que a veces falla en Spaces) y usamos el split normal.
+    """
+    ds = load_dataset("Salesforce/wikisql", split=split)
+    if index < 0 or index >= len(ds):
+        raise IndexError(f"TABLE_INDEX fuera de rango (0..{len(ds)-1}).")
+    ex = ds[index]
     header = [str(h) for h in ex["table"]["header"]]
     rows = ex["table"]["rows"][:max_rows]
+    df = pd.DataFrame(rows, columns=header)
+    # Normaliza nombres de columnas a string
+    df.columns = [str(c) for c in df.columns]
+    return df
+# --------- Esquema de petición ----------
+class NLQuery(BaseModel):
+    nl_query: str
+# --------- Mini UI en / ----------
+INDEX_HTML = """
+<!doctype html>
+<html lang="es">
+<meta charset="utf-8">
+<title>NL→SQL (TAPEX + WikiSQL)</title>
+<style>
+  body{font-family:system-ui,-apple-system,Segoe UI,Roboto,Ubuntu,sans-serif;max-width:860px;margin:30px auto;padding:0 16px;color:#eaeaea;background:#0f1115}
+  h1{font-size:1.6rem;margin:0 0 8px}
+  .card{background:#171923;border:1px solid #232736;border-radius:12px;padding:16px;margin:18px 0}
+  input,button,select{font-size:1rem}
+  input{width:100%;padding:10px 12px;border-radius:8px;border:1px solid #2a3042;background:#0f1115;color:#eaeaea}
+  button{padding:10px 16px;border-radius:8px;border:1px solid #2a3042;background:#1f2433;color:#eaeaea;cursor:pointer}
+  pre{white-space:pre-wrap;background:#0f1115;border:1px solid #2a3042;padding:12px;border-radius:8px}
+  .row{display:flex;gap:12px;align-items:center}
+</style>
+<h1>🧠 NL → SQL (TAPEX + WikiSQL)</h1>
+<div class="card">
+  <p><b>Backend:</b> este Space ofrece endpoints REST; prueba una consulta:</p>
+  <div class="row">
+    <input id="q" placeholder="Ej.: Muestra los jugadores que son Guards." />
+    <button onclick="run()">Generar SQL</button>
+    <button onclick="prev()">Ver preview tabla</button>
+  </div>
+  <p style="font-size:.9rem;opacity:.8">Swagger: <a href="./docs" target="_blank">/docs</a> · Salud: <a href="./api/health" target="_blank">/api/health</a></p>
+  <pre id="out">Listo para generar...</pre>
+</div>
+<script>
+async function run(){
+  const q = document.getElementById('q').value.trim();
+  const r = await fetch('./api/nl2sql', {method:'POST', headers:{'Content-Type':'application/json'}, body: JSON.stringify({nl_query: q}) });
+  document.getElementById('out').textContent = JSON.stringify(await r.json(), null, 2);
+}
+async function prev(){
+  const r = await fetch('./api/preview');
+  document.getElementById('out').textContent = JSON.stringify(await r.json(), null, 2);
+}
+</script>
+</html>
+"""
+@app.get("/", response_class=HTMLResponse)
+def home():
+    return INDEX_HTML
+# --------- Rutas API ----------
 @app.get("/api/health")
 def health():
     return {"ok": True, "model": HF_MODEL_ID, "split": TABLE_SPLIT, "index": TABLE_INDEX}
 @app.get("/api/preview")
 def preview():
+    try:
+        df = get_table(TABLE_SPLIT, TABLE_INDEX, MAX_ROWS)
+        # Regresa primeras 8 filas para no saturar
+        data = df.head(8).to_dict(orient="records")
+        return {"columns": list(df.columns), "rows": data}
+    except Exception as e:
+        # Envía mensaje simple (para no saturar logs)
+        return JSONResponse(status_code=500, content={"error": str(e)})
 @app.post("/api/nl2sql")
 def nl2sql(q: NLQuery):
+    nl = (q.nl_query or "").strip()
+    if not nl:
+        raise HTTPException(status_code=400, detail="Consulta vacía.")
+    # Traducción ES→EN si detectamos caracteres no ASCII
     try:
+        is_ascii = all(ord(c) < 128 for c in nl)
+        nl_en = nl if is_ascii else GoogleTranslator(source="auto", target="en").translate(nl)
+    except Exception:
+        # Si la traducción falla, seguimos con el texto original
+        nl_en = nl
+    try:
+        df = get_table(TABLE_SPLIT, TABLE_INDEX, MAX_ROWS)
+        tok, model = get_model_and_tokenizer()
+        # Tokenización limitada
+        enc = tok(table=df, query=nl_en, return_tensors="pt", truncation=True, max_length=512)
+        with torch.inference_mode():
+            out = model.generate(**enc, max_length=160, num_beams=1)
+        sql = tok.batch_decode(out, skip_special_tokens=True)[0]
+        return {"consulta_original": nl, "consulta_traducida": nl_en, "sql_generado": sql}
     except Exception as e:
         raise HTTPException(status_code=500, detail=str(e))