""" WikiSQL Data Processing Script v5 Build schema-aware inputs for WikiSQL to fine-tune T5/CodeT5/BART. Input: WikiSQL/{train,dev,test}.{jsonl,tables.jsonl} Output: processed/wikisql/{train,dev,test}.json Usage: python src/process_wikisql.py python src/process_wikisql.py --no_column_types """ 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 HAS_SQLGLOT = True except ImportError: HAS_SQLGLOT = False print("WARNING: sqlglot not installed. Run: pip install sqlglot") AGG_OPS = ['', 'MAX', 'MIN', 'COUNT', 'SUM', 'AVG'] COND_OPS = ['=', '>', '<', '!=', 'LIKE'] TABLE_NAME = "data_table" WIKISQL_TYPE_MAP = {'text': 'text', 'real': 'number'} SQL_KEYWORDS = { 'select', 'from', 'where', 'and', 'or', 'not', 'in', 'is', 'null', 'like', 'order', 'by', 'group', 'having', 'limit', 'join', 'left', 'right', 'inner', 'outer', 'on', 'as', 'distinct', 'count', 'sum', 'avg', 'max', 'min', 'between', 'case', 'when', 'then', 'else', 'end', 'union', 'all', 'insert', 'update', 'delete', 'drop', 'create', 'table', 'index', 'view', 'exists', 'with', 'into', 'values', 'set', 'alter', 'add', 'column', 'primary', 'key', 'foreign', 'references', 'check', 'default', 'unique', 'cast', 'true', 'false', } SPECIAL_CHARS = set(" /-.(),%#'\"&+*?:;!@$\\<>=[]{}|~`^") # ── SQL utilities ────────────────────────────────────────────────────────────── def normalize_sql(sql): """Normalize SQL for Exact Match — lowercase everything (WikiSQL DB is already lowercase).""" if not sql: return "" if HAS_SQLGLOT: try: return sqlglot.parse_one(sql, dialect='sqlite').sql( dialect='sqlite', pretty=False ).lower().strip() except Exception: pass sql = sql.strip().lower() 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 check_sql_validity(sql): if not sql or not sql.strip(): return False if HAS_SQLGLOT: try: sqlglot.parse_one(sql, dialect='sqlite') return True except Exception: return False return sql.strip().upper().startswith('SELECT') # ── Quote helpers ────────────────────────────────────────────────────────────── def quote_column(col_name): """Double-quote column name if it contains special characters or is an SQL keyword.""" if not col_name or not col_name.strip(): return '"unknown_column"' needs_quote = ( any(c in SPECIAL_CHARS for c in col_name) or col_name[0].isdigit() or col_name.lower() in SQL_KEYWORDS ) return f'"{col_name}"' if needs_quote else col_name def quote_value(val): """Format WHERE value — lowercase string, numeric as-is.""" if isinstance(val, str): return f"'{val.lower().replace(chr(39), chr(39)*2)}'" return str(val) # ── Core conversion ──────────────────────────────────────────────────────────── def convert_sql_to_string(sql_json, header): """Reconstruct SQL string from WikiSQL structured annotation {sel, agg, conds}.""" sel_col = quote_column(header[sql_json['sel']]) agg = AGG_OPS[sql_json['agg']] select = f"SELECT {agg}({sel_col})" if agg else f"SELECT {sel_col}" conds = [ f"{quote_column(header[c[0]])} {COND_OPS[c[1]]} {quote_value(c[2])}" for c in sql_json['conds'] ] where = " WHERE " + " AND ".join(conds) if conds else "" return (select + f" FROM {TABLE_NAME}" + where).lower() def build_schema_string(header, types=None, use_column_types=True): col_parts = [] for i, col_name in enumerate(header): col = quote_column(col_name) if use_column_types and types and i < len(types): col += f":{WIKISQL_TYPE_MAP.get(types[i], types[i])}" col_parts.append(col) return f"{TABLE_NAME} ( {' , '.join(col_parts)} )" def build_model_input(question, header, types=None, use_column_types=True): schema = build_schema_string(header, types=types, use_column_types=use_column_types) return f"question: {question.strip()} | schema: {schema}" # ── Data loading ─────────────────────────────────────────────────────────────── def load_tables(tables_path): tables = {} with open(tables_path, 'r', encoding='utf-8') as f: for line in f: if line.strip(): t = json.loads(line) tables[t['id']] = {'header': t['header'], 'types': t.get('types', [])} print(f" Loaded {len(tables):,} tables from {Path(tables_path).name}") return tables # ── Process split ────────────────────────────────────────────────────────────── def process_split(data_path, tables_path, split_name, use_column_types=True): print(f"\n[{split_name.upper()}] Processing...") tables = load_tables(tables_path) dataset, errors, not_found, invalid_sql = [], 0, 0, 0 with open(data_path, 'r', encoding='utf-8') as f: for idx, line in enumerate(f): if not line.strip(): continue try: item = json.loads(line) table_id = item['table_id'] if table_id not in tables or not tables[table_id]['header']: not_found += 1 continue header = tables[table_id]['header'] types = tables[table_id].get('types', []) sql_str = convert_sql_to_string(item['sql'], header) valid = check_sql_validity(sql_str) if not valid: invalid_sql += 1 dataset.append({ 'input': build_model_input(item['question'], header, types=types, use_column_types=use_column_types), 'output': sql_str, 'normalized_output': normalize_sql(sql_str), 'is_valid_sql': valid, 'question': item['question'].strip(), 'table_id': table_id, 'phase': item.get('phase', 1), }) except Exception as e: errors += 1 if errors <= 5: print(f" WARNING line {idx}: {e}") print(f" Processed: {len(dataset):,} | Not found: {not_found} | " f"Errors: {errors} | Invalid SQL: {invalid_sql}") return dataset # ── Main ─────────────────────────────────────────────────────────────────────── def main(): parser = argparse.ArgumentParser(description='Process WikiSQL v5') parser.add_argument('--data_dir', default=r'C:\Users\ADMIN\Documents\Text to SQL\WikiSQL') parser.add_argument('--output_dir', default=r'C:\Users\ADMIN\Documents\Text to SQL\processed\wikisql') parser.add_argument('--no_column_types', action='store_true', default=False, help='Ablation: remove :type annotation from schema') args = parser.parse_args() data_dir = Path(args.data_dir) output_dir = Path(args.output_dir) use_column_types = not args.no_column_types output_dir.mkdir(parents=True, exist_ok=True) splits = { 'train': ('train.jsonl', 'train.tables.jsonl'), 'dev': ('dev.jsonl', 'dev.tables.jsonl'), 'test': ('test.jsonl', 'test.tables.jsonl'), } for split_name, (data_file, tables_file) in splits.items(): data_path = data_dir / data_file tables_path = data_dir / tables_file if not data_path.exists() or not tables_path.exists(): print(f"WARNING: Missing files for {split_name}, skipping") continue dataset = process_split(str(data_path), str(tables_path), split_name, use_column_types=use_column_types) 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()