ritvik360/nl2sql-bench

1"""2generate_edge_cases.py3======================4Targeted edge-case data generator for the 4 failure patterns found in eval:5  1. ROW_NUMBER vs RANK vs DENSE_RANK  (tie-breaking semantics)6  2. strftime month as INTEGER (not '%Y-%m' string)7  3. SELECT column discipline  (no unrequested extras)8  4. LAG/LEAD period-over-period9  5. HAVING vs WHERE placement10  6. COUNT(DISTINCT) vs COUNT11 12Produces: edge_cases.jsonl  (same chat format as nl2sql_cleaned_ready_to_train.jsonl)13Run:  python generate_edge_cases.py14"""15 16import os, sys, json, re, hashlib17from tqdm import tqdm18from transformers import AutoModelForCausalLM, AutoTokenizer, AwqConfig, BitsAndBytesConfig19import torch20import transformers.activations21# Yeh line AutoAWQ ko bewakoof banayegi taaki wo crash na ho22if not hasattr(transformers.activations, 'PytorchGELUTanh'):23    transformers.activations.PytorchGELUTanh = transformers.activations.NewGELUActivation24os.environ["CUDA_VISIBLE_DEVICES"] = "3,1,6,7"25 26PROJECT_ROOT = os.path.abspath(os.path.dirname(__file__))27if PROJECT_ROOT not in sys.path:28    sys.path.insert(0, PROJECT_ROOT)29 30from data_factory.schemas import SCHEMA_CONTEXT31 32quantization_config = BitsAndBytesConfig(33    load_in_4bit=True,34    bnb_4bit_compute_dtype=torch.bfloat16, 35    bnb_4bit_use_double_quant=True,36    bnb_4bit_quant_type="nf4"37)38 39MODEL_NAME  = "Qwen/Qwen2.5-72B-Instruct"40OUTPUT_FILE = "edge_cases.jsonl"41BATCH_SIZE  = 8          # smaller — edge prompts are long42SAMPLES_PER_PATTERN = 715 # ~6 batches per pattern → 5005 total edge samples43 44SYSTEM_PROMPT = (45    "You are a Senior SQL Architect. "46    "Output ONLY the SQL query. Use SQLite syntax."47)48 49# ── Edge-case prompt templates ──────────────────────────────────────────────50# Each entry: (pattern_tag, user_prompt_template)51# {schema} is filled at runtime with a random domain schema.52 53EDGE_PATTERNS = [54 55    # 1. ROW_NUMBER tie-breaking — the #1 failure56    ("row_number_tiebreak", """SCHEMA:57{schema}58 59Generate exactly 8 NL2SQL pairs that REQUIRE ROW_NUMBER() (not RANK or DENSE_RANK) \60because the question explicitly says "pick one winner when there is a tie" \61using a tiebreaker column (e.g. lower id, earlier date).62 63Output ONLY a valid JSON array:64[65  {{"nl": "...", "sql": "SELECT ..."}},66  ...67]68 69Rules:70- Every SQL must use ROW_NUMBER() OVER (...) not RANK().71- The OVER clause ORDER BY must include the tiebreaker column.72- WHERE rn = 1 must appear in an outer query or CTE.73- No markdown. No explanation. Just the JSON array."""),74 75    # 2. RANK / DENSE_RANK — when ties SHOULD persist76    ("rank_dense_rank", """SCHEMA:77{schema}78 79Generate exactly 8 NL2SQL pairs where RANK() or DENSE_RANK() is the CORRECT choice \80because the question says "show all tied records at the same rank".81 82Output ONLY a valid JSON array:83[84  {{"nl": "...", "sql": "SELECT ..."}},85  ...86]87 88Rules:89- Use RANK() when question implies gaps after ties, DENSE_RANK() when no gaps.90- NL must make the tie-semantics explicit ("same rank", "tied positions").91- No markdown. No explanation. Just the JSON array."""),92 93    # 3. strftime integer month output94    ("strftime_integer_month", """SCHEMA:95{schema}96 97Generate exactly 8 NL2SQL pairs where the question asks for a numeric month number \98(1–12), NOT a 'YYYY-MM' string.99 100Output ONLY a valid JSON array:101[102  {{"nl": "...", "sql": "SELECT ..."}},103  ...104]105 106Rules:107- SQL must use CAST(strftime('%m', <col>) AS INTEGER) to produce integer month.108- Do NOT use strftime('%Y-%m', ...) when the question asks for month number.109- NL questions must say "month number", "which month (1–12)", or similar.110- No markdown. No explanation. Just the JSON array."""),111 112    # 4. SELECT column discipline113    ("select_column_discipline", """SCHEMA:114{schema}115 116Generate exactly 8 NL2SQL pairs where the question explicitly names ONLY the columns \117to return. The SQL must select EXACTLY those columns — no extras like avg_salary, \118row counts, or intermediate aggregates.119 120Output ONLY a valid JSON array:121[122  {{"nl": "...", "sql": "SELECT ..."}},123  ...124]125 126Rules:127- NL must say "return only X, Y, Z" or "show me only the name and total".128- SQL SELECT list must contain only those columns.129- If aggregation is needed internally (e.g. for HAVING), do NOT expose it in SELECT.130- No markdown. No explanation. Just the JSON array."""),131 132    # 5. LAG / LEAD period-over-period133    ("lag_lead_period", """SCHEMA:134{schema}135 136Generate exactly 8 NL2SQL pairs that require LAG() or LEAD() window functions \137for period-over-period comparison (e.g. month-over-month revenue change, \138previous order amount, next appointment date).139 140Output ONLY a valid JSON array:141[142  {{"nl": "...", "sql": "SELECT ..."}},143  ...144]145 146Rules:147- Use LAG(<col>, 1) OVER (ORDER BY ...) or LEAD(...) correctly.148- NL must imply comparison with previous or next row/period.149- No markdown. No explanation. Just the JSON array."""),150 151    # 6. HAVING vs WHERE152    ("having_vs_where", """SCHEMA:153{schema}154 155Generate exactly 8 NL2SQL pairs that test correct placement of filter conditions:156- Conditions on raw columns → WHERE157- Conditions on aggregates → HAVING158Include 4 pairs where a wrong model might put an aggregate condition in WHERE (trap).159 160Output ONLY a valid JSON array:161[162  {{"nl": "...", "sql": "SELECT ..."}},163  ...164]165 166Rules:167- SQL must never filter an aggregate (COUNT, SUM, AVG) inside WHERE.168- SQL must never put a raw column filter inside HAVING.169- No markdown. No explanation. Just the JSON array."""),170 171    # 7. COUNT(DISTINCT) vs COUNT172    ("count_distinct", """SCHEMA:173{schema}174 175Generate exactly 8 NL2SQL pairs where the question specifically asks for \176"unique", "distinct", or "different" counts — requiring COUNT(DISTINCT col).177Also include 2 pairs where COUNT(*) is correct to reinforce the contrast.178 179Output ONLY a valid JSON array:180[181  {{"nl": "...", "sql": "SELECT ..."}},182  ...183]184 185Rules:186- When NL says "unique/distinct", SQL must use COUNT(DISTINCT <col>).187- When NL says "total orders placed" (not distinct), use COUNT(*) or COUNT(id).188- No markdown. No explanation. Just the JSON array."""),189]190 191 192# ── Helpers ──────────────────────────────────────────────────────────────────193 194def extract_json_array(text: str) -> str:195    text = text.strip()196    # strip code fences if model leaks them197    text = re.sub(r"```(?:json)?\n?(.*?)```", r"\1", text, flags=re.DOTALL).strip()198    s, e = text.find("["), text.rfind("]")199    return text[s:e+1] if s != -1 and e != -1 else "[]"200 201def get_hash(text: str) -> str:202    return hashlib.md5(text.lower().strip().encode()).hexdigest()203 204def build_record(nl: str, sql: str, domain: str) -> dict:205    return {206        "prompt": [207            {"role": "system", "content": SYSTEM_PROMPT},208            {"role": "user",   "content": f"SCHEMA: {SCHEMA_CONTEXT[domain]}\nQUESTION: {nl}"}209        ],210        "sql": sql211    }212 213 214# ── Main ─────────────────────────────────────────────────────────────────────215 216def main():217    print(f"Loading {MODEL_NAME}...")218    tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME, padding_side="left")219    tokenizer.pad_token = tokenizer.eos_token220    custom_memory = {0:"30GiB",1:"75GiB",2:"45GiB",3:"45GiB"}221    model = AutoModelForCausalLM.from_pretrained(222        MODEL_NAME,223        device_map="auto",              224        max_memory=custom_memory,       225        quantization_config=quantization_config, 226        torch_dtype=torch.bfloat16,227        low_cpu_mem_usage=True,228        attn_implementation = "sdpa"229    )230 231    domains = list(SCHEMA_CONTEXT.keys())232    seen    = set()233    total   = 0234 235    out = open(OUTPUT_FILE, "a", encoding="utf-8")236 237    for pattern_tag, prompt_tmpl in EDGE_PATTERNS:238        print(f"\n[PATTERN] {pattern_tag}")239        collected = 0240        domain_idx = 0241        pbar = tqdm(total=SAMPLES_PER_PATTERN, desc=pattern_tag)242 243        while collected < SAMPLES_PER_PATTERN:244            # Build a batch of prompts, cycling through domains245            batch_domains = []246            batch_prompts = []247            for _ in range(BATCH_SIZE):248                domain = domains[domain_idx % len(domains)]249                domain_idx += 1250                user_msg = prompt_tmpl.format(schema=SCHEMA_CONTEXT[domain])251                msgs = [252                    {"role": "system", "content": "You output only valid JSON arrays. No markdown."},253                    {"role": "user",   "content": user_msg}254                ]255                batch_prompts.append(256                    tokenizer.apply_chat_template(msgs, tokenize=False, add_generation_prompt=True)257                )258                batch_domains.append(domain)259 260            inputs = tokenizer(261                batch_prompts, return_tensors="pt", padding=True, truncation=True262            ).to(model.device)263 264            try:265                with torch.no_grad():266                    outputs = model.generate(267                        **inputs,268                        max_new_tokens=2048,269                        do_sample=True,270                        temperature=0.5,271                        top_p=0.9,272                        pad_token_id=tokenizer.eos_token_id273                    )274                responses = tokenizer.batch_decode(275                    outputs[:, inputs.input_ids.shape[1]:], skip_special_tokens=True276                )277 278                for resp, domain in zip(responses, batch_domains):279                    raw = extract_json_array(resp)280                    try:281                        pairs = json.loads(raw)282                    except Exception:283                        continue284 285                    for pair in pairs:286                        nl  = pair.get("nl",  "").strip()287                        sql = pair.get("sql", "").strip()288                        if not nl or not sql:289                            continue290                        # strip fences in sql just in case291                        sql = re.sub(r"```(?:sql)?\n?(.*?)```", r"\1", sql, flags=re.DOTALL).strip()292 293                        h = get_hash(nl + sql)294                        if h in seen:295                            continue296                        seen.add(h)297 298                        record = build_record(nl, sql, domain)299                        out.write(json.dumps(record, ensure_ascii=False) + "\n")300                        out.flush()301                        collected += 1302                        total     += 1303                        pbar.update(1)304 305                        if collected >= SAMPLES_PER_PATTERN:306                            break307 308            except Exception as e:309                tqdm.write(f"[WARN] Batch failed: {e}")310                continue311 312        pbar.close()313 314    out.close()315    print(f"\nDone! {total} edge-case records saved to {OUTPUT_FILE}")316 317 318if __name__ == "__main__":319    main()