Update app.py

app.py

@@ -1,113 +1,154 @@
-# 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
-import torch
-import pandas as pd
 from functools import lru_cache
-
-from datasets import load_dataset
-from deep_translator import GoogleTranslator
+from huggingface_hub import hf_hub_download
 from transformers import TapexTokenizer, BartForConditionalGeneration
+from deep_translator import GoogleTranslator
+import os, json, pandas as pd, torch
 
-# ------------ Config ------------
+# ------------------------
+# 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 ------------
+WIKISQL_REPO = os.getenv("WIKISQL_REPO", "Salesforce/wikisql")  # dataset oficial
+SPLIT = os.getenv("TABLE_SPLIT", "validation")                  # "validation" == dev en WikiSQL
+INDEX = int(os.getenv("TABLE_INDEX", "10"))
+MAX_ROWS = int(os.getenv("MAX_ROWS", "12"))
+
+# ------------------------
+# App
+# ------------------------
 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=["*"],
+    allow_origins=["*"], allow_methods=["*"], allow_headers=["*"], allow_credentials=True,
 )
 
-# ------------ 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
+class NLQuery(BaseModel):
+    nl_query: str
+
+# ------------------------
+# Modelo
+# ------------------------
+tok = TapexTokenizer.from_pretrained(HF_MODEL_ID)
+model = BartForConditionalGeneration.from_pretrained(HF_MODEL_ID)
+if torch.cuda.is_available():
+    model = model.to("cuda")
+
+# ------------------------
+# Util: carga WikiSQL (JSONL)
+# ------------------------
+def _read_jsonl(path):
+    with open(path, "r", encoding="utf-8") as f:
+        for line in f:
+            if line.strip():
+                yield json.loads(line)
+
+def _download_file(filename: str) -> str:
+    # descarga desde el dataset hug
+    return hf_hub_download(repo_id=WIKISQL_REPO, filename=filename, repo_type="dataset")
 
 @lru_cache(maxsize=32)
-def get_table(split: str, index: int, max_rows: int) -> pd.DataFrame:
+def get_table_from_wikisql(split: str, index: int, max_rows: int) -> pd.DataFrame:
     """
-    Carga tabla de WikiSQL usando archivos parquet (sin script local).
-    Compatible con Spaces free (CPU).
+    Carga la tabla de WikiSQL sin scripts, usando directamente los JSONL del repo:
+      - dev.jsonl            (validation = 'dev' en terminología original)
+      - dev.tables.jsonl
+    Si cambias split a 'train' o 'test', intenta los nombres equivalentes.
     """
-    import pyarrow.parquet as pq
-    import tempfile
-    from huggingface_hub import hf_hub_download
-
-    # Descarga directa del dataset en parquet
-    parquet_path = hf_hub_download(
-        repo_id="Salesforce/wikisql",
-        filename=f"data/{split}-00000-of-00001.parquet",
-        repo_type="dataset"
-    )
-
-    # Lee parquet con pyarrow/pandas
-    table = pq.read_table(parquet_path)
-    df_full = table.to_pandas()
-
-    # Cada registro tiene una columna "table" con estructura compleja → parseamos
-    example = df_full.iloc[index]
-    header = [str(h) for h in example["table"]["header"]]
-    rows = example["table"]["rows"][:max_rows]
+    # Mapeo simple: validation->dev, train->train, test->test
+    split_map = {"validation": "dev", "dev": "dev", "train": "train", "test": "test"}
+    base = split_map.get(split.lower(), "dev")
+
+    # Posibles nombres de archivo en el repo (algunos mirrors usan variantes)
+    qa_candidates = [f"data/{base}.jsonl", f"data/{base}.json", f"{base}.jsonl"]
+    tbl_candidates = [f"data/{base}.tables.jsonl", f"{base}.tables.jsonl"]
+
+    qa_path = None
+    tbl_path = None
+
+    # Descarga QA
+    for cand in qa_candidates:
+        try:
+            qa_path = _download_file(cand)
+            break
+        except Exception:
+            continue
+    if qa_path is None:
+        raise RuntimeError(f"No se encontró el archivo QA para split={split}. Intentos: {qa_candidates}")
+
+    # Descarga tablas
+    for cand in tbl_candidates:
+        try:
+            tbl_path = _download_file(cand)
+            break
+        except Exception:
+            continue
+    if tbl_path is None:
+        raise RuntimeError(f"No se encontró el archivo de tablas para split={split}. Intentos: {tbl_candidates}")
+
+    # Leemos la pregunta N (para tomar su table_id) — si no necesitas la pregunta, puedes omitir esto
+    qa_list = list(_read_jsonl(qa_path))
+    if not (0 <= index < len(qa_list)):
+        raise IndexError(f"index={index} fuera de rango (0..{len(qa_list)-1}) para split={split}")
+    table_id = qa_list[index].get("table_id") or qa_list[index].get("table", {}).get("id")
+    if table_id is None:
+        raise RuntimeError("No se pudo extraer 'table_id' del registro de QA.")
+
+    # Buscamos esa tabla en dev.tables.jsonl
+    header, rows = None, None
+    for obj in _read_jsonl(tbl_path):
+        if obj.get("id") == table_id:
+            header = [str(h) for h in obj["header"]]
+            rows = obj["rows"]
+            break
+    if header is None or rows is None:
+        raise RuntimeError(f"No se encontró la tabla con id={table_id} en {os.path.basename(tbl_path)}")
+
+    # recortamos filas
+    rows = 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 ------------
+# ------------------------
+# Endpoints
+# ------------------------
 @app.get("/api/health")
 def health():
-    return {"ok": True, "model": HF_MODEL_ID, "split": TABLE_SPLIT, "index": TABLE_INDEX}
+    return {"ok": True, "model": HF_MODEL_ID, "split": SPLIT, "index": 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")}
+        df = get_table_from_wikisql(SPLIT, INDEX, MAX_ROWS)
+        return {"columns": df.columns.tolist(), "rows": df.head(8).to_dict(orient="records")}
     except Exception as e:
-        return JSONResponse(status_code=500, content={"error": str(e)})
+        return {"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)
+        text = (q.nl_query or "").strip()
+        if not text:
+            raise ValueError("Consulta vacía.")
+
+        # Traducción ES->EN si detectamos acentos u otros
+        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 = get_table_from_wikisql(SPLIT, INDEX, MAX_ROWS)
+        enc = tok(table=df, query=query_en, return_tensors="pt", truncation=True)
+        if torch.cuda.is_available():
+            enc = {k: v.to("cuda") for k, v in enc.items()}
+        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}
+
+        return {
+            "consulta_original": text,
+            "consulta_traducida": query_en,
+            "sql_generado": sql
+        }
     except Exception as e:
         raise HTTPException(status_code=500, detail=str(e))