custom_train.py

"""
merge_and_train.py
==================
1. Merges nl2sql_cleaned_ready_to_train.jsonl + edge_cases.jsonl
2. Shuffles the combined dataset
3. Retrains using the same GRPO setup as train.py

Run:
    python merge_and_train.py

Flags (env vars):
    EDGE_FILE   — path to edge cases jsonl   (default: edge_cases.jsonl)
    BASE_FILE   — path to existing cleaned    (default: nl2sql_cleaned_ready_to_train.jsonl)
    MERGED_FILE — merged output path          (default: nl2sql_merged_final.jsonl)
    SKIP_MERGE  — set "1" to skip merge step and go straight to training
"""

import os, sys, json, random
import torch
from datasets import Dataset
from transformers import AutoModelForCausalLM, AutoTokenizer
from peft import LoraConfig
from trl import GRPOConfig, GRPOTrainer

os.environ["CUDA_VISIBLE_DEVICES"] = "0,5,1,6"

sys.path.insert(0, "./server")
from environment import NL2SQLEnvironment
from models import NL2SQLAction
from tasks import all_task_names, get_task

# ── Config ───────────────────────────────────────────────────────────────────
BASE_FILE   = os.getenv("BASE_FILE",   "nl2sql_cleaned_ready_to_train.jsonl")
EDGE_FILE   = os.getenv("EDGE_FILE",   "edge_cases.jsonl")
MERGED_FILE = os.getenv("MERGED_FILE", "nl2sql_merged_final.jsonl")
SKIP_MERGE  = os.getenv("SKIP_MERGE",  "0") == "1"

MODEL_NAME  = "Qwen/Qwen2.5-Coder-7B-Instruct"
OUTPUT_DIR  = "./qwen-7b-coder-nl2sql-grpo-v2"

SYSTEM_PROMPT = """You are a Senior Database Architect and an expert in SQLite.
Your task is to translate natural language questions into highly optimized, correct SQLite SELECT queries.

STRICT RULES:
1. Output EXACTLY ONE valid SQLite query.
2. DO NOT wrap the query in markdown formatting (no ```sql or ```).
3. DO NOT output any explanations, conversational text, or preambles.
4. ONLY use standard SQLite functions.
5. If the question implies ordering, use the correct ORDER BY clause.
6. SELECT only the columns explicitly requested — no extras.

Your output must be executable directly against the database as-is."""


# ── Step 1: Merge ─────────────────────────────────────────────────────────────

def merge_datasets():
    if SKIP_MERGE:
        print(f"[SKIP_MERGE=1] Using existing {MERGED_FILE}")
        return

    print(f"Loading base:  {BASE_FILE}")
    print(f"Loading edges: {EDGE_FILE}")

    base_lines = []
    with open(BASE_FILE, "r", encoding="utf-8") as f:
        for line in f:
            line = line.strip()
            if line:
                base_lines.append(line)

    edge_lines = []
    with open(EDGE_FILE, "r", encoding="utf-8") as f:
        for line in f:
            line = line.strip()
            if line:
                edge_lines.append(line)

    combined = base_lines + edge_lines
    random.shuffle(combined)

    with open(MERGED_FILE, "w", encoding="utf-8") as f:
        for line in combined:
            f.write(line + "\n")

    print(
        f"Merged: {len(base_lines)} base + {len(edge_lines)} edge "
        f"= {len(combined)} total → {MERGED_FILE}"
    )


# ── Step 2: Build HF Dataset ──────────────────────────────────────────────────

def build_dataset():
    """
    Primary source: merged JSONL (base + edge cases).
    Fallback: task examples from server/tasks/ (same as original train.py).
    Both are combined so GRPO sees everything.
    """
    data = []

    # Load merged JSONL
    with open(MERGED_FILE, "r", encoding="utf-8") as f:
        for line in f:
            line = line.strip()
            if not line:
                continue
            rec = json.loads(line)
            # rec has "prompt" (list of messages) and "sql"
            # GRPO needs "prompt" and "task_name" — we use a synthetic task_name
            data.append({
                "prompt":    rec["prompt"],
                "task_name": "merged_jsonl"   # grader falls back to execution-based reward
            })

    # Also keep the original task examples so GRPO reward env works for them
    for t_name in all_task_names():
        task  = get_task(t_name)
        schema = task.schema_context()
        for ex in task.examples:
            data.append({
                "prompt": [
                    {"role": "system", "content": SYSTEM_PROMPT},
                    {"role": "user",   "content": f"SCHEMA:\n{schema}\n\nQUESTION: {ex.question}"}
                ],
                "task_name": t_name
            })

    random.shuffle(data)
    print(f"Dataset size: {len(data)} samples")
    return Dataset.from_list(data)


# ── Step 3: Reward function ───────────────────────────────────────────────────

def sql_reward_func(prompts, completions, task_name, **kwargs):
    rewards = []
    env     = NL2SQLEnvironment()

    for idx, completion in enumerate(completions):
        generated = (
            completion[0]["content"] if isinstance(completion, list) else completion
        )
        # Strip code fences defensively
        import re
        generated = re.sub(r"```(?:sql)?\n?(.*?)```", r"\1", generated, flags=re.DOTALL).strip()

        t = task_name[idx] if isinstance(task_name, list) else task_name

        # For merged_jsonl rows the env won't have a matching task →
        # reward purely on execution (non-empty result set = +1, error = 0)
        if t == "merged_jsonl":
            rewards.append(_execution_reward(generated, prompts[idx]))
            continue

        env.reset(task_name=t)
        try:
            obs = env.step(NL2SQLAction(query=generated))
            rewards.append(float(obs.reward))
        except Exception:
            rewards.append(0.0)

    return rewards


def _execution_reward(sql: str, prompt) -> float:
    """Simple execution check for merged_jsonl samples."""
    import sqlite3, re as _re

    # Extract schema from the user message
    user_content = ""
    for msg in (prompt if isinstance(prompt, list) else []):
        if isinstance(msg, dict) and msg.get("role") == "user":
            user_content = msg.get("content", "")
            break

    schema_match = _re.search(r"SCHEMA:\s*(.*?)\nQUESTION:", user_content, _re.DOTALL)
    if not schema_match:
        return 0.5  # can't verify, neutral reward

    schema_sql = schema_match.group(1).strip()
    try:
        conn = sqlite3.connect(":memory:")
        conn.executescript(schema_sql)
        rows = conn.execute(sql).fetchall()
        conn.close()
        return 1.0 if rows else 0.3   # ran cleanly but empty → partial credit
    except Exception:
        return 0.0


# ── Step 4: Train ─────────────────────────────────────────────────────────────

def main():
    merge_datasets()
    dataset = build_dataset()

    tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME, padding_side="right")
    if tokenizer.pad_token is None:
        tokenizer.pad_token = tokenizer.eos_token

    model = AutoModelForCausalLM.from_pretrained(
        MODEL_NAME,
        torch_dtype=torch.bfloat16,
        attn_implementation="sdpa"
    )

    peft_config = LoraConfig(
        r=128,
        lora_alpha=256,
        target_modules=["q_proj","k_proj","v_proj","o_proj","gate_proj","up_proj","down_proj"],
        bias="none",
        task_type="CAUSAL_LM"
    )

    training_args = GRPOConfig(
        output_dir=OUTPUT_DIR,
        learning_rate=1e-5,              # lower LR for fine-grained edge case tuning
        per_device_train_batch_size=2,
        gradient_accumulation_steps=4,
        max_completion_length=256,
        num_generations=8,
        temperature=0.5,
        bf16=True,
        logging_steps=5,
        num_train_epochs=5,              # fewer epochs — base knowledge already there
        report_to="none",
        remove_unused_columns=False,
        ddp_find_unused_parameters=False
    )

    trainer = GRPOTrainer(
        model=model,
        reward_funcs=sql_reward_func,
        args=training_args,
        train_dataset=dataset,
        peft_config=peft_config,
        processing_class=tokenizer
    )

    trainer.train()

    if trainer.accelerator.is_main_process:
        trainer.model.save_pretrained(f"{OUTPUT_DIR}/final")
        tokenizer.save_pretrained(f"{OUTPUT_DIR}/final")
        print(f"\nSaved to {OUTPUT_DIR}/final")


if __name__ == "__main__":
    main()