Upload src/process_spider.py with huggingface_hub

src/process_spider.py

@@ -0,0 +1,338 @@
+"""
+Spider Data Processing Script v5
+Build schema-aware inputs for Spider to fine-tune T5/CodeT5/BART.
+
+Input:  spider_data/{train_spider,dev,test}.json + tables.json + test_tables.json
+Output: processed/spider/{train,dev,test}.json
+
+Usage:
+    python src/process_spider.py
+    python src/process_spider.py --no_foreign_keys
+    python src/process_spider.py --no_column_types
+    python src/process_spider.py --no_primary_keys
+    python src/process_spider.py --no_lowercase_identifiers
+    python src/process_spider.py --oracle_schema_link --output_dir processed/spider_oracle
+        (RQ6 ablation: filter schema to keep only tables used in gold SQL — ceiling-effect
+        study for WRONG_TABLE error; not applied to test split which has no gold SQL)
+"""
+
+import re
+import json
+import argparse
+import sys
+from pathlib import Path
+
+try:
+    sys.stdout.reconfigure(encoding='utf-8')
+except Exception:
+    pass
+
+try:
+    import sqlglot
+    import sqlglot.expressions as exp
+    HAS_SQLGLOT = True
+except ImportError:
+    HAS_SQLGLOT = False
+    print("WARNING: sqlglot not installed. Run: pip install sqlglot")
+
+
+# ── SQL utilities ──────────────────────────────────────────────────────────────
+
+def normalize_sql(sql):
+    """Normalize SQL for Exact Match — preserve string literals (Spider DB is mixed-case)."""
+    if not sql:
+        return ""
+    if HAS_SQLGLOT:
+        try:
+            return sqlglot.parse_one(sql, read='sqlite').sql(
+                dialect='sqlite', pretty=False
+            ).strip()
+        except Exception:
+            pass
+    sql = sql.strip()
+    sql = re.sub(r'\s+', ' ', sql)
+    sql = re.sub(r'\s*,\s*', ' , ', sql)
+    sql = re.sub(r'\s*(=|!=|<>|>=|<=|>|<)\s*', r' \1 ', sql)
+    return sql.strip()
+
+
+def normalize_question(question):
+    return re.sub(r'\s+', ' ', question.strip()).rstrip('.?!').strip()
+
+
+def normalize_identifier(name, lowercase=True):
+    return name.strip().lower() if lowercase else name.strip()
+
+
+def normalize_primary_keys(primary_keys):
+    pk_indices = set()
+    for pk in primary_keys:
+        if isinstance(pk, list):
+            pk_indices.update(pk)
+        else:
+            pk_indices.add(pk)
+    return pk_indices
+
+
+def check_sql_validity(sql):
+    if not sql or not sql.strip():
+        return False
+    if HAS_SQLGLOT:
+        try:
+            sqlglot.parse_one(sql, read='sqlite')
+            return True
+        except Exception:
+            return False
+    return sql.strip().upper().startswith('SELECT')
+
+
+# ── Schema builder ─────────────────────────────────────────────────────────────
+
+def build_schema_string(db, use_foreign_keys=True, use_column_types=True,
+                        use_primary_keys=True, lowercase_identifiers=True):
+    table_names = [normalize_identifier(n, lowercase_identifiers)
+                   for n in db['table_names_original']]
+    col_names   = [(ti, normalize_identifier(cn, lowercase_identifiers))
+                   for ti, cn in db['column_names_original']]
+    col_types    = db.get('column_types', [])
+    primary_keys = normalize_primary_keys(db.get('primary_keys', []))
+
+    table_cols = {i: [] for i in range(len(table_names))}
+    for col_idx, (ti, cn) in enumerate(col_names):
+        if ti == -1:
+            continue
+        col = cn
+        if use_primary_keys and col_idx in primary_keys:
+            col += "*"
+        if use_column_types and col_idx < len(col_types):
+            col += f":{col_types[col_idx]}"
+        table_cols[ti].append(col)
+
+    parts = []
+    for i, tname in enumerate(table_names):
+        cols = table_cols.get(i, [])
+        parts.append(f"{tname} ( {' , '.join(cols)} )" if cols else tname)
+    schema = " | ".join(parts)
+
+    if use_foreign_keys and db.get('foreign_keys'):
+        fk_parts = []
+        for c1, c2 in db['foreign_keys']:
+            t1 = table_names[col_names[c1][0]]
+            t2 = table_names[col_names[c2][0]]
+            fk_parts.append(f"{t1}.{col_names[c1][1]} = {t2}.{col_names[c2][1]}")
+        if fk_parts:
+            schema += " | foreign_keys: " + " , ".join(fk_parts)
+
+    return schema
+
+
+def build_model_input(question, db, use_foreign_keys=True, use_column_types=True,
+                      use_primary_keys=True, lowercase_identifiers=True):
+    schema = build_schema_string(db, use_foreign_keys, use_column_types,
+                                  use_primary_keys, lowercase_identifiers)
+    return f"question: {normalize_question(question)} | schema: {schema}"
+
+
+# ── Oracle schema-linking (RQ6 ablation) ────────────────────────────────────────
+# Ceiling-effect study: use gold SQL to filter schema to only tables actually used.
+# Measures WRONG_TABLE error reduction under perfect schema-linking.
+# NOT a real inference system (requires gold SQL in advance) — train/dev only.
+
+def extract_gold_tables(sql, table_names_original):
+    """Extract table names used in gold SQL via sqlglot AST (handles T1/T2 aliases
+    and nested subqueries). Falls back to all original tables on parse error or no
+    match — safe, never produces an empty schema."""
+    valid = {t.lower() for t in table_names_original}
+    if not HAS_SQLGLOT:
+        return valid
+    try:
+        tree = sqlglot.parse_one(sql, read='sqlite')
+        used = {t.name.lower() for t in tree.find_all(exp.Table)}
+    except Exception:
+        return valid
+    matched = used & valid
+    return matched if matched else valid
+
+
+def filter_db_to_tables(db, keep_tables_lower):
+    """Filter db dict to keep only tables in keep_tables_lower, remapping all
+    column/table indices so build_schema_string() works unchanged.
+    Foreign keys are kept only when both columns belong to retained tables."""
+    table_names = db['table_names_original']
+    keep_indices = [i for i, t in enumerate(table_names) if t.lower() in keep_tables_lower]
+    if not keep_indices:
+        return db  # an toan: khong loc duoc gi thi giu nguyen schema goc
+
+    index_map = {old: new for new, old in enumerate(keep_indices)}
+    new_table_names = [table_names[i] for i in keep_indices]
+
+    col_names    = db['column_names_original']
+    col_types    = db.get('column_types', [])
+    new_col_names, new_col_types, col_index_map = [], [], {}
+    for old_idx, (ti, cn) in enumerate(col_names):
+        if ti == -1 or ti in index_map:
+            new_ti = -1 if ti == -1 else index_map[ti]
+            col_index_map[old_idx] = len(new_col_names)
+            new_col_names.append([new_ti, cn])
+            if old_idx < len(col_types):
+                new_col_types.append(col_types[old_idx])
+
+    new_primary_keys = []
+    for pk in db.get('primary_keys', []):
+        if isinstance(pk, list):
+            filtered = [col_index_map[c] for c in pk if c in col_index_map]
+            if filtered:
+                new_primary_keys.append(filtered if len(filtered) > 1 else filtered[0])
+        elif pk in col_index_map:
+            new_primary_keys.append(col_index_map[pk])
+
+    new_foreign_keys = [
+        [col_index_map[c1], col_index_map[c2]]
+        for c1, c2 in db.get('foreign_keys', [])
+        if c1 in col_index_map and c2 in col_index_map
+    ]
+
+    return {
+        **db,
+        'table_names_original':  new_table_names,
+        'column_names_original': new_col_names,
+        'column_types':          new_col_types,
+        'primary_keys':          new_primary_keys,
+        'foreign_keys':          new_foreign_keys,
+    }
+
+
+# ── Data loading ───────────────────────────────────────────────────────────────
+
+def load_schemas(tables_path):
+    with open(tables_path, 'r', encoding='utf-8') as f:
+        tables = json.load(f)
+    schemas = {db['db_id']: db for db in tables}
+    print(f"  Loaded {len(schemas):,} schemas from {Path(tables_path).name}")
+    return schemas
+
+
+# ── Process split ──────────────────────────────────────────────────────────────
+
+def process_split(data_path, schemas, split_name,
+                  use_foreign_keys=True, use_column_types=True,
+                  use_primary_keys=True, lowercase_identifiers=True,
+                  oracle_schema_link=False):
+    print(f"\n[{split_name.upper()}] Processing...")
+    with open(data_path, 'r', encoding='utf-8') as f:
+        data = json.load(f)
+
+    dataset, errors, not_found, invalid_sql = [], 0, 0, 0
+    is_test = (split_name == 'test')
+    # Oracle filtering only applies when gold SQL is available (train/dev).
+    # Test split has no gold SQL — skip.
+    apply_oracle = oracle_schema_link and not is_test
+    tables_before, tables_after = [], []
+
+    for idx, item in enumerate(data):
+        try:
+            db_id    = item['db_id']
+            question = normalize_question(item['question'])
+            if db_id not in schemas:
+                not_found += 1
+                continue
+
+            db  = schemas[db_id]
+            sql = item.get('query', item.get('SQL', '')).strip() if not is_test else ''
+
+            db_for_schema = db
+            if apply_oracle and sql:
+                keep_tables = extract_gold_tables(sql, db['table_names_original'])
+                db_for_schema = filter_db_to_tables(db, keep_tables)
+                tables_before.append(len(db['table_names_original']))
+                tables_after.append(len(db_for_schema['table_names_original']))
+
+            record = {
+                'input':    build_model_input(question, db_for_schema, use_foreign_keys,
+                                               use_column_types, use_primary_keys,
+                                               lowercase_identifiers),
+                'question': question,
+                'db_id':    db_id,
+            }
+
+            if not is_test:
+                normalized_sql = normalize_sql(sql)
+                valid          = check_sql_validity(sql)
+                if not valid:
+                    invalid_sql += 1
+                record['output']            = normalized_sql
+                record['raw_output']        = sql
+                record['normalized_output'] = normalized_sql
+                record['is_valid_sql']      = valid
+                if 'query_toks' in item:
+                    record['query_toks'] = item['query_toks']
+
+            dataset.append(record)
+        except Exception as e:
+            errors += 1
+            if errors <= 5:
+                print(f"  WARNING idx {idx}: {e}")
+
+    print(f"  Processed: {len(dataset):,} | Not found: {not_found} | "
+          f"Errors: {errors} | Invalid SQL: {invalid_sql}")
+    if apply_oracle and tables_before:
+        avg_before = sum(tables_before) / len(tables_before)
+        avg_after  = sum(tables_after) / len(tables_after)
+        reduction  = (1 - avg_after / avg_before) * 100 if avg_before else 0
+        print(f"  [oracle_schema_link] Avg tables/schema: {avg_before:.2f} -> "
+              f"{avg_after:.2f}  (-{reduction:.1f}%)")
+    return dataset
+
+
+# ── Main ───────────────────────────────────────────────────────────────────────
+
+def main():
+    parser = argparse.ArgumentParser(description='Process Spider v5')
+    parser.add_argument('--data_dir',                  default=r'C:\Users\ADMIN\Documents\Text  to SQL\spider_data')
+    parser.add_argument('--output_dir',                default=r'C:\Users\ADMIN\Documents\Text  to SQL\processed\spider')
+    parser.add_argument('--no_foreign_keys',           action='store_true', default=False)
+    parser.add_argument('--no_column_types',           action='store_true', default=False)
+    parser.add_argument('--no_primary_keys',           action='store_true', default=False)
+    parser.add_argument('--no_lowercase_identifiers',  action='store_true', default=False)
+    parser.add_argument('--oracle_schema_link',        action='store_true', default=False,
+                        help='RQ6 ablation: filter schema to keep only tables in gold SQL '
+                             '(ceiling-effect study for WRONG_TABLE error). Not applied to '
+                             'test split (no gold SQL available).')
+    args = parser.parse_args()
+
+    data_dir   = Path(args.data_dir)
+    output_dir = Path(args.output_dir)
+    output_dir.mkdir(parents=True, exist_ok=True)
+
+    use_fk                = not args.no_foreign_keys
+    use_column_types      = not args.no_column_types
+    use_primary_keys      = not args.no_primary_keys
+    lowercase_identifiers = not args.no_lowercase_identifiers
+    oracle_schema_link    = args.oracle_schema_link
+
+    schemas      = load_schemas(str(data_dir / 'tables.json'))
+    test_schemas = (load_schemas(str(data_dir / 'test_tables.json'))
+                    if (data_dir / 'test_tables.json').exists() else schemas)
+
+    splits = [
+        ('train', data_dir / 'train_spider.json', schemas),
+        ('dev',   data_dir / 'dev.json',          schemas),
+        ('test',  data_dir / 'test.json',          test_schemas),
+    ]
+    for split_name, fpath, split_schemas in splits:
+        if not fpath.exists():
+            print(f"WARNING: {fpath} not found, skipping")
+            continue
+        dataset = process_split(str(fpath), split_schemas, split_name,
+                                use_fk, use_column_types,
+                                use_primary_keys, lowercase_identifiers,
+                                oracle_schema_link)
+        out = output_dir / f"{split_name}.json"
+        with open(out, 'w', encoding='utf-8') as f:
+            json.dump(dataset, f, ensure_ascii=False, indent=2)
+        size = out.stat().st_size / 1e6
+        print(f"  Saved -> {out}  ({len(dataset):,} samples, {size:.1f} MB)")
+
+
+if __name__ == '__main__':
+    main()