# Text-to-SQL fine-tune of Qwen3.5-9B with Unsloth (QLoRA), measured.
# Phases: base eval on 200 held-out questions -> LoRA training -> same eval.
# Metrics: normalized exact match against the reference SQL, and the share
# that runs in SQLite against the question's own CREATE TABLE (valid SQL).
# Everything goes to /workspace/sqlft/; results.json at the end.
import json, os, re, sqlite3, sys, time
os.makedirs("/workspace/sqlft", exist_ok=True)
OUT = "/workspace/sqlft"
N_TEST, N_TRAIN, SEED = 200, int(os.environ.get("N_TRAIN", "8000")), 3407
MODEL = "unsloth/Qwen3.5-9B"
res = {"model": MODEL, "n_test": N_TEST, "n_train": N_TRAIN, "seed": SEED}

def log(*a):
    print(time.strftime("%H:%M:%S"), *a, flush=True)

from unsloth import FastModel
import torch
from datasets import load_dataset

ds = load_dataset("b-mc2/sql-create-context", split="train").shuffle(seed=SEED)
test = ds.select(range(N_TEST))
train = ds.select(range(N_TEST, N_TEST + N_TRAIN))
log("dataset", len(ds), "test", len(test), "train", len(train))

SYSTEM = "You translate questions into SQLite queries. Answer with the SQL query only, on one line, no explanation."
def user_msg(ex):
    return f"Tables:\n{ex['context']}\n\nQuestion: {ex['question']}"

t0 = time.time()
model, tok = FastModel.from_pretrained(MODEL, max_seq_length=1024, load_in_4bit=True, full_finetuning=False)
res["load_s"] = round(time.time() - t0, 1)
log("loaded in", res["load_s"], "s; vram", round(torch.cuda.memory_reserved() / 2**30, 2), "GB")
text_tok = getattr(tok, "tokenizer", tok)

def chat(msgs, gen):
    kw = dict(tokenize=False, add_generation_prompt=gen)
    try:
        return text_tok.apply_chat_template(msgs, enable_thinking=False, **kw)
    except TypeError:
        return text_tok.apply_chat_template(msgs, **kw)

def norm(s):
    s = s.strip().strip("`").strip()
    s = re.sub(r"^sql\s+", "", s, flags=re.I)
    s = s.split("\n")[0] if "\n" in s.strip() and s.strip().lower().startswith("select") else s
    s = s.replace('"', "'").rstrip(";").strip()
    return re.sub(r"\s+", " ", s).lower()

def runs(sql, context):
    db = sqlite3.connect(":memory:")
    try:
        for stmt in context.split(";"):
            if stmt.strip():
                db.execute(stmt)
        db.execute(sql)
        return True
    except Exception:
        return False
    finally:
        db.close()

def evaluate(tag):
    from unsloth import FastModel as FM
    FM.for_inference(model)
    text_tok.padding_side = "left"
    if text_tok.pad_token is None:
        text_tok.pad_token = text_tok.eos_token
    preds, t = [], time.time()
    for i in range(0, N_TEST, 25):
        batch = [test[j] for j in range(i, min(i + 25, N_TEST))]
        prompts = [chat([{"role": "system", "content": SYSTEM}, {"role": "user", "content": user_msg(ex)}], True) for ex in batch]
        enc = text_tok(prompts, return_tensors="pt", padding=True).to("cuda")
        with torch.no_grad():
            out = model.generate(**enc, max_new_tokens=160, do_sample=False, pad_token_id=text_tok.pad_token_id)
        for k, ex in enumerate(batch):
            txt = text_tok.decode(out[k][enc["input_ids"].shape[1]:], skip_special_tokens=True)
            txt = re.sub(r"<think>.*?</think>", "", txt, flags=re.S).strip()
            preds.append({"q": ex["question"], "ref": ex["answer"], "pred": txt})
    secs = time.time() - t
    em = sum(norm(p["pred"]) == norm(p["ref"]) for p in preds)
    ok = sum(runs(p["pred"].strip().strip("`").rstrip(";"), test[i]["context"]) for i, p in enumerate(preds))
    json.dump(preds, open(f"{OUT}/preds_{tag}.json", "w"), indent=1)
    r = {"exact_match": em, "valid_sql": ok, "of": N_TEST, "eval_s": round(secs, 1)}
    log(tag, r)
    return r

res["base"] = evaluate("base")
json.dump(res, open(f"{OUT}/results.json", "w"), indent=1)

# LoRA on the language layers only (the vision tower is untouched).
model = FastModel.get_peft_model(
    model, finetune_vision_layers=False, finetune_language_layers=True,
    finetune_attention_modules=True, finetune_mlp_modules=True,
    r=16, lora_alpha=16, lora_dropout=0, bias="none", random_state=SEED,
)
def fmt(ex):
    msgs = [{"role": "system", "content": SYSTEM}, {"role": "user", "content": user_msg(ex)},
            {"role": "assistant", "content": ex["answer"]}]
    return {"text": chat(msgs, False)}
train_txt = train.map(fmt, remove_columns=train.column_names)

from trl import SFTTrainer, SFTConfig
text_tok.padding_side = "right"
trainer = SFTTrainer(
    model=model, tokenizer=text_tok, train_dataset=train_txt,
    args=SFTConfig(
        dataset_text_field="text", max_seq_length=1024, per_device_train_batch_size=8,
        gradient_accumulation_steps=2, num_train_epochs=1, learning_rate=2e-4,
        warmup_steps=20, lr_scheduler_type="linear", optim="adamw_8bit", weight_decay=0.01,
        logging_steps=25, output_dir=f"{OUT}/ckpt", save_strategy="no", report_to="none", seed=SEED,
        bf16=True,
    ),
)
from unsloth.chat_templates import train_on_responses_only
try:
    trainer = train_on_responses_only(trainer, instruction_part="<|im_start|>user\n", response_part="<|im_start|>assistant\n")
except Exception as e:
    log("responses-only masking skipped:", e)
torch.cuda.reset_peak_memory_stats()
t = time.time()
stats = trainer.train()
res["train_s"] = round(time.time() - t, 1)
res["train_steps"] = stats.global_step
res["train_loss"] = round(stats.training_loss, 4)
res["peak_vram_gb"] = round(torch.cuda.max_memory_reserved() / 2**30, 2)
res["train_tokens"] = int(sum(len(text_tok(x)["input_ids"]) for x in train_txt["text"]))
log("trained", {k: res[k] for k in ("train_s", "train_steps", "train_loss", "peak_vram_gb", "train_tokens")})
json.dump(res, open(f"{OUT}/results.json", "w"), indent=1)

res["tuned"] = evaluate("tuned")
model.save_pretrained(f"{OUT}/lora")
text_tok.save_pretrained(f"{OUT}/lora")
json.dump(res, open(f"{OUT}/results.json", "w"), indent=1)
log("DONE", json.dumps(res))