Update app.py

app.py

@@ -2,19 +2,18 @@ from fastapi import FastAPI, HTTPException
 from fastapi.middleware.cors import CORSMiddleware
 from pydantic import BaseModel
 from functools import lru_cache
-from huggingface_hub import hf_hub_download
 from transformers import TapexTokenizer, BartForConditionalGeneration
 from deep_translator import GoogleTranslator
+from pathlib import Path
 import os, json, pandas as pd, torch
 
 # ------------------------
 # Config
 # ------------------------
-HF_MODEL_ID = os.getenv("HF_MODEL_ID", "stvnnnnnn/tapex-wikisql-best")
-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"))
+HF_MODEL_ID  = os.getenv("HF_MODEL_ID",  "stvnnnnnn/tapex-wikisql-best")
+SPLIT        = os.getenv("TABLE_SPLIT",  "validation")   # "validation" ~ "dev"
+INDEX        = int(os.getenv("TABLE_INDEX", "10"))
+MAX_ROWS     = int(os.getenv("MAX_ROWS",    "12"))
 
 # ------------------------
 # App
@@ -37,80 +36,57 @@ if torch.cuda.is_available():
     model = model.to("cuda")
 
 # ------------------------
-# Util: carga WikiSQL (JSONL)
+# Utilidades de carga robustas
 # ------------------------
-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")
+def _read_json_or_jsonl(p: Path) -> dict:
+    """
+    Lee un JSON normal (.json) o un JSONL (.jsonl) y devuelve el primer objeto.
+    """
+    txt = p.read_text(encoding="utf-8").strip()
+    if p.suffix.lower() == ".jsonl":
+        for line in txt.splitlines():
+            s = line.strip()
+            if s:
+                return json.loads(s)
+        raise ValueError(f"{p} está vacío.")
+    return json.loads(txt)
 
 @lru_cache(maxsize=32)
-def get_table_from_wikisql(split: str, index: int, max_rows: int) -> pd.DataFrame:
+def get_table(split: str, index: int, max_rows: int) -> pd.DataFrame:
     """
-    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.
+    1) Intenta cargar ./data/<split>.json o ./data/<split>.jsonl (mapeando 'validation'->'dev').
+    2) Si no existe, hace fallback a un ejemplo de WikiSQL (conversión Parquet oficial).
     """
-    # 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
+    base_dir = Path(__file__).parent
+    data_dir = base_dir / "data"
+
+    # Normalizamos nombre local (para demo usamos 'dev')
+    local_name = "dev" if split.lower() in ("validation", "dev") else split.lower()
+
+    # 1) Buscar archivo local
+    for candidate in (data_dir / f"{local_name}.json", data_dir / f"{local_name}.jsonl"):
+        if candidate.exists():
+            js = _read_json_or_jsonl(candidate)
+            header = [str(h) for h in js["header"]]
+            rows   = js["rows"][:max_rows]
+            df = pd.DataFrame(rows, columns=header)
+            df.columns = [str(c) for c in df.columns]
+            return df
+
+    # 2) Fallback: cargar un ejemplo del dataset WikiSQL (Parquet convertido)
+    try:
+        from datasets import load_dataset  # import diferido para arrancar más rápido
+        ds = load_dataset("Salesforce/wikisql", split="validation", revision="refs/convert/parquet")
+        if not (0 <= index < len(ds)):
+            index = 0  # seguridad
+        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
+    except Exception as e:
+        raise RuntimeError(f"No se pudo obtener una tabla: {e}")
 
 # ------------------------
 # Endpoints
@@ -122,7 +98,7 @@ def health():
 @app.get("/api/preview")
 def preview():
     try:
-        df = get_table_from_wikisql(SPLIT, INDEX, MAX_ROWS)
+        df = get_table(SPLIT, INDEX, MAX_ROWS)
         return {"columns": df.columns.tolist(), "rows": df.head(8).to_dict(orient="records")}
     except Exception as e:
         return {"error": str(e)}
@@ -134,14 +110,15 @@ def nl2sql(q: NLQuery):
         if not text:
             raise ValueError("Consulta vacía.")
 
-        # Traducción ES->EN si detectamos acentos u otros
+        # Traducción ES->EN si detectamos caracteres no ASCII
         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)
+        df  = get_table(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]