# /// script
# requires-python = ">=3.10"
# dependencies = ["unsloth", "datasets", "transformers", "huggingface_hub"]
# ///
"""Stage 1 + stage 2 mixed bf16 LoRA run for Qwen 3.8 27B on Hugging Face Jobs (Opus item E, 2026-10-06). NOT started: no job without Kral's go.

Memory test first (a few dollars, finds the largest sequence length that fits, no data needed):
  hf jobs uv run --flavor h200 --timeout 40m --secrets HF_TOKEN train/hf_train_bf16.py -- --memory-test --sweep 16000,32000,48000
Real run (after the memory test and Kral's decision on the ratio):
  hf jobs uv run --flavor <flavor> --timeout 10h --secrets HF_TOKEN train/hf_train_bf16.py -- \\
      --stage1 erhankeseli/abap-stage1-data --stage2 erhankeseli/abap-stage2-data --s1-epochs 1 --s2-epochs 3 --out erhankeseli/abap-mixed-adapter

Data: stage 1 rows have `text` (loss on every token); stage 2 rows have `text` and `assistant_spans` (loss only inside the spans: assistant turns;
none on the system turn with the tool schemas, the user turn or the tool results). A sample longer than --max-seq is skipped, never cut.
Base model in bf16 (no nf4): the adapter then fits the MLX 4-bit base better (stage 1 finding, train/STATE.md).
"""
import argparse
import json
import os
import random
import time


def tokenize_masked(tok, text, spans, max_len):
    """input_ids and labels; labels are -100 outside the spans (spans=None: loss on every token). None when longer than max_len."""
    enc = tok(text, add_special_tokens=False, return_offsets_mapping=True)
    ids = enc["input_ids"]
    if len(ids) > max_len:
        return None
    if spans is None:
        return ids, list(ids)
    labels = []
    for tid, (a, b) in zip(ids, enc["offset_mapping"]):
        labels.append(tid if any(a >= s and b <= e for s, e in spans) else -100)
    return ids, labels


def build_examples(tok, stage1, stage2, s1_epochs, s2_epochs, max_len, seed):
    rnd = random.Random(seed)
    rows = []
    for ep in range(int(s1_epochs)):
        rows += [("s1", r["text"], None, 1.0) for r in stage1]
    frac = s1_epochs - int(s1_epochs)
    if frac:
        rows += [("s1", r["text"], None, 1.0) for r in rnd.sample(stage1, int(len(stage1) * frac))]
    for ep in range(int(s2_epochs)):
        for r in stage2:
            rows += [("s2", r["text"], r["assistant_spans"], 1.0)] * max(1, round(float(r.get("weight", 1.0))))
    rnd.shuffle(rows)
    out, skipped = [], {"s1": 0, "s2": 0}
    for src, text, spans, _ in rows:
        t = tokenize_masked(tok, text, spans, max_len)
        if t is None:
            skipped[src] += 1
            continue
        out.append({"input_ids": t[0], "labels": t[1], "src": src})
    return out, skipped


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--model", default="Qwen/Qwen3.8-27B")
    ap.add_argument("--stage1", default="erhankeseli/abap-stage1-data")
    ap.add_argument("--stage2", default="erhankeseli/abap-stage2-data")
    ap.add_argument("--s1-file", default="train.jsonl")
    ap.add_argument("--s2-file", default="stage2_train.jsonl")
    ap.add_argument("--s2-valid", default="stage2_valid.jsonl")
    ap.add_argument("--s1-epochs", type=float, default=1.0)
    ap.add_argument("--s2-epochs", type=float, default=3.0)
    ap.add_argument("--max-seq", type=int, default=48000)
    ap.add_argument("--rank", type=int, default=16)
    ap.add_argument("--alpha", type=int, default=32)
    ap.add_argument("--lr", type=float, default=5e-5)
    ap.add_argument("--grad-accum", type=int, default=4)
    ap.add_argument("--warmup-steps", type=int, default=20)
    ap.add_argument("--no-offload", action="store_true", help="gradient checkpointing on the GPU instead of the Unsloth offload")
    ap.add_argument("--out", default="erhankeseli/abap-mixed-adapter")
    ap.add_argument("--save-every", type=int, default=0)
    ap.add_argument("--memory-test", action="store_true")
    ap.add_argument("--sweep", default="16000,32000,48000", help="memory test: sequence lengths, tried in this order, stops at the first OOM")
    ap.add_argument("--steps", type=int, default=3, help="memory test: optimizer steps per length")
    a = ap.parse_args()

    import torch
    from unsloth import FastLanguageModel
    from huggingface_hub import HfApi, hf_hub_download
    token = os.environ["HF_TOKEN"]
    model, tok = FastLanguageModel.from_pretrained(a.model, max_seq_length=a.max_seq, load_in_4bit=False, dtype=torch.bfloat16, token=token)
    model = FastLanguageModel.get_peft_model(
        model, r=a.rank, lora_alpha=a.alpha, lora_dropout=0.0, bias="none",
        target_modules=["q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj", "in_proj_qkv", "in_proj_z", "out_proj"],
        use_gradient_checkpointing=True if a.no_offload else "unsloth", random_state=20261006)
    print("GPU", torch.cuda.get_device_name(0), "x", torch.cuda.device_count(), "| weights allocated GB", round(torch.cuda.memory_allocated() / 2**30, 1), flush=True)

    if a.memory_test:
        text = "CLASS zcl_demo DEFINITION PUBLIC. METHODS run IMPORTING iv TYPE i. ENDCLASS.\n" * 4000
        ids_all = tok(text, add_special_tokens=False)["input_ids"]
        opt = torch.optim.AdamW([p for p in model.parameters() if p.requires_grad], lr=1e-5)
        model.train()
        result = {"gpu": torch.cuda.get_device_name(0), "gpus": torch.cuda.device_count(), "offload": not a.no_offload, "sweep": []}
        for n in [int(x) for x in a.sweep.split(",")]:
            ids = torch.tensor([(ids_all * (n // len(ids_all) + 1))[:n]], device="cuda")
            labels = ids.clone()
            labels[:, : int(n * 0.68)] = -100     # about 32 % loss tokens, as in the stage 2 data
            torch.cuda.reset_peak_memory_stats()
            try:
                times = []
                for _ in range(a.steps):
                    t0 = time.time()
                    loss = model(input_ids=ids, labels=labels).loss
                    loss.backward()
                    opt.step()
                    opt.zero_grad(set_to_none=True)
                    torch.cuda.synchronize()
                    times.append(round(time.time() - t0, 1))
                row = {"seq_len": n, "ok": True, "peak_gb": round(torch.cuda.max_memory_allocated() / 2**30, 1),
                       "reserved_gb": round(torch.cuda.max_memory_reserved() / 2**30, 1), "step_seconds": times, "loss": float(loss)}
            except torch.cuda.OutOfMemoryError as e:
                row = {"seq_len": n, "ok": False, "error": "OOM", "peak_gb": round(torch.cuda.max_memory_allocated() / 2**30, 1)}
                result["sweep"].append(row)
                print("MEMTEST", json.dumps(row), flush=True)
                break
            result["sweep"].append(row)
            print("MEMTEST", json.dumps(row), flush=True)
            torch.cuda.empty_cache()
        print("RESULT", json.dumps(result), flush=True)
        HfApi(token=token).upload_file(path_or_fileobj=json.dumps(result, indent=1).encode(), path_in_repo="memtest_%s_%dx.json" % (
            result["gpu"].replace(" ", "_"), result["gpus"]), repo_id=a.out, repo_type="model", create_pr=False) if a.out else None
        return

    def rows(repo, fname):
        p = hf_hub_download(repo, fname, repo_type="dataset", token=token)
        return [json.loads(l) for l in open(p)]
    s1, s2 = rows(a.stage1, a.s1_file), rows(a.stage2, a.s2_file)
    ex, skipped = build_examples(tok, s1, s2, a.s1_epochs, a.s2_epochs, a.max_seq, 20261006)
    loss_tok = sum(sum(1 for x in e["labels"] if x != -100) for e in ex)
    print("EXAMPLES", len(ex), "skipped", skipped, "loss tokens", loss_tok, "stage 2 share of loss tokens",
          round(sum(sum(1 for x in e["labels"] if x != -100) for e in ex if e["src"] == "s2") / max(loss_tok, 1), 2), flush=True)
    from datasets import Dataset
    from transformers import Trainer, TrainingArguments

    class Collate:
        def __call__(self, batch):      # batch size 1; no padding needed
            b = batch[0]
            return {"input_ids": torch.tensor([b["input_ids"]]), "labels": torch.tensor([b["labels"]]),
                    "attention_mask": torch.ones(1, len(b["input_ids"]), dtype=torch.long)}
    args = TrainingArguments(output_dir="out", per_device_train_batch_size=1, gradient_accumulation_steps=a.grad_accum, num_train_epochs=1,
                             learning_rate=a.lr, lr_scheduler_type="cosine", warmup_steps=a.warmup_steps, optim="adamw_8bit", bf16=True,
                             logging_steps=1, save_strategy="no", report_to="none", seed=20261006, remove_unused_columns=False)
    tr = Trainer(model=model, args=args, train_dataset=Dataset.from_list(ex), data_collator=Collate())
    t0 = time.time()
    tr.train()
    info = {"examples": len(ex), "skipped": skipped, "train_seconds": round(time.time() - t0, 1), "rank": a.rank, "alpha": a.alpha, "lr": a.lr,
            "peak_gpu_gb": round(torch.cuda.max_memory_allocated() / 2**30, 1), "s1_epochs": a.s1_epochs, "s2_epochs": a.s2_epochs}
    print("RESULT", json.dumps(info), flush=True)
    model.save_pretrained("adapter")
    json.dump(info, open("adapter/job_result.json", "w"))
    HfApi(token=token).upload_folder(folder_path="adapter", repo_id=a.out, repo_type="model")
    print("PUSHED", a.out)


if __name__ == "__main__":
    main()