File size: 1,468 Bytes
a6b7c38
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
"""
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)