ritvik360/nl2sql-bench

1import os2import sys3import json4import torch5import hashlib6from pathlib import Path7from tqdm import tqdm8from transformers import AutoModelForCausalLM, AutoTokenizer9 10# GPU CONFIG - All 4 H100s engaged11os.environ["CUDA_VISIBLE_DEVICES"] = "0,1,2,3,4,7"12 13PROJECT_ROOT = os.path.abspath(os.path.dirname(__file__))14if PROJECT_ROOT not in sys.path:15    sys.path.insert(0, PROJECT_ROOT)16 17from data_factory.schemas import SCHEMA_CONTEXT18from data_factory.validator import SQLValidator19 20# CONFIG21MODEL_NAME = "Qwen/Qwen2.5-72B-Instruct"22TARGET_TEMPLATES = 1000023OUTPUT_FILE = "llm_10k_base_templates.json"24BATCH_SIZE = 6425 26PROMPT_TEMPLATE = """27You are a senior expert in SQLite schema design and NL2SQL dataset generation.28 29TASK30Generate exactly 10 UNIQUE, COMPLEX, and FULLY VALID SQLite SQL SELECT queries for the given schema.31For each query, also write a natural language question that a real user might ask.32 33HARD RULES34- Output ONLY a valid JSON array.35- Do NOT wrap output in markdown, code fences, or explanations.36- Every item must be a JSON object with exactly these keys:37  - "sql"38  - "base_nl"39  - "difficulty"40  - "has_order"41- All SQL must be a single SELECT statement.42- Do NOT use INSERT, UPDATE, DELETE, DROP, CREATE, ALTER, PRAGMA, ATTACH, DETACH, or any DDL/DML.43- Every table and column used in SQL must exist in the provided schema.44- Do NOT invent columns, tables, aliases, or constraints.45- SQL must be valid for SQLite.46- Prefer queries that are meaningfully different from each other.47- Avoid repetitive templates.48- Each SQL should test a different reasoning pattern.49- Each base_nl should sound natural and distinct from the others.50- Use advanced SQL patterns where appropriate:51  - multiple JOINs52  - CTEs53  - subqueries54  - window functions such as ROW_NUMBER, RANK, DENSE_RANK, LAG, LEAD55  - GROUP BY and HAVING56  - conditional aggregation57  - anti-joins / exclusion logic58  - top-N per group59  - time-based filtering60- Exactly 3 of the 10 queries must be "easy" (basic filtering, simple lookups, 1-2 tables).61- Exactly 3 of the 10 queries must be "medium" (moderate complexity, standard JOINs, basic aggregation).62- Exactly 4 of the 10 queries must be genuinely "hard" (advanced patterns, CTEs, subqueries, window functions).63- Ensure the "difficulty" key strictly contains one of these exact string values: "easy", "medium", or "hard".64 65QUALITY TARGETS66- The SQL should be executable as written.67- The question should be answerable from the schema alone.68- Prefer business-like, realistic analytics questions.69- Prefer queries that require combining 2 to 4 tables.70- If a query uses aggregation, ensure the NL clearly implies aggregation.71- If a query uses ordering, include "has_order": true.72- If a query does not require ordering, set "has_order": false.73- Make the 10 queries cover diverse intent types:74  1. ranking75  2. comparison against average or median76  3. top/bottom-N77  4. grouped aggregation78  5. time filtering79  6. multi-join analysis80  7. exclusion / NOT EXISTS81  8. window-function based analysis82  9. conditional counting83  10. trend or interval-based logic84 85SCHEMA86{schema}87 88OUTPUT FORMAT89Return ONLY a valid JSON array of 10 objects.90 91Example structure:92[93  {{94    "sql": "SELECT ...",95    "base_nl": "Show ...",96    "difficulty": "hard",97    "has_order": true98  }}99]100 101FINAL SELF-CHECK BEFORE RESPONDING102- Confirm the output is valid JSON.103- Confirm there are exactly 10 objects.104- Confirm every SQL is a single SELECT.105- Confirm no hallucinated schema elements exist.106- Confirm the 10 questions are not paraphrases of each other.107"""108 109def extract_json(raw_text):110    text = raw_text.strip()111    if text.startswith("```json"):112        text = text[7:-3].strip()113    elif text.startswith("```"):114        text = text[3:-3].strip()115    start = text.find("[")116    end = text.rfind("]")117    if start != -1 and end != -1:118        return text[start:end+1]119    return None120 121def main():122    print("Loading Model Qwen-72B (SDPA) for 10K Mining...")123 124    tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)125    custom_max_memory = {126        0: "60GiB",  # System GPU 0 (Has 13GB used, ~67GB free)127        1: "75GiB",  # System GPU 1 (Fully free)128        2: "75GiB",  # System GPU 2 (Fully free)129        3: "75GiB",  # System GPU 3 (Fully free)130        4: "75GiB",  # System GPU 4 (Fully free)131        5: "45GiB"   # System GPU 7 (Has 25GB used, ~55GB free)132    }133    model = AutoModelForCausalLM.from_pretrained(134        MODEL_NAME,135        device_map="auto",136        max_memory = custom_max_memory,137        torch_dtype=torch.bfloat16,138        attn_implementation="sdpa"139    )140 141    domains = list(SCHEMA_CONTEXT.keys())142    valid_templates = []143    seen_sql_hashes = set()144 145    # Resume support: Load existing templates to prevent duplicates146    if os.path.exists(OUTPUT_FILE):147        with open(OUTPUT_FILE, "r") as f:148            valid_templates = json.load(f)149            for t in valid_templates:150                seen_sql_hashes.add(hashlib.md5(t["sql"].lower().encode()).hexdigest())151 152    pbar = tqdm(total=TARGET_TEMPLATES, initial=len(valid_templates), desc="Mining 10K Base Templates")153 154    validators = {}155    domain_idx = 0156 157    while len(valid_templates) < TARGET_TEMPLATES:158        batch_prompts = []159        batch_domains = []160 161        # Prepare Batch162        for _ in range(BATCH_SIZE):163            domain = domains[domain_idx % len(domains)]164            schema_string = SCHEMA_CONTEXT[domain]165            domain_idx += 1166            167            messages = [168                {"role": "system", "content": "You output only valid JSON arrays. Do not include markdown."},169                {"role": "user", "content": PROMPT_TEMPLATE.format(schema=schema_string)}170            ]171            chat_text = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)172            batch_prompts.append(chat_text)173            batch_domains.append(domain)174 175        inputs = tokenizer(batch_prompts, return_tensors="pt", padding=True, truncation=True).to(model.device)176 177        try:178            tqdm.write(f"\n[DEBUG] Sending batch of {BATCH_SIZE} to model.generate(). Please wait...")179            with torch.no_grad():180                outputs = model.generate(181                    **inputs,182                    max_new_tokens=5000, 183                    do_sample=True,184                    temperature=0.55,185                    top_p=0.9,186                    pad_token_id=tokenizer.eos_token_id187                )188            tqdm.write("[DEBUG] Model generation finished. Decoding responses...")189 190            # Output Slicing191            input_length = inputs.input_ids.shape[1]192            generated_tokens = outputs[:, input_length:]193            responses = tokenizer.batch_decode(generated_tokens, skip_special_tokens=True)194 195            batch_added = 0196            for i, (response, domain) in enumerate(zip(responses, batch_domains)):197                tqdm.write(f"\n[DEBUG] Processing Response {i+1}/{BATCH_SIZE} for domain: {domain}")198                199                json_text = extract_json(response)200                if not json_text:201                    tqdm.write(f"[DEBUG] extract_json failed. Raw text snippet: {response[:200]}...")202                    continue203 204                try:205                    generated_data = json.loads(json_text)206                    tqdm.write(f"[DEBUG] JSON loaded successfully. Found {len(generated_data)} items.")207                except Exception as e:208                    tqdm.write(f"[DEBUG] json.loads failed. Error: {e}")209                    tqdm.write(f"[DEBUG] Bad JSON snippet: {json_text[:200]}...")210                    continue211 212                if domain not in validators:213                    validators[domain] = SQLValidator(domain, seed=42)214                validator = validators[domain]215 216                for item in generated_data:217                    if not isinstance(item, dict): continue218                    219                    sql = item.get("sql", "").strip()220                    if not sql: continue221 222                    # Check for duplicates using hash223                    sql_hash = hashlib.md5(sql.lower().encode()).hexdigest()224                    if sql_hash in seen_sql_hashes:225                        tqdm.write("[DEBUG] Duplicate query skipped.")226                        continue227 228                    val_result = validator.validate(sql)229 230                    # Hard validation rule: SQL must execute AND return rows231                    if val_result.passed and val_result.row_count > 0:232                        tqdm.write(f"[DEBUG] SQL Passed (Rows: {val_result.row_count}): {sql[:50]}...")233                        item["domain"] = domain234                        item["id"] = f"base_{len(valid_templates)}"235                        valid_templates.append(item)236                        seen_sql_hashes.add(sql_hash)237                        batch_added += 1238                    else:239                        tqdm.write(f"[DEBUG] SQL Failed Validation or 0 Rows (Passed: {val_result.passed}, Rows: {val_result.row_count}): {sql[:50]}...")240 241            if batch_added > 0:242                pbar.update(batch_added)243                tqdm.write(f"[DEBUG] Auto-saving {batch_added} new templates to JSON...")244                # Auto-save after every successful batch245                with open(OUTPUT_FILE, "w") as f:246                    json.dump(valid_templates, f, indent=2)247 248            if len(valid_templates) >= TARGET_TEMPLATES:249                break250 251        except Exception as e:252            tqdm.write(f"\n[DEBUG] CRITICAL EXCEPTION CAUGHT: {e}")253            continue254 255    # Close validators256    for v in validators.values():257        v.close()258 259    pbar.close()260    print(f"\nBoom! Generated {len(valid_templates)} Elite Base Templates!")261 262if __name__ == "__main__":263    main()