Spaces:

stvnnnnnn
/

nl2sql-tapex-backend

Running

App Files Files Community

stvnnnnnn commited on about 1 month ago

Commit

619ce97

verified ·

1 Parent(s): c1c647f

Update app.py

Browse files

Files changed (1) hide show

app.py +14 -83

app.py CHANGED Viewed

@@ -1,8 +1,8 @@
-# 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
@@ -14,151 +14,82 @@ 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:

+# app.py — NL→SQL (TAPEX + WikiSQL) backend (solo API)
 from fastapi import FastAPI, HTTPException
 from fastapi.middleware.cors import CORSMiddleware
+from fastapi.responses import JSONResponse
 from pydantic import BaseModel
 import os
 from deep_translator import GoogleTranslator
 from transformers import TapexTokenizer, BartForConditionalGeneration
+# ------------ 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", "100"))
+# Caché HF escribible en Spaces
 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 (API)")
 app.add_middleware(
     CORSMiddleware,
+    allow_origins=["*"],      # cámbialo al dominio de Vercel cuando lo tengas
     allow_credentials=True,
     allow_methods=["*"],
     allow_headers=["*"],
 )
+# ------------ Carga perezosa ------------
 @lru_cache(maxsize=1)
 def get_model_and_tokenizer():
     tok = TapexTokenizer.from_pretrained(HF_MODEL_ID)
     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:
+    ds = load_dataset("Salesforce/wikisql", split=split)  # evita revision parquet
     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)
     df.columns = [str(c) for c in df.columns]
     return df
+# ------------ Schemas ------------
 class NLQuery(BaseModel):
     nl_query: str
+# ------------ Endpoints ------------
 @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)
+        return {"columns": list(df.columns), "rows": df.head(8).to_dict(orient="records")}
     except Exception as e:
         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 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:
         nl_en = nl
     try:
         df = get_table(TABLE_SPLIT, TABLE_INDEX, MAX_ROWS)
         tok, model = get_model_and_tokenizer()
         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: