text2sql-demo / src /app /schema_builder.py
minimew's picture
Upload src/app/schema_builder.py with huggingface_hub
a6b7c38 verified
Raw
History Blame Contribute Delete
1.47 kB
"""
Schema builder cho demo — tai schema tu spider_data/tables.json va build dung
input string format giong luc training (tai su dung build_model_input() /
build_schema_string() trong process_spider.py, dam bao khop 100% format).
"""
import sys
import json
from pathlib import Path
_SRC_DIR = Path(__file__).resolve().parent.parent # src/
if str(_SRC_DIR) not in sys.path:
sys.path.insert(0, str(_SRC_DIR))
from process_spider import build_model_input, build_schema_string # noqa: E402
_TABLES_PATH = _SRC_DIR.parent / "spider_data" / "tables.json"
_schemas_cache = None
def _load_all_schemas() -> dict:
global _schemas_cache
if _schemas_cache is None:
with open(_TABLES_PATH, "r", encoding="utf-8") as f:
tables = json.load(f)
_schemas_cache = {db["db_id"]: db for db in tables}
return _schemas_cache
def load_db_schema(db_id: str) -> dict:
schemas = _load_all_schemas()
if db_id not in schemas:
raise ValueError(f"db_id '{db_id}' not found in {_TABLES_PATH.name}")
return schemas[db_id]
def get_schema_string(db_id: str) -> str:
"""Schema string thuan (khong co prefix 'question: ... | schema: ') -- dung de hien thi UI."""
db = load_db_schema(db_id)
return build_schema_string(db)
def build_input(question: str, db_id: str) -> str:
"""Full model input string -- dung truoc khi generate SQL."""
db = load_db_schema(db_id)
return build_model_input(question, db)