generate_data.py

import os
import sys
import json
import torch
import hashlib
from pathlib import Path
from tqdm import tqdm
from transformers import AutoModelForCausalLM, AutoTokenizer

# GPU CONFIG - All 4 H100s engaged
os.environ["CUDA_VISIBLE_DEVICES"] = "0,1,2,3,4,7"

PROJECT_ROOT = os.path.abspath(os.path.dirname(__file__))
if PROJECT_ROOT not in sys.path:
    sys.path.insert(0, PROJECT_ROOT)

from data_factory.schemas import SCHEMA_CONTEXT
from data_factory.validator import SQLValidator

# CONFIG
MODEL_NAME = "Qwen/Qwen2.5-72B-Instruct"
TARGET_TEMPLATES = 10000
OUTPUT_FILE = "llm_10k_base_templates.json"
BATCH_SIZE = 64

PROMPT_TEMPLATE = """
You are a senior expert in SQLite schema design and NL2SQL dataset generation.

TASK
Generate exactly 10 UNIQUE, COMPLEX, and FULLY VALID SQLite SQL SELECT queries for the given schema.
For each query, also write a natural language question that a real user might ask.

HARD RULES
- Output ONLY a valid JSON array.
- Do NOT wrap output in markdown, code fences, or explanations.
- Every item must be a JSON object with exactly these keys:
  - "sql"
  - "base_nl"
  - "difficulty"
  - "has_order"
- All SQL must be a single SELECT statement.
- Do NOT use INSERT, UPDATE, DELETE, DROP, CREATE, ALTER, PRAGMA, ATTACH, DETACH, or any DDL/DML.
- Every table and column used in SQL must exist in the provided schema.
- Do NOT invent columns, tables, aliases, or constraints.
- SQL must be valid for SQLite.
- Prefer queries that are meaningfully different from each other.
- Avoid repetitive templates.
- Each SQL should test a different reasoning pattern.
- Each base_nl should sound natural and distinct from the others.
- Use advanced SQL patterns where appropriate:
  - multiple JOINs
  - CTEs
  - subqueries
  - window functions such as ROW_NUMBER, RANK, DENSE_RANK, LAG, LEAD
  - GROUP BY and HAVING
  - conditional aggregation
  - anti-joins / exclusion logic
  - top-N per group
  - time-based filtering
- Exactly 3 of the 10 queries must be "easy" (basic filtering, simple lookups, 1-2 tables).
- Exactly 3 of the 10 queries must be "medium" (moderate complexity, standard JOINs, basic aggregation).
- Exactly 4 of the 10 queries must be genuinely "hard" (advanced patterns, CTEs, subqueries, window functions).
- Ensure the "difficulty" key strictly contains one of these exact string values: "easy", "medium", or "hard".

QUALITY TARGETS
- The SQL should be executable as written.
- The question should be answerable from the schema alone.
- Prefer business-like, realistic analytics questions.
- Prefer queries that require combining 2 to 4 tables.
- If a query uses aggregation, ensure the NL clearly implies aggregation.
- If a query uses ordering, include "has_order": true.
- If a query does not require ordering, set "has_order": false.
- Make the 10 queries cover diverse intent types:
  1. ranking
  2. comparison against average or median
  3. top/bottom-N
  4. grouped aggregation
  5. time filtering
  6. multi-join analysis
  7. exclusion / NOT EXISTS
  8. window-function based analysis
  9. conditional counting
  10. trend or interval-based logic

SCHEMA
{schema}

OUTPUT FORMAT
Return ONLY a valid JSON array of 10 objects.

Example structure:
[
  {{
    "sql": "SELECT ...",
    "base_nl": "Show ...",
    "difficulty": "hard",
    "has_order": true
  }}
]

FINAL SELF-CHECK BEFORE RESPONDING
- Confirm the output is valid JSON.
- Confirm there are exactly 10 objects.
- Confirm every SQL is a single SELECT.
- Confirm no hallucinated schema elements exist.
- Confirm the 10 questions are not paraphrases of each other.
"""

def extract_json(raw_text):
    text = raw_text.strip()
    if text.startswith("```json"):
        text = text[7:-3].strip()
    elif text.startswith("```"):
        text = text[3:-3].strip()
    start = text.find("[")
    end = text.rfind("]")
    if start != -1 and end != -1:
        return text[start:end+1]
    return None

def main():
    print("Loading Model Qwen-72B (SDPA) for 10K Mining...")

    tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
    custom_max_memory = {
        0: "60GiB",  # System GPU 0 (Has 13GB used, ~67GB free)
        1: "75GiB",  # System GPU 1 (Fully free)
        2: "75GiB",  # System GPU 2 (Fully free)
        3: "75GiB",  # System GPU 3 (Fully free)
        4: "75GiB",  # System GPU 4 (Fully free)
        5: "45GiB"   # System GPU 7 (Has 25GB used, ~55GB free)
    }
    model = AutoModelForCausalLM.from_pretrained(
        MODEL_NAME,
        device_map="auto",
        max_memory = custom_max_memory,
        torch_dtype=torch.bfloat16,
        attn_implementation="sdpa"
    )

    domains = list(SCHEMA_CONTEXT.keys())
    valid_templates = []
    seen_sql_hashes = set()

    # Resume support: Load existing templates to prevent duplicates
    if os.path.exists(OUTPUT_FILE):
        with open(OUTPUT_FILE, "r") as f:
            valid_templates = json.load(f)
            for t in valid_templates:
                seen_sql_hashes.add(hashlib.md5(t["sql"].lower().encode()).hexdigest())

    pbar = tqdm(total=TARGET_TEMPLATES, initial=len(valid_templates), desc="Mining 10K Base Templates")

    validators = {}
    domain_idx = 0

    while len(valid_templates) < TARGET_TEMPLATES:
        batch_prompts = []
        batch_domains = []

        # Prepare Batch
        for _ in range(BATCH_SIZE):
            domain = domains[domain_idx % len(domains)]
            schema_string = SCHEMA_CONTEXT[domain]
            domain_idx += 1
            
            messages = [
                {"role": "system", "content": "You output only valid JSON arrays. Do not include markdown."},
                {"role": "user", "content": PROMPT_TEMPLATE.format(schema=schema_string)}
            ]
            chat_text = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
            batch_prompts.append(chat_text)
            batch_domains.append(domain)

        inputs = tokenizer(batch_prompts, return_tensors="pt", padding=True, truncation=True).to(model.device)

        try:
            tqdm.write(f"\n[DEBUG] Sending batch of {BATCH_SIZE} to model.generate(). Please wait...")
            with torch.no_grad():
                outputs = model.generate(
                    **inputs,
                    max_new_tokens=5000, 
                    do_sample=True,
                    temperature=0.55,
                    top_p=0.9,
                    pad_token_id=tokenizer.eos_token_id
                )
            tqdm.write("[DEBUG] Model generation finished. Decoding responses...")

            # Output Slicing
            input_length = inputs.input_ids.shape[1]
            generated_tokens = outputs[:, input_length:]
            responses = tokenizer.batch_decode(generated_tokens, skip_special_tokens=True)

            batch_added = 0
            for i, (response, domain) in enumerate(zip(responses, batch_domains)):
                tqdm.write(f"\n[DEBUG] Processing Response {i+1}/{BATCH_SIZE} for domain: {domain}")
                
                json_text = extract_json(response)
                if not json_text:
                    tqdm.write(f"[DEBUG] extract_json failed. Raw text snippet: {response[:200]}...")
                    continue

                try:
                    generated_data = json.loads(json_text)
                    tqdm.write(f"[DEBUG] JSON loaded successfully. Found {len(generated_data)} items.")
                except Exception as e:
                    tqdm.write(f"[DEBUG] json.loads failed. Error: {e}")
                    tqdm.write(f"[DEBUG] Bad JSON snippet: {json_text[:200]}...")
                    continue

                if domain not in validators:
                    validators[domain] = SQLValidator(domain, seed=42)
                validator = validators[domain]

                for item in generated_data:
                    if not isinstance(item, dict): continue
                    
                    sql = item.get("sql", "").strip()
                    if not sql: continue

                    # Check for duplicates using hash
                    sql_hash = hashlib.md5(sql.lower().encode()).hexdigest()
                    if sql_hash in seen_sql_hashes:
                        tqdm.write("[DEBUG] Duplicate query skipped.")
                        continue

                    val_result = validator.validate(sql)

                    # Hard validation rule: SQL must execute AND return rows
                    if val_result.passed and val_result.row_count > 0:
                        tqdm.write(f"[DEBUG] SQL Passed (Rows: {val_result.row_count}): {sql[:50]}...")
                        item["domain"] = domain
                        item["id"] = f"base_{len(valid_templates)}"
                        valid_templates.append(item)
                        seen_sql_hashes.add(sql_hash)
                        batch_added += 1
                    else:
                        tqdm.write(f"[DEBUG] SQL Failed Validation or 0 Rows (Passed: {val_result.passed}, Rows: {val_result.row_count}): {sql[:50]}...")

            if batch_added > 0:
                pbar.update(batch_added)
                tqdm.write(f"[DEBUG] Auto-saving {batch_added} new templates to JSON...")
                # Auto-save after every successful batch
                with open(OUTPUT_FILE, "w") as f:
                    json.dump(valid_templates, f, indent=2)

            if len(valid_templates) >= TARGET_TEMPLATES:
                break

        except Exception as e:
            tqdm.write(f"\n[DEBUG] CRITICAL EXCEPTION CAUGHT: {e}")
            continue

    # Close validators
    for v in validators.values():
        v.close()

    pbar.close()
    print(f"\nBoom! Generated {len(valid_templates)} Elite Base Templates!")

if __name__ == "__main__":
    main()