src/process_wikisql.py

"""
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()