m2w-gemma4-e2b-pointer / code /m2w_train_lora.py
senpaicq's picture
Add LoRA adapter, pointer head, model card, code and results
a882958 verified
Raw History Blame Contribute Delete
19.1 kB
#!/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()