"""Golden AI file-grounding stage: QLoRA on Qwen3-4B, free Colab T4, then GGUF Q4_K_M.

Why this is its own script and not colab_train.py: that one trains the retired
1.5B line at SEQ 1024 with its own hardcoded system prompt. This stage trains
the model production actually serves (Qwen3-4B), and every record carries its
OWN system message, because the whole point is the ATTACHED FILE block that
site/app.html pins there. Truncating it, or replacing it with a generic
identity prompt, would train the opposite of the behaviour we want.

Measured against the real Qwen3 tokenizer over data/splits/sft_file_grounding.jsonl:
243 records, median 771 tokens, p90 990, max 2873, answers median 66 tokens.
SEQ is 3072 so nothing is cut; a record longer than that is DROPPED with a
warning rather than silently losing its answer off the end.

Run inside Colab on a GPU runtime, with the dataset files next to it:

    GOLDEN_UPLOAD_URL=https://<tunnel> python colab_file_grounding.py

Produces /content/golden-ai-file-grounding-Q4_K_M.gguf plus the LoRA adapter.
"""

import json
import math
import os
import random
import subprocess
import sys
import tarfile
import time

os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True")

import torch
from torch.optim import AdamW
from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
from peft import LoraConfig, get_peft_model, prepare_model_for_kbit_training

BASE = os.environ.get("GOLDEN_BASE", "Qwen/Qwen3-4B")
TRAIN_FILE = os.environ.get("GOLDEN_TRAIN", "sft_file_grounding.jsonl")
VAL_FILE = os.environ.get("GOLDEN_VAL", "sft_file_grounding_val.jsonl")
OUT_NAME = os.environ.get("GOLDEN_OUT", "golden-ai-file-grounding")
QUANT = os.environ.get("GOLDEN_QUANT", "Q4_K_M")          # match what the VM serves

SEQ = int(os.environ.get("GOLDEN_SEQ", "3072"))
BATCH = 1
ACCUM = int(os.environ.get("GOLDEN_ACCUM", "4"))
EPOCHS = float(os.environ.get("GOLDEN_EPOCHS", "3"))
LR = float(os.environ.get("GOLDEN_LR", "1.5e-4"))
# v0.10 is the reason for this floor: 5e-5 for 38 steps moved a 1.5B not at all
# and cost a night. 243 records x 3 epochs / accum 4 is about 180 steps.
WARMUP = int(os.environ.get("GOLDEN_WARMUP", "10"))
RANK = int(os.environ.get("GOLDEN_RANK", "32"))
ALPHA = int(os.environ.get("GOLDEN_ALPHA", "64"))

UPLOAD_URL = os.environ.get("GOLDEN_UPLOAD_URL", "").rstrip("/")
SAVE_EVERY = int(os.environ.get("GOLDEN_SAVE_EVERY", "60"))
RESUME_ADAPTER = os.environ.get("GOLDEN_RESUME_ADAPTER", "")
SKIP_GGUF = os.environ.get("GOLDEN_SKIP_GGUF", "") == "1"
DRIVE_DIR = os.environ.get("GOLDEN_DRIVE_DIR", "")
# /content on Colab; overridden by --smoke so the laptop run has somewhere to write
OUT_DIR = os.environ.get("GOLDEN_OUT_DIR", "/content")

IGNORE = -100
# Two laptop modes, both meant to be run BEFORE booking a session, because a
# bug found on Colab costs the whole session:
#   --dry-run  tokenise and mask the dataset, assert the spans, stop
#   --smoke    additionally run the real loop for a few steps on CPU with a
#              tiny base model, so train/evaluate/checkpoint are exercised
DRY_RUN = "--dry-run" in sys.argv
SMOKE = "--smoke" in sys.argv
# recovery: re-run only merge -> GGUF -> quantise from a saved adapter
MERGE_ONLY = "--merge-only" in sys.argv
DEVICE = "cpu" if (DRY_RUN or SMOKE) else "cuda"

if SMOKE:
    BASE = os.environ.get("GOLDEN_SMOKE_BASE", "Qwen/Qwen3-0.6B")
    SEQ, ACCUM, EPOCHS, SAVE_EVERY = 1024, 2, 1, 0
    OUT_DIR = os.environ.get("GOLDEN_OUT_DIR", os.path.join(os.getcwd(), "build", "smoke"))
    os.makedirs(OUT_DIR, exist_ok=True)

if not (DRY_RUN or SMOKE):
    assert torch.cuda.is_available(), "select a GPU runtime first (Runtime > Change runtime type)"
    print("GPU:", torch.cuda.get_device_name(0),
          f"({torch.cuda.get_device_properties(0).total_memory / 1e9:.1f} GB)", flush=True)


def vram(tag=""):
    if DEVICE != "cuda":
        return
    used = torch.cuda.max_memory_allocated() / 1e9
    print(f"[vram] {tag} peak {used:.2f} GB", flush=True)


def template_ids(tok, messages, add_gen):
    # enable_thinking=False matters twice over. It puts the empty
    # <think></think> scaffold into the PROMPT, which is exactly what the VM
    # sends (llama.cpp --reasoning off; the live server returns no think block
    # and no reasoning_content), and it keeps that scaffold out of the labelled
    # span, so the model is not taught to emit one. Without it the generation
    # prompt stops at "assistant\n" and the scaffold lands in the answer.
    try:
        out = tok.apply_chat_template(messages, tokenize=True, add_generation_prompt=add_gen,
                                      enable_thinking=False)
    except TypeError:      # template does not take the kwarg
        out = tok.apply_chat_template(messages, tokenize=True, add_generation_prompt=add_gen)
    if hasattr(out, "input_ids"):
        out = out.input_ids
    elif isinstance(out, dict):
        out = out["input_ids"]
    if out and isinstance(out[0], (list, tuple)):
        out = out[0]
    return list(out)


def make_tarball(dest, root, name):
    r"""tarfile rather than the tar binary: on Windows `tar czf C:\...` reads the
    drive letter as a remote host spec and fails, and on Colab this removes a
    dependency on a binary being present at the one moment it matters."""
    with tarfile.open(dest, "w:gz") as tar:
        tar.add(os.path.join(root, name), arcname=name)
    return dest


def checkpoint(model, tok, tag):
    path = os.path.join(OUT_DIR, f"ckpt-{tag}")
    model.save_pretrained(path)
    tok.save_pretrained(path)
    if UPLOAD_URL:
        tar = os.path.join(OUT_DIR, f"ckpt-{tag}.tgz")
        make_tarball(tar, OUT_DIR, f"ckpt-{tag}")
        subprocess.Popen(["curl", "-s", "-T", tar, f"{UPLOAD_URL}/ckpt-{tag}.tgz"])
    print(f"[ckpt] saved {tag}", flush=True)


def load_model():
    if SMOKE:
        # bitsandbytes needs CUDA, so the CPU smoke test runs the same LoRA and
        # the same loop against a small fp32 model instead
        model = AutoModelForCausalLM.from_pretrained(BASE)
        model.config.use_cache = False
        if RESUME_ADAPTER:
            # the recovery path matters: free sessions get reclaimed, and both
            # earlier runs died mid-SFT. Exercise it here rather than
            # discovering it is broken at 2am on a reclaimed VM.
            from peft import PeftModel
            model = PeftModel.from_pretrained(model, RESUME_ADAPTER, is_trainable=True)
            print(f"resumed adapter from {RESUME_ADAPTER}", flush=True)
        else:
            lora = LoraConfig(r=8, lora_alpha=16, lora_dropout=0.05, bias="none",
                              task_type="CAUSAL_LM",
                              target_modules=["q_proj", "k_proj", "v_proj", "o_proj"])
            model = get_peft_model(model, lora)
        model.print_trainable_parameters()
        return model
    quant = BitsAndBytesConfig(load_in_4bit=True, bnb_4bit_quant_type="nf4",
                               bnb_4bit_use_double_quant=True,
                               bnb_4bit_compute_dtype=torch.float16)
    try:
        model = AutoModelForCausalLM.from_pretrained(BASE, quantization_config=quant,
                                                     device_map={"": 0}, dtype=torch.float16)
    except TypeError:      # transformers < 5 spells it torch_dtype
        model = AutoModelForCausalLM.from_pretrained(BASE, quantization_config=quant,
                                                     device_map={"": 0}, torch_dtype=torch.float16)
    model.config.use_cache = False
    model = prepare_model_for_kbit_training(model, use_gradient_checkpointing=True)
    if RESUME_ADAPTER:
        from peft import PeftModel
        model = PeftModel.from_pretrained(model, RESUME_ADAPTER, is_trainable=True)
        print(f"resumed adapter from {RESUME_ADAPTER}", flush=True)
    else:
        lora = LoraConfig(r=RANK, lora_alpha=ALPHA, lora_dropout=0.05, bias="none",
                          task_type="CAUSAL_LM",
                          target_modules=["q_proj", "k_proj", "v_proj", "o_proj",
                                          "gate_proj", "up_proj", "down_proj"])
        model = get_peft_model(model, lora)
    model.print_trainable_parameters()
    return model


def build_examples(tok, path, label):
    """(input_ids, labels, weight) with loss on assistant tokens only.

    The record's own system message is kept as-is: it holds the ATTACHED FILE
    block, which is the thing being taught."""
    examples, dropped, no_sys = [], 0, 0
    if not os.path.exists(path):
        print(f"{label}: {path} not found", flush=True)
        return examples
    for line in open(path, encoding="utf-8"):
        if not line.strip():
            continue
        rec = json.loads(line)
        msgs = rec["messages"]
        if not any(m["role"] == "system" for m in msgs):
            no_sys += 1
        full = template_ids(tok, msgs, False)
        if len(full) > SEQ:
            # truncation would cut the ANSWER, not the file block: refuse it
            dropped += 1
            continue
        labels = [IGNORE] * len(full)
        for i, m in enumerate(msgs):
            if m["role"] != "assistant":
                continue
            a = len(template_ids(tok, msgs[:i], True))
            b = len(template_ids(tok, msgs[:i + 1], False))
            if not (0 <= a < b <= len(full)):
                print(f"{label}: non-prefix-stable turn, record skipped", flush=True)
                labels = None
                break
            labels[a:b] = full[a:b]
        if labels is None or all(l == IGNORE for l in labels):
            dropped += 1
            continue
        examples.append((full, labels, float(rec.get("metadata", {}).get("weight", 1.0))))
    lens = sorted(len(e[0]) for e in examples)
    ans = sum(sum(1 for l in e[1] if l != IGNORE) for e in examples)
    print(f"{label}: {len(examples)} examples, {dropped} dropped, {no_sys} without a system message; "
          f"tokens median {lens[len(lens) // 2] if lens else 0} max {lens[-1] if lens else 0}, "
          f"{ans:,} supervised answer tokens", flush=True)
    return examples


def forward_loss(model, example, scale_by):
    ids = torch.tensor([example[0]], device=DEVICE)
    labels = torch.tensor([example[1]], device=DEVICE)
    logits = model(input_ids=ids, attention_mask=torch.ones_like(ids)).logits
    shift_logits = logits[:, :-1, :]
    shift_labels = labels[:, 1:]
    keep = shift_labels != IGNORE
    # Only the answer tokens carry loss, so slice BEFORE cross entropy: at
    # SEQ 3072 with a 151k vocab the full-width contiguous copy is 0.9 GB of
    # VRAM for nothing.
    sel_logits = shift_logits[keep]
    sel_labels = shift_labels[keep]
    loss = torch.nn.functional.cross_entropy(sel_logits.float(), sel_labels, reduction="sum")
    return loss * example[2] / max(1, scale_by)


def evaluate(model, examples, label):
    if not examples:
        return float("nan")
    model.eval()
    total, tokens = 0.0, 0
    with torch.no_grad():
        for ex in examples:
            n = sum(1 for l in ex[1] if l != IGNORE)
            loss = forward_loss(model, (ex[0], ex[1], 1.0), 1)
            total += loss.item()
            tokens += n
    model.train()
    value = total / max(1, tokens)
    print(f"[eval] {label} loss {value:.4f} over {tokens} answer tokens", flush=True)
    return value


def train(model, tok, examples, val):
    steps_per_epoch = math.ceil(len(examples) / (BATCH * ACCUM))
    total_steps = max(1, int(steps_per_epoch * EPOCHS))
    params = [p for p in model.parameters() if p.requires_grad]
    optimizer = AdamW(params, lr=LR)

    def lr_at(step):
        if step < WARMUP:
            return (step + 1) / max(1, WARMUP)
        progress = (step - WARMUP) / max(1, total_steps - WARMUP)
        return 0.5 * (1 + math.cos(math.pi * min(1.0, progress)))

    schedule = torch.optim.lr_scheduler.LambdaLR(optimizer, lr_at)
    print(f"training: {len(examples)} examples, {total_steps} optimizer steps "
          f"(batch {BATCH} x accum {ACCUM}), lr {LR}", flush=True)
    evaluate(model, val, "val@0")

    model.train()
    step, t0, epoch = 0, time.time(), 0
    oom = 0
    while step < total_steps:
        epoch += 1
        order = list(range(len(examples)))
        random.Random(8159 + epoch).shuffle(order)
        for start in range(0, len(order), BATCH * ACCUM):
            if step >= total_steps:
                break
            window = order[start:start + BATCH * ACCUM]
            total_tokens = sum(sum(1 for l in examples[i][1] if l != IGNORE) for i in window)
            if total_tokens == 0:
                continue
            optimizer.zero_grad(set_to_none=True)
            shown = 0.0
            for i in window:
                try:
                    loss = forward_loss(model, examples[i], total_tokens)
                    loss.backward()
                    shown += loss.item()
                except torch.cuda.OutOfMemoryError:
                    # one long record must not end a two hour session
                    oom += 1
                    print(f"[oom] skipped a {len(examples[i][0])} token example", flush=True)
                    torch.cuda.empty_cache()
            torch.nn.utils.clip_grad_norm_(params, 1.0)
            optimizer.step()
            schedule.step()
            step += 1
            if step % 5 == 0 or step == total_steps:
                rate = (time.time() - t0) / step
                print(f"[sft] step {step}/{total_steps} loss {shown:.3f} "
                      f"lr {schedule.get_last_lr()[0]:.2e} {rate:.1f}s/step "
                      f"eta {(total_steps - step) * rate / 60:.0f}m", flush=True)
            if SAVE_EVERY > 0 and step % SAVE_EVERY == 0:
                checkpoint(model, tok, f"sft-{step}")
                evaluate(model, val, f"val@{step}")
    vram("after training")
    if oom:
        print(f"WARNING: {oom} examples were skipped for VRAM", flush=True)
    return model


def to_gguf(tok):
    print("merging at fp16...", flush=True)
    try:
        import torchao  # noqa: F401
        ok = tuple(int(x) for x in torchao.__version__.split(".")[:2]) >= (0, 16)
        if not ok:
            print("WARNING: torchao " + torchao.__version__ + " is older than transformers wants. "
                  "If the merge raises ImportError, run `pip install -U torchao` and rerun with "
                  "--merge-only; the adapter is already saved.", flush=True)
    except Exception:
        pass
    torch.cuda.empty_cache()
    from peft import PeftModel
    try:
        base = AutoModelForCausalLM.from_pretrained(BASE, dtype=torch.float16, device_map={"": 0})
    except TypeError:
        base = AutoModelForCausalLM.from_pretrained(BASE, torch_dtype=torch.float16, device_map={"": 0})
    merged = PeftModel.from_pretrained(base, os.path.join(OUT_DIR, "golden-adapter")).merge_and_unload()
    # shard small: Colab's free VM has 12.7 GB of RAM and an 8 GB single shard
    # is how you lose the run at the very last step
    merged_dir = os.path.join(OUT_DIR, "golden-merged")
    merged.save_pretrained(merged_dir, safe_serialization=True, max_shard_size="2GB")
    tok.save_pretrained(merged_dir)
    del merged, base
    torch.cuda.empty_cache()

    if not os.path.exists("llama.cpp"):
        subprocess.run(["git", "clone", "--depth", "1",
                        "https://github.com/ggml-org/llama.cpp"], check=True)
    subprocess.run([sys.executable, "-m", "pip", "-q", "install", "gguf", "sentencepiece"], check=True)
    f16 = os.path.join(OUT_DIR, f"{OUT_NAME}-f16.gguf")
    print("converting to GGUF f16...", flush=True)
    subprocess.run([sys.executable, "llama.cpp/convert_hf_to_gguf.py", merged_dir,
                    "--outfile", f16, "--outtype", "f16"], check=True)

    out = os.path.join(OUT_DIR, f"{OUT_NAME}-{QUANT}.gguf")
    print(f"quantising to {QUANT} (matching the VM)...", flush=True)
    if not os.path.exists("llama.cpp/build/bin/llama-quantize"):
        subprocess.run(["cmake", "-S", "llama.cpp", "-B", "llama.cpp/build",
                        "-DGGML_NATIVE=OFF", "-DLLAMA_CURL=OFF", "-DCMAKE_BUILD_TYPE=Release"], check=True)
        subprocess.run(["cmake", "--build", "llama.cpp/build", "--target", "llama-quantize",
                        "-j", str(os.cpu_count() or 2)], check=True)
    subprocess.run(["llama.cpp/build/bin/llama-quantize", f16, out, QUANT], check=True)
    os.remove(f16)
    size = os.path.getsize(out) / 1e9
    print(f"DONE: {out} ({size:.2f} GB)", flush=True)
    return out


def dry_run(tok, train_examples, val_examples):
    """Prove the masking is right before a session is spent on it: the labelled
    span must decode back to exactly the record's assistant text, and every
    file-grounding record's prompt must still contain its ATTACHED FILE block.
    Rebuilds each example from the raw record so a dropped one cannot shift the
    pairing."""
    checked, with_file, anchors, skipped = 0, 0, 0, 0
    for line in open(TRAIN_FILE, encoding="utf-8"):
        if not line.strip():
            continue
        rec = json.loads(line)
        rid = rec.get("metadata", {}).get("record_id", "?")
        msgs = rec["messages"]
        full = template_ids(tok, msgs, False)
        if len(full) > SEQ:
            skipped += 1
            continue
        a = len(template_ids(tok, msgs[:-1], True))
        b = len(full)
        body = tok.decode(full[a:b]).replace("<|im_end|>", "").strip()
        want = msgs[-1]["content"].strip()
        assert body == want, f"masking mismatch in {rid}:\n  got  {body[:140]!r}\n  want {want[:140]!r}"
        prompt = tok.decode(full[:a])
        assert prompt.endswith("<think>\n\n</think>\n\n"), (
            f"{rid}: prompt does not end with the empty think scaffold the VM sends: {prompt[-40:]!r}")
        if rid.startswith("file_grounding"):
            assert "=== ATTACHED FILE" in prompt, f"file block missing from the prompt of {rid}"
            with_file += 1
        else:
            anchors += 1
        checked += 1
    steps = math.ceil(len(train_examples) / (BATCH * ACCUM)) * EPOCHS
    print(f"dry run OK: {checked} records mask to exactly their assistant text "
          f"({with_file} carry a file block, {anchors} are anchors, {skipped} over SEQ)", flush=True)
    print(f"plan: {len(train_examples)} train / {len(val_examples)} val, "
          f"{int(steps)} optimizer steps at lr {LR}, SEQ {SEQ}, base {BASE}", flush=True)


def main():
    if MERGE_ONLY:
        # Recovery path. The 2026-08-25 run trained all 183 steps and then died
        # in the merge with "incompatible version of torchao 0.10.0, only
        # versions above 0.16.0 are supported" -- a Colab image problem, not a
        # training one. The adapter survives, so this re-does only the tail:
        #   !pip -q install -U torchao
        #   !python colab_file_grounding.py --merge-only
        adapter = os.path.join(OUT_DIR, "golden-adapter")
        if not os.path.isdir(adapter):
            raise SystemExit("no adapter at " + adapter + "; nothing to merge")
        print("merge-only: using the adapter already on disk at " + adapter, flush=True)
        to_gguf(AutoTokenizer.from_pretrained(BASE))
        return

    tok = AutoTokenizer.from_pretrained(BASE)
    train_examples = build_examples(tok, TRAIN_FILE, "train")
    val_examples = build_examples(tok, VAL_FILE, "val")
    if not train_examples:
        raise SystemExit(f"no training examples in {TRAIN_FILE}")
    if DRY_RUN:
        dry_run(tok, train_examples, val_examples)
        return
    if SMOKE:
        # the point is to exercise the loop, not to learn anything
        train_examples = train_examples[:6]
        val_examples = val_examples[:2]
        print("SMOKE: CPU run over %d train / %d val examples on %s"
              % (len(train_examples), len(val_examples), BASE), flush=True)

    model = load_model()
    model = train(model, tok, train_examples, val_examples)
    evaluate(model, val_examples, "val@final")
    checkpoint(model, tok, "sft-final")

    adapter_dir = os.path.join(OUT_DIR, "golden-adapter")
    model.save_pretrained(adapter_dir)
    tok.save_pretrained(adapter_dir)
    tarball = os.path.join(OUT_DIR, "golden-adapter.tgz")
    make_tarball(tarball, OUT_DIR, "golden-adapter")
    print(f"adapter: {tarball} ({os.path.getsize(tarball) / 1e6:.0f} MB)", flush=True)
    if UPLOAD_URL:
        subprocess.run(["curl", "-s", "-T", tarball,
                        f"{UPLOAD_URL}/golden-adapter.tgz"], check=False)

    del model
    if DEVICE == "cuda":
        torch.cuda.empty_cache()
    if SMOKE:
        print("SMOKE OK: build -> train -> eval -> checkpoint -> adapter all ran", flush=True)
        return
    if SKIP_GGUF:
        print("GOLDEN_SKIP_GGUF=1: stopping before the merge", flush=True)
        return
    out = to_gguf(tok)
    if DRIVE_DIR:
        os.makedirs(DRIVE_DIR, exist_ok=True)
        subprocess.run(["cp", out, DRIVE_DIR], check=True)
        print(f"copied to {DRIVE_DIR}", flush=True)


if __name__ == "__main__":
    main()