code/m2w_train_lora.py

#!/usr/bin/env python3
"""Mind2Web: Kev tarzı pointer kafası + LoRA, Gemma 4 E2B üzerinde (v1: YALNIZCA METİN).

Girdi: m2w_ranker.py çıktısı (m2w_top50/*.jsonl.gz).
Tek ileri geçiş: [görev + geçmiş] [50 aday satırı] [görev + geçmiş + "Next element:"].
Sorgu = son token'ın gizli durumu (katman --layer). Aday anahtarı = o adayın token'larının ortalaması.
Skor = Wq(q)·Wk(k_i)/sqrt(dk), 50 aday + öğrenilen "hiçbiri" anahtarı üzerinde softmax.
Ayrıca sorgudan işlem türü kafası (CLICK/TYPE/SELECT).

!! TEST EDİLMEDİ: Gemma 4 / peft ile bu betik sandbox'ta çalıştırılamadı (torch yok). Colab'da
ilk çalıştırmada hata çıkabilir; --max_train 40 --max_val 40 ile önce küçük deneyin.

Kullanım:
    python m2w_train_lora.py --data m2w_top50 --out run_L30 --layer 30 --epochs 2
    python m2w_train_lora.py --data m2w_smoke_top50 --out smoke --max_train 40 --max_val 40 --epochs 1
"""
import argparse
import gzip
import json
import math
import os
import re
import time

import numpy as np
import torch
from torch import nn

import m2w_metrics as M

MODEL_ID = "google/gemma-4-E2B"
REV = "d29ff6b45f081a49ee2733a859c9c9c2d95d1a6f"
OPS = {"CLICK": 0, "TYPE": 1, "SELECT": 2}
MAX_CAND_TOK = 24
CAND_TOK = MAX_CAND_TOK   # --max_cand_tok ile değişir
MAX_SEQ = 0               # --max_seq: bir adımın toplam token sayısı bunu aşarsa aday satırları o adım için eşit oranda kısaltılır (0 = sınırsız)
MAX_HIST = 5


def read(path):
    with gzip.open(path, "rt", encoding="utf-8") as f:
        return [json.loads(l) for l in f]


# --------------------------------------------------------------- girdi kurma
def build_ids(tok, ex, order, max_cand_tok=MAX_CAND_TOK, max_seq=None):
    """(input_ids, aday aralıkları, sorgu indeksi). Her parça ayrı tokenlanır, aralıklar kesin bilinir.
    max_seq doluysa ve toplam aşılıyorsa: aday satırlarının hepsi için ortak bir üst sınır bulunur (en büyük c, toplam <= bütçe)."""
    max_seq = MAX_SEQ if max_seq is None else max_seq
    hist = " ; ".join(h[:80] for h in ex["history"][-MAX_HIST:]) or "(none)"
    head = f"Task: {ex['task']}\nPrevious actions: {hist}\nCandidates:\n"
    tail = f"Task: {ex['task']}\nPrevious actions: {hist}\nNext element to interact with:"
    ids = list(tok(head, add_special_tokens=True)["input_ids"])
    tail_ids = list(tok(tail, add_special_tokens=False)["input_ids"])
    toks = [tok("- " + ex["cands"][j]["text"] + "\n", add_special_tokens=False, truncation=True,
                max_length=max_cand_tok)["input_ids"] for j in order]
    if max_seq:
        budget = max_seq - len(ids) - len(tail_ids)
        if sum(len(t) for t in toks) > budget:
            lo, hi = 4, max_cand_tok
            while lo < hi:
                mid = (lo + hi + 1) // 2
                if sum(min(len(t), mid) for t in toks) <= budget:
                    lo = mid
                else:
                    hi = mid - 1
            toks = [t[:lo] for t in toks]
    spans = []
    for t in toks:
        spans.append((len(ids), len(ids) + len(t)))
        ids += list(t)
    ids += tail_ids
    return ids, spans, len(ids) - 1


def collate(tok, exs, orders, device):
    items = [build_ids(tok, ex, od, CAND_TOK, MAX_SEQ) for ex, od in zip(exs, orders)]
    L = max(len(i[0]) for i in items)
    pad = tok.pad_token_id if tok.pad_token_id is not None else 0
    ids = torch.full((len(items), L), pad, dtype=torch.long)
    att = torch.zeros((len(items), L), dtype=torch.long)
    for b, (x, _, _) in enumerate(items):
        ids[b, :len(x)] = torch.tensor(x)
        att[b, :len(x)] = 1
    return ids.to(device), att.to(device), [i[1] for i in items], [i[2] for i in items]


# --------------------------------------------------------------------- kafa
class Head(nn.Module):
    def __init__(self, d, dk=256):
        super().__init__()
        self.ln = nn.LayerNorm(d)
        self.wq = nn.Linear(d, dk)
        self.wk = nn.Linear(d, dk)
        self.null = nn.Parameter(torch.randn(dk) * 0.02)   # "hiçbiri" anahtarı
        self.op = nn.Linear(d, len(OPS))
        self.dk = dk

    def forward(self, h, spans, qidx):
        """h: [B,L,d]. Dönen: aday+hiçbiri logit listesi (B tane [n+1]) ve işlem logitleri [B,3]."""
        hq = self.ln(h[torch.arange(h.shape[0], device=h.device), torch.tensor(qidx, device=h.device)].float())
        q = self.wq(hq)
        logits = []
        for b, sp in enumerate(spans):
            v = torch.stack([h[b, s:e].float().mean(0) for s, e in sp])
            k = self.wk(self.ln(v))
            k = torch.cat([k, self.null[None]], 0)
            logits.append((k @ q[b]) / math.sqrt(self.dk))
        return logits, self.op(hq)


def lora_names(model, max_layer, mlp=False):
    names = []
    targets = ("q_proj", "k_proj", "v_proj", "o_proj") + (("gate_proj", "up_proj", "down_proj") if mlp else ())
    for n, m in model.named_modules():
        if isinstance(m, nn.Linear) and n.split(".")[-1] in targets \
                and not any(s in n for s in ("vision", "audio")):
            mm = re.search(r"layers\.(\d+)\.", n)
            if mm and int(mm.group(1)) < max_layer:
                names.append(n)
    return names


def run_model(model, ids, att, layer):
    kw = dict(input_ids=ids, attention_mask=att, output_hidden_states=True)
    try:
        out = model(logits_to_keep=1, **kw)
    except TypeError:
        out = model(**kw)
    return out.hidden_states[layer]


def loss_fn(logits, ex, order, op_logits):
    """Çoklu pozitif: -log sum_{p in P} softmax. Pozitif listede yoksa hedef = hiçbiri."""
    pos = set(ex["pos_idx"])
    tgt = [r for r, j in enumerate(order) if j in pos] or [len(order)]
    l = -(torch.logsumexp(logits[tgt], 0) - torch.logsumexp(logits, 0))
    return l


# ---------------------------------------------------------------- değerlendirme
@torch.no_grad()
def evaluate(model, head, tok, rows, layer, device, bs=4, mode="ranker", rng=None):
    model.eval(); head.eval()
    res = []
    for i in range(0, len(rows), bs):
        exs = rows[i:i + bs]
        if mode == "reverse":
            orders = [list(range(len(e["cands"])))[::-1] for e in exs]
        else:
            orders = [list(range(len(e["cands"]))) for e in exs]
        ids, att, spans, qidx = collate(tok, exs, orders, device)
        with torch.autocast("cuda", dtype=torch.bfloat16, enabled=device == "cuda"):
            h = run_model(model, ids, att, layer)
        logits, opl = head(h, spans, qidx)
        for b, ex in enumerate(exs):
            p = torch.softmax(logits[b], 0).cpu().numpy()
            n = len(orders[b])
            j = int(p.argmax())
            pred_bid = ex["cands"][orders[b][j]]["bid"] if j < n else "NULL"
            pos = set(ex["pos_idx"])
            correct = j < n and orders[b][j] in pos
            res.append({"conf": float(p[j]), "correct": bool(correct), "pred": pred_bid,
                        "null": j == n, "has_pos": len(pos) > 0,
                        "op_ok": int(opl[b].argmax()) == OPS.get(ex["op"], 0),
                        "p_null": float(p[n])})
    return res


def report(res):
    conf = [r["conf"] for r in res]; cor = [r["correct"] for r in res]
    s = M.summarize(conf, cor)
    has = [r for r in res if r["has_pos"]]
    s["acc_pos_in_list"] = float(np.mean([r["correct"] for r in has])) if has else float("nan")
    s["null_pred_rate"] = float(np.mean([r["null"] for r in res]))
    s["op_acc"] = float(np.mean([r["op_ok"] for r in res]))
    nohas = [r for r in res if not r["has_pos"]]
    s["abstain_ok_when_no_pos"] = float(np.mean([r["null"] for r in nohas])) if nohas else float("nan")
    return s


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--data", default="m2w_top50")
    ap.add_argument("--out", default="run")
    ap.add_argument("--layer", type=int, default=30, help="hidden_states[layer] okunur (0=gömme, 35=son)")
    ap.add_argument("--epochs", type=int, default=2)
    ap.add_argument("--mb", type=int, default=2, help="mikro-batch")
    ap.add_argument("--accum", type=int, default=4)
    ap.add_argument("--lr_lora", type=float, default=2e-4)
    ap.add_argument("--lr_head", type=float, default=1e-3)
    ap.add_argument("--rank", type=int, default=16)
    ap.add_argument("--max_train", type=int, default=0)
    ap.add_argument("--max_val", type=int, default=0)
    ap.add_argument("--no_shuffle", action="store_true", help="eğitimde aday sırasını karıştırma")
    ap.add_argument("--grad_ckpt", action="store_true",
                    help="Gradient checkpointing (KV paylaşımlı katmanlarda hata riski var, varsayılan kapalı)")
    ap.add_argument("--eval_splits", nargs="+", default=["val", "test_task", "test_website", "test_domain"])
    ap.add_argument("--seed", type=int, default=0)
    ap.add_argument("--max_cand_tok", type=int, default=MAX_CAND_TOK, help="aday satırı başına en çok token")
    ap.add_argument("--lora_mlp", action="store_true", help="LoRA'yı MLP (gate/up/down) projeksiyonlarına da koy")
    ap.add_argument("--always_pos", action="store_true",
                    help="EĞİTİMDE süzgecin attığı doğru adayı (dropped_pos) listeye geri ekle; val/test aynı kalır")
    ap.add_argument("--train_sample", type=int, default=0, help="eğitimden rastgele bu kadar adım (sabit tohum 123, tüm denemelerde aynı)")
    ap.add_argument("--train_eval_n", type=int, default=300, help="sonda eğitim kesitinde ölçülen adım sayısı (aşırı uyum göstergesi)")
    ap.add_argument("--no_reverse", action="store_true", help="val'de ters sıra geçişini atla (zaman kazandırır, sıra değişim oranı yazılmaz)")
    ap.add_argument("--max_seq", type=int, default=0, help="bir adımın en çok token sayısı (0 = sınırsız); aşılırsa aday satırları kısaltılır (bellek koruması)")
    args = ap.parse_args()
    global CAND_TOK, MAX_SEQ
    CAND_TOK = args.max_cand_tok; MAX_SEQ = args.max_seq

    from peft import LoraConfig, get_peft_model
    from transformers import AutoModelForMultimodalLM, AutoTokenizer

    device = "cuda" if torch.cuda.is_available() else "cpu"
    torch.manual_seed(args.seed); rng = np.random.default_rng(args.seed)
    os.makedirs(args.out, exist_ok=True)

    data = {s: read(os.path.join(args.data, f"{s}.jsonl.gz")) for s in ["train"] + args.eval_splits
            if os.path.exists(os.path.join(args.data, f"{s}.jsonl.gz"))}
    train = data["train"]
    if "val" not in data:   # küçük denemede val bölümü olmayabilir
        n = max(2, int(0.2 * len(train)))
        data["val"], train = train[-n:], train[:-n]
        print(f"val yok: eğitimin son {n} adımı val yapıldı")
    if args.train_sample and args.train_sample < len(train):
        idx = np.random.default_rng(123).permutation(len(train))[: args.train_sample]
        train = [train[i] for i in sorted(idx)]
    if args.max_train:
        train = train[: args.max_train]
    train_orig = train
    if args.always_pos:
        fixed, added = [], 0
        for e in train:
            dp = e.get("dropped_pos") or []
            if dp:
                e = dict(e); n0 = len(e["cands"])
                e["cands"] = list(e["cands"]) + list(dp)
                e["pos_idx"] = list(e["pos_idx"]) + list(range(n0, n0 + len(dp))); added += 1
            fixed.append(e)
        train = fixed
        print(f"doğruyu ekle: {added}/{len(train)} eğitim adımına atılan doğru aday geri eklendi; hâlâ doğrusu olmayan: "
              f"{sum(1 for e in train if not e['pos_idx'])}")
    for s in list(data):
        if s != "train" and args.max_val:
            data[s] = data[s][: args.max_val]
    print({s: len(v) for s, v in data.items()}, "| eğitim kullanılan:", len(train))

    tok = AutoTokenizer.from_pretrained(MODEL_ID, revision=REV)
    model = AutoModelForMultimodalLM.from_pretrained(MODEL_ID, revision=REV, dtype=torch.bfloat16).to(device)
    names = lora_names(model, args.layer, args.lora_mlp)
    if not names:
        qs = [(n, type(m).__name__) for n, m in model.named_modules() if n.endswith("q_proj")]
        print("q_proj modülleri (ilk 12):", qs[:12])
        raise SystemExit("LoRA hedef modül bulunamadı: yukarıdaki adlara göre lora_names() düzeltilmeli")
    print("LoRA hedef modül sayısı:", len(names), "| örnek:", names[0], names[-1])
    model = get_peft_model(model, LoraConfig(r=args.rank, lora_alpha=2 * args.rank, lora_dropout=0.05,
                                             target_modules=names, bias="none"))
    if args.grad_ckpt:
        model.gradient_checkpointing_enable(); model.enable_input_require_grads()
    model.print_trainable_parameters()
    try:
        d = model.config.get_text_config().hidden_size
    except Exception:
        d = 1536
    head = Head(d).to(device)

    lora_params = [p for p in model.parameters() if p.requires_grad]
    opt = torch.optim.AdamW([{"params": lora_params, "lr": args.lr_lora, "weight_decay": 0.0},
                             {"params": head.parameters(), "lr": args.lr_head, "weight_decay": 0.01}])
    steps_per_epoch = math.ceil(len(train) / (args.mb * args.accum))
    total = steps_per_epoch * args.epochs
    warm = max(1, int(0.03 * total))
    sched = torch.optim.lr_scheduler.LambdaLR(
        opt, lambda s: (s + 1) / warm if s < warm else 0.5 * (1 + math.cos(math.pi * (s - warm) / max(1, total - warm))))

    best, log = -1.0, []
    n_oom = 0
    t0 = time.time(); step = 0
    for ep in range(args.epochs):
        model.train(); head.train()
        perm = rng.permutation(len(train))
        run_loss, run_n = 0.0, 0
        for i in range(0, len(perm), args.mb):
            exs = [train[j] for j in perm[i:i + args.mb]]
            orders = [list(rng.permutation(len(e["cands"]))) if not args.no_shuffle else list(range(len(e["cands"])))
                      for e in exs]
            ids, att, spans, qidx = collate(tok, exs, orders, device)
            if ep == 0 and i == 0:
                print(f"ilk batch: girdi {tuple(ids.shape)} | aday/adım {[len(sp) for sp in spans]} "
                      f"| aday başına ort. token {np.mean([e - s for sp in spans for s, e in sp]):.1f}", flush=True)
            try:
                with torch.autocast("cuda", dtype=torch.bfloat16, enabled=device == "cuda"):
                    h = run_model(model, ids, att, args.layer)
                logits, opl = head(h, spans, qidx)
                l_sel = torch.stack([loss_fn(logits[b], exs[b], orders[b], opl[b]) for b in range(len(exs))]).mean()
                tgt_op = torch.tensor([OPS.get(e["op"], 0) for e in exs], device=device)
                loss = l_sel + 0.3 * nn.functional.cross_entropy(opl, tgt_op)
                (loss / args.accum).backward()
                run_loss += float(l_sel.detach()); run_n += 1
            except torch.cuda.OutOfMemoryError:
                # çok uzun nadir adım: atla, bu birikim penceresinin yarım gradyanını da sil, devam et
                n_oom += 1
                h = logits = opl = l_sel = loss = None
                opt.zero_grad(set_to_none=True); torch.cuda.empty_cache()
                print(f"!!! bellek taştı (girdi {tuple(ids.shape)}), adım atlandı (toplam {n_oom})", flush=True)
            if (i // args.mb + 1) % args.accum == 0 or i + args.mb >= len(perm):
                nn.utils.clip_grad_norm_(lora_params + list(head.parameters()), 1.0)
                opt.step(); sched.step(); opt.zero_grad(set_to_none=True); step += 1
                if step % 25 == 0:
                    print(f"ep {ep} adım {step}/{total} | seçim kaybı {run_loss / run_n:.3f} | {time.time() - t0:.0f}s", flush=True)
                    run_loss, run_n = 0.0, 0
        mem = torch.cuda.max_memory_allocated() / 2 ** 30 if device == "cuda" else 0
        print(f"epoch {ep} eğitim süresi {time.time() - t0:.0f}s ({(time.time() - t0) / max(1, len(perm) * (ep + 1)):.2f} sn/örnek) "
              f"| tepe bellek {mem:.1f} GB", flush=True)
        if "val" in data:
            r = report(evaluate(model, head, tok, data["val"], args.layer, device))
            print(f"== epoch {ep} val: acc {r['acc']:.3f} (ece {r['ece']:.3f}, auroc {r['auroc']:.3f}, "
                  f"op {r['op_acc']:.3f}, liste içi {r['acc_pos_in_list']:.3f})", flush=True)
            log.append({"epoch": ep, **{k: v for k, v in r.items() if k != "acc_ci95"}})
            if r["acc"] > best:
                best = r["acc"]
                model.save_pretrained(os.path.join(args.out, "lora"))
                torch.save(head.state_dict(), os.path.join(args.out, "head.pt"))

    # en iyi epoch'u yükle ve tüm bölümlerde değerlendir
    if os.path.exists(os.path.join(args.out, "head.pt")):
        head.load_state_dict(torch.load(os.path.join(args.out, "head.pt"), map_location=device))
        from peft import set_peft_model_state_dict
        try:
            from safetensors.torch import load_file
            set_peft_model_state_dict(model, load_file(os.path.join(args.out, "lora", "adapter_model.safetensors")))
        except Exception as e:
            print("LoRA en iyi epoch yüklenemedi, son epoch ile devam:", e)
    train_time = time.time() - t0
    final = {}
    for s in args.eval_splits:
        if s not in data:
            continue
        res = evaluate(model, head, tok, data[s], args.layer, device)
        final[s] = report(res)
        if s == "val" and not args.no_reverse:
            rev = evaluate(model, head, tok, data[s], args.layer, device, mode="reverse")
            flips = np.mean([a["pred"] != b["pred"] for a, b in zip(res, rev)])
            final[s]["sira_degisim_orani_ters_siralama"] = float(flips)
        with open(os.path.join(args.out, f"tahmin_{s}.json"), "w") as f:
            json.dump([{k: v for k, v in r.items()} for r in res], f)
        print(f"\n[{s}] " + json.dumps({k: (round(v, 3) if isinstance(v, float) else v) for k, v in final[s].items()
                                       if k != 'acc_ci95'}, ensure_ascii=False))
        print("   acc %95 aralığı:", [round(x, 3) for x in final[s]["acc_ci95"]])
    if args.train_eval_n:
        sub = [train_orig[i] for i in np.random.default_rng(7).permutation(len(train_orig))[: args.train_eval_n]]
        tr = report(evaluate(model, head, tok, sub, args.layer, device))
        final["egitim_kesiti"] = {k: v for k, v in tr.items() if k != "acc_ci95"}
        print("\n[egitim kesiti, aşırı uyum göstergesi]", {k: round(v, 3) for k, v in final["egitim_kesiti"].items()
                                                         if k in ("acc", "acc_pos_in_list", "ece")})
    with open(os.path.join(args.out, "sonuc.json"), "w", encoding="utf-8") as f:
        json.dump({"args": vars(args), "val_log": log, "final": final, "egitim_suresi_sn": round(train_time),
                   "n_egitim": len(train), "n_lora_modul": len(names), "oom_atlanan_adim": n_oom}, f, indent=2, ensure_ascii=False)
    print("\nyazıldı:", args.out)


if __name__ == "__main__":
    main()