Spaces:
Running
Running
| """Train one WebMix LoRA adapter on replayed browser-use rows (runs in the `webmix` env). | |
| python -m miniwebagent.train --adapter drag # family adapter: CE on its rows, KL anchor on the rest | |
| python -m miniwebagent.train --adapter pooled # one adapter on every family's rows, no anchor | |
| Rows: data/webmix/replay/<task>/rows.jsonl (miniwebgen.replay; browser-use's own prompt + screenshot -> | |
| the flash-mode JSON it answered). Each row is a family's by the macro segment it came from. | |
| Loss (WebMix's row-disjoint KL anchor, method/hyperWeb/integrations/llamafactory-kl.patch | |
| `disjoint_kl_loss`, re-implemented here because that patch targets a LLaMA-Factory fork): a row is | |
| either a CE row (the adapter's own families: mean -log p(target token)) or a KL row (other families, | |
| `--anchor-ratio` of them per CE row, redrawn every epoch: the k3 estimate exp(d) - d - 1 with | |
| d = log p_ref - log p_theta on the target's action tokens (not the templated memory, see | |
| `action_mask`), clamped to [-20, 8], weighted by --beta). The | |
| reference is the same model with the adapter disabled, so no second copy is loaded. | |
| Recipe (WebMix default): rank 16, alpha 32, dropout 0.05, attention + MLP projections of the language | |
| model only (targets.py `attn_mlp`: q/k/v/o on the full-attention layers, gate/up/down on every | |
| layer; no vision tower, no DeltaNet projections), lr 1e-4 cosine, AdamW. Batch 1 x grad-accum. | |
| Only the target positions' logits are computed (`logits_to_keep`): a 248k-token vocabulary over a | |
| ~6.6k-token prompt would not fit otherwise. | |
| Multi-GPU (granite, 4x A800, 2026-09-30): `torchrun --nproc_per_node N -m miniwebagent.train ...` gives each rank a | |
| disjoint share of every epoch's rows and sums the LoRA gradients across ranks at each optimizer step (one | |
| flat all-reduce; no DDP wrapper, so KL rows' adapter-off forwards need no special handling). --grad-accum | |
| stays the GLOBAL batch in rows, so a run's recipe does not depend on the GPU count. Rows are encoded | |
| (screenshot decode + tokenization) in background threads ahead of the GPU. | |
| """ | |
| from __future__ import annotations | |
| from helpers.paths import REPO_ROOT | |
| import argparse | |
| import hashlib | |
| import json | |
| import os | |
| import math | |
| import random | |
| import time | |
| from pathlib import Path | |
| from miniwebagent import REPLAY_DIR | |
| ROOT = REPO_ROOT | |
| ROWS_DIR = REPLAY_DIR | |
| OUT_DIR = ROOT / "data" / "webmix" / "adapters" | |
| MODEL = os.environ.get("WEBMIX_BASE_MODEL", "Qwen/Qwen3.5-4B") # e.g. microsoft/Fara1.5-4B (RQ2 init swap), Qwen/Qwen3.5-9B | |
| FAMILIES = ["Drag & gesture", "Discrete selection", "Text entry", "Form transaction", "Navigation", | |
| "Out-of-page I/O", "Reasoning base"] | |
| # the one-5090 pilot (design doc "Pilot on one 5090"): three family adapters + one pooled over the rest | |
| ADAPTERS = { | |
| "pooled": FAMILIES, | |
| "drag": ["Drag & gesture"], | |
| "io": ["Out-of-page I/O"], | |
| "select": ["Discrete selection"], | |
| "rest": ["Text entry", "Form transaction", "Navigation", "Reasoning base"], | |
| # reasoning steps and final answers (user, 2026-09-30): rows from rejection sampling on the train-site human | |
| # reasoning tasks (miniwebagent.evaluate --record-rows), KL-anchored on the synthetic families' rows | |
| "reasoning": ["Reasoning base"], | |
| # the planner (user, 2026-10-06: "lets try to train our planner"): its own rows by construction (miniwebgen.plan_data), | |
| # trained with --rows-dir on them and --beta 0 (no executor rows to anchor on) | |
| "planner": ["planner"], | |
| } | |
| TARGET_LEAVES = ["q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj"] | |
| TARGET_RE = r".*language_model.*\.(" + "|".join(TARGET_LEAVES) + r")$" | |
| def save_adapter(model, out): | |
| """PEFT writes the training regex into adapter_config.json; servers (vLLM) expect module names. | |
| The weights only cover language-model modules, so the leaf-name list is equivalent.""" | |
| model.save_pretrained(out, selected_adapters=["default"]) # not the frozen anchor copy (--init-adapter) | |
| cfg_path = Path(out) / "adapter_config.json" | |
| cfg = json.loads(cfg_path.read_text()) | |
| cfg["target_modules"] = TARGET_LEAVES | |
| cfg_path.write_text(json.dumps(cfg, indent=2)) | |
| def is_val(task_id, frac=0.05): | |
| """A fixed ~5% of trajectories (by task) held out for validation loss.""" | |
| return int(hashlib.md5(task_id.encode()).hexdigest()[:8], 16) / 0xFFFFFFFF < frac | |
| def load_rows(rows_dir=ROWS_DIR): | |
| rows = [] | |
| for p in sorted(rows_dir.glob("*/rows.jsonl")): | |
| if p.parent.name.startswith("_"): | |
| continue | |
| try: | |
| parts = json.loads((p.parent / "episode.json").read_text()).get("parts") | |
| except (OSError, ValueError): | |
| parts = None | |
| for line in p.read_text().split("\n"): | |
| if line.strip(): | |
| r = json.loads(line) | |
| if r.get("inject") or r.get("aug_prefix"): | |
| continue # a recovery episode's injected mistake, and its prefix (a copy of the plain replay) | |
| r["dir"] = str(p.parent) | |
| if parts: | |
| r["parts"] = parts | |
| rows.append(r) | |
| return rows | |
| def in_frac(r, frac): | |
| """Scaling runs (2026-10-01): keep a fixed `frac` of the generated trajectories, whole episodes at a time (a | |
| recovery variant follows its source, a chain needs all its parts); nested across fractions, separate from the | |
| validation split.""" | |
| if frac >= 1: | |
| return True | |
| key = lambda t: int(hashlib.md5(f"frac:{t}".encode()).hexdigest()[:8], 16) / 0xFFFFFFFF # noqa: E731 | |
| return all(key(t) < frac for t in (r.get("parts") or [_base_task(r["task_id"])])) | |
| def in_ids(r, ids): | |
| """--episode-ids (skill-mix ablation, 2026-10-02): keep a row when its trajectory is listed (a recovery variant follows | |
| its source, a chain needs all its parts).""" | |
| return ids is None or all(t in ids for t in (r.get("parts") or [_base_task(r["task_id"])])) | |
| def row_is_val(r): | |
| """Held out when its trajectory is (a recovery variant follows its source, a chain any of its parts).""" | |
| return any(is_val(t) for t in (r.get("parts") or [_base_task(r["task_id"])])) | |
| # ── examples ────────────────────────────────────────────────────────────────── | |
| def to_hf(row): | |
| """OpenAI-format messages (image parts point at files next to rows.jsonl) -> (HF messages, images).""" | |
| from PIL import Image | |
| msgs, imgs = [], [] | |
| for m in row["messages"]: | |
| c = m["content"] | |
| if isinstance(c, str): | |
| msgs.append({"role": m["role"], "content": [{"type": "text", "text": c}]}) | |
| continue | |
| parts = [] | |
| for x in c: | |
| if x.get("type") == "text": | |
| parts.append({"type": "text", "text": x["text"]}) | |
| elif x.get("type") == "image_url": | |
| imgs.append(Image.open(Path(row["dir"]) / x["image_url"]["url"]).convert("RGB")) | |
| parts.append({"type": "image"}) | |
| msgs.append({"role": m["role"], "content": parts}) | |
| return msgs, imgs | |
| def encode(processor, row): | |
| """-> model inputs + labels (-100 on the prompt) + n_target (supervised tokens incl. <|im_end|>).""" | |
| msgs, imgs = to_hf(row) | |
| prompt = processor.apply_chat_template(msgs, tokenize=False, add_generation_prompt=True, enable_thinking=False) | |
| full = prompt + row["target"] + "<|im_end|>" | |
| kw = {"images": imgs} if imgs else {} | |
| enc = processor(text=[full], return_tensors="pt", **kw) | |
| n_prompt = processor(text=[prompt], return_tensors="pt", **kw)["input_ids"].shape[1] | |
| labels = enc["input_ids"].clone() | |
| labels[:, :n_prompt] = -100 | |
| enc["labels"] = labels | |
| enc["n_target"] = n = enc["input_ids"].shape[1] - n_prompt | |
| enc["action_mask"] = action_mask(processor.tokenizer, row["target"], n) | |
| return enc | |
| def action_mask(tokenizer, target, n): | |
| """(n,) bool over the target tokens: True from the "action" field on (incl. <|im_end|>). The KL | |
| anchor applies there only (the memory is written by the base model itself since cycle 3, so it needs | |
| no anchor), and --ce-on action restricts the CE to it. All True when the | |
| target's standalone tokenization does not line up with its tokens in context.""" | |
| import torch | |
| t = tokenizer(target + "<|im_end|>", add_special_tokens=False, return_offsets_mapping=True) | |
| at = target.find('"action"') | |
| if len(t["input_ids"]) != n or at < 0: | |
| return torch.ones(n, dtype=torch.bool) | |
| return torch.tensor([end > at for _, end in t["offset_mapping"]], dtype=torch.bool) | |
| def target_logprobs(model, enc, device): | |
| """log p(target token) at each supervised position, computing logits only for those positions.""" | |
| import torch | |
| n = enc["n_target"] | |
| inputs = {k: v.to(device) for k, v in enc.items() if k not in ("labels", "n_target", "action_mask")} | |
| out = model(**inputs, logits_to_keep=n + 1, use_cache=False) | |
| logits = out.logits[:, :-1, :].float() # positions predicting the n target tokens | |
| tgt = enc["labels"][:, -n:].to(device) | |
| return torch.log_softmax(logits, dim=-1).gather(-1, tgt.unsqueeze(-1)).squeeze(-1) # (1, n) | |
| def prefetched(processor, items, workers=4, ahead=8): | |
| """Yield (item, encode(row)) in order, encoding up to `ahead` rows in background threads.""" | |
| from collections import deque | |
| from concurrent.futures import ThreadPoolExecutor | |
| with ThreadPoolExecutor(workers) as ex: | |
| q, it = deque(), iter(items) | |
| for x in it: | |
| q.append((x, ex.submit(encode, processor, x[0]))) | |
| if len(q) >= ahead: | |
| break | |
| while q: | |
| x, fut = q.popleft() | |
| nxt = next(it, None) | |
| if nxt is not None: | |
| q.append((nxt, ex.submit(encode, processor, nxt[0]))) | |
| yield x, fut.result() | |
| def _dist(): | |
| """-> (rank, world, local_rank); initialises NCCL under torchrun.""" | |
| import os | |
| world = int(os.environ.get("WORLD_SIZE", "1")) | |
| if world == 1: | |
| return 0, 1, 0 | |
| import torch | |
| import torch.distributed as dist | |
| local = int(os.environ["LOCAL_RANK"]) | |
| torch.cuda.set_device(local) | |
| from datetime import timedelta | |
| dist.init_process_group("nccl", timeout=timedelta(hours=1)) # rank 0 evaluates while the others wait | |
| return dist.get_rank(), world, local | |
| def _allreduce_grads(params, world): | |
| import torch | |
| import torch.distributed as dist | |
| grads = [p.grad if p.grad is not None else torch.zeros_like(p) for p in params] | |
| flat = torch.cat([g.reshape(-1) for g in grads]) | |
| dist.all_reduce(flat) | |
| flat /= world | |
| i = 0 | |
| for p, g in zip(params, grads): | |
| n = g.numel() | |
| p.grad = flat[i:i + n].view_as(p) | |
| i += n | |
| # ── training ────────────────────────────────────────────────────────────────── | |
| def main(argv=None): | |
| ap = argparse.ArgumentParser(prog="python -m miniwebagent.train") | |
| ap.add_argument("--adapter", required=True, choices=sorted(ADAPTERS)) | |
| ap.add_argument("--epochs", type=float, default=3) | |
| ap.add_argument("--lr", type=float, default=1e-4) | |
| ap.add_argument("--rank", type=int, default=16) | |
| ap.add_argument("--alpha", type=int, default=32) | |
| ap.add_argument("--grad-accum", type=int, default=8) | |
| ap.add_argument("--beta", type=float, default=0.1, help="KL anchor weight (0 = no anchor rows)") | |
| ap.add_argument("--anchor-ratio", type=float, default=0.5, help="KL rows per CE row") | |
| ap.add_argument("--limit", type=int, default=0, help="cap CE rows (smoke tests)") | |
| ap.add_argument("--max-steps", type=int, default=0, help="stop after this many rows (smoke tests)") | |
| ap.add_argument("--out", default=None) | |
| ap.add_argument("--rows-dir", default=None, help="rows to train on (default: the executor replay, REPLAY_DIR)") | |
| ap.add_argument("--seed", type=int, default=0) | |
| ap.add_argument("--wandb-project", default="webmix-pilot", help="wandb project ('' = no wandb)") | |
| ap.add_argument("--ce-on", choices=["all", "action"], default="all", | |
| help="CE on the whole target (memory + action) or on the action tokens only") | |
| ap.add_argument("--exclude-variants", default="", help="comma list of replay variants to leave out (recovery,chain)") | |
| ap.add_argument("--frac", type=float, default=1.0, | |
| help="scaling: train on this fraction of the generated episodes (whole episodes, nested)") | |
| ap.add_argument("--episode-ids", default=None, | |
| help="train only on the trajectories listed in this file (one task id per line; variants follow)") | |
| ap.add_argument("--init-adapter", default=None, | |
| help="branch-then-specialize (2026-10-01): start from this adapter's weights (same rank/alpha) and " | |
| "anchor the KL rows to it (a frozen copy) instead of the base model") | |
| ap.add_argument("--drop-done", action="store_true", | |
| help="review round 1 control (2026-10-05): leave out every episode's final `done` step, so practice does " | |
| "not teach the single agent to stop after one primitive (the `macro_done` before it stays)") | |
| ap.add_argument("--save-every", type=int, default=5000, | |
| help="also save the adapter every N rows (out/rows_<n>; 0 = epoch boundaries only)") | |
| args = ap.parse_args(argv) | |
| import torch | |
| from peft import LoraConfig, get_peft_model | |
| from transformers import AutoModelForImageTextToText, AutoProcessor | |
| rank, world, local = _dist() | |
| main0 = rank == 0 | |
| random.seed(args.seed) | |
| torch.manual_seed(args.seed) | |
| fams = set(ADAPTERS[args.adapter]) | |
| rows = load_rows(Path(args.rows_dir)) if args.rows_dir else load_rows() | |
| if args.exclude_variants: | |
| drop = set(args.exclude_variants.split(",")) | |
| rows = [r for r in rows if _variant(r["task_id"]) not in drop] | |
| if args.drop_done: | |
| n0 = len(rows) | |
| rows = [r for r in rows if r.get("op") != "done"] | |
| if main0: | |
| print(f"--drop-done: left out {n0 - len(rows)} final done steps", flush=True) | |
| ids = set(Path(args.episode_ids).read_text().split()) if args.episode_ids else None | |
| train = [r for r in rows if not row_is_val(r) and in_frac(r, args.frac) and in_ids(r, ids)] | |
| val = [r for r in rows if row_is_val(r)] | |
| ce_rows = [r for r in train if r["family"] in fams] | |
| kl_pool = [r for r in train if r["family"] not in fams] if args.beta > 0 else [] | |
| if args.limit: | |
| random.shuffle(ce_rows) | |
| ce_rows = ce_rows[:args.limit] | |
| out = Path(args.out) if args.out else OUT_DIR / args.adapter | |
| out.mkdir(parents=True, exist_ok=True) | |
| if main0: | |
| print(f"adapter {args.adapter}: {len(ce_rows)} CE rows ({sorted(fams)}), anchor pool {len(kl_pool)}, " | |
| f"val {len(val)}, {world} GPU(s)", flush=True) | |
| device = f"cuda:{local}" | |
| processor = AutoProcessor.from_pretrained(MODEL) | |
| model = AutoModelForImageTextToText.from_pretrained(MODEL, dtype=torch.bfloat16, device_map=device, | |
| attn_implementation="sdpa") | |
| model.config.use_cache = False | |
| model.gradient_checkpointing_enable(gradient_checkpointing_kwargs={"use_reentrant": False}) | |
| model.enable_input_require_grads() | |
| model = get_peft_model(model, LoraConfig(r=args.rank, lora_alpha=args.alpha, lora_dropout=0.05, bias="none", | |
| target_modules=TARGET_RE, task_type="CAUSAL_LM")) | |
| anchor = None | |
| if args.init_adapter: | |
| # specialists trained from base on their family's rows alone lost the general habits (scrolling, finding, | |
| # navigating) pooled_v4 learned from every family; branching from pooled_v4 keeps them | |
| from peft import set_peft_model_state_dict | |
| from safetensors.torch import load_file | |
| sd = load_file(str(Path(args.init_adapter) / "adapter_model.safetensors")) | |
| res = set_peft_model_state_dict(model, sd) | |
| missing = [k for k in getattr(res, "missing_keys", []) if "lora_" in k] | |
| assert not missing, f"--init-adapter left {len(missing)} LoRA weights unset, e.g. {missing[:2]}" | |
| if args.beta > 0: | |
| model.load_adapter(args.init_adapter, adapter_name="anchor", is_trainable=False) | |
| model.set_adapter("default") | |
| anchor = "anchor" | |
| if main0: | |
| print(f"initialized from {args.init_adapter} ({len(sd)} tensors); KL anchor: " | |
| f"{'that adapter (frozen)' if anchor else 'none'}", flush=True) | |
| if main0: | |
| model.print_trainable_parameters() | |
| n_kl = int(round(len(ce_rows) * args.anchor_ratio)) if kl_pool else 0 | |
| per_epoch = (len(ce_rows) + n_kl) // world * world # every rank takes the same number of rows | |
| total = int(math.ceil(per_epoch * args.epochs)) // world * world | |
| if args.max_steps: | |
| total = min(total, args.max_steps // world * world) | |
| accum = max(1, args.grad_accum // world) # rows per rank per optimizer step | |
| opt_steps = max(1, total // (accum * world)) | |
| params = [p for p in model.parameters() if p.requires_grad] | |
| opt = torch.optim.AdamW(params, lr=args.lr, weight_decay=0.0) | |
| warm = max(1, int(0.03 * opt_steps)) | |
| sched = torch.optim.lr_scheduler.LambdaLR( | |
| opt, lambda s: min(1.0, (s + 1) / warm) * 0.5 * (1 + math.cos(math.pi * min(1.0, s / opt_steps)))) | |
| log = open(out / "train_log.jsonl", "a") if main0 else None | |
| val_fam = [r for r in val if r["family"] in fams] | |
| random.Random(1).shuffle(val_fam) # a fixed sample across macros (the file order is alphabetical by macro) | |
| wb = None | |
| if args.wandb_project and main0: | |
| import wandb | |
| wb = wandb.init(project=args.wandb_project, name=f"{out.name}-{time.strftime('%m%d-%H%M')}", dir=str(out), | |
| config={**vars(args), "families": sorted(fams), "base_model": MODEL, "ce_rows": len(ce_rows), | |
| "anchor_pool": len(kl_pool), "anchor_rows_per_epoch": n_kl, "val_rows": len(val_fam), | |
| "rows_total": total, "optimizer_steps": opt_steps, "gpus": world}) | |
| model.train() | |
| done, t0, tokens, acc = 0, time.time(), 0, {"ce": [], "kl": []} # done = rows over all ranks | |
| saved_at = 0 | |
| epoch = 0 | |
| while done < total: | |
| batch = [(r, "ce") for r in ce_rows] + [(r, "kl") for r in random.sample(kl_pool, min(n_kl, len(kl_pool)))] | |
| random.shuffle(batch) # the same order on every rank (same seed, same calls) | |
| batch = batch[:len(batch) // world * world][rank::world] | |
| batch = batch[:max(0, (total - done) // world)] | |
| k = 0 | |
| for (r, kind), enc in prefetched(processor, batch): | |
| logp = target_logprobs(model, enc, device) | |
| if kind == "ce": | |
| if args.ce_on == "action": | |
| m = enc["action_mask"].to(device).unsqueeze(0).float() | |
| loss = -(logp * m).sum() / m.sum().clamp(min=1.0) | |
| else: | |
| loss = -logp.mean() | |
| else: | |
| with torch.no_grad(): | |
| if anchor: # the initial adapter, frozen (branch-then-specialize) | |
| model.set_adapter(anchor) | |
| ref = target_logprobs(model, enc, device) | |
| model.set_adapter("default") | |
| else: | |
| with model.disable_adapter(): | |
| ref = target_logprobs(model, enc, device) | |
| d = (ref - logp).clamp(-20.0, 8.0) | |
| m = enc["action_mask"].to(device).unsqueeze(0).float() | |
| loss = args.beta * ((torch.exp(d) - d - 1.0) * m).sum() / m.sum().clamp(min=1.0) | |
| (loss / accum).backward() | |
| acc[kind].append(loss.item() / (args.beta if kind == "kl" else 1.0)) | |
| tokens += enc["input_ids"].shape[1] * world | |
| done += world | |
| k += 1 | |
| if k % accum == 0 or done >= total: | |
| if world > 1: | |
| _allreduce_grads(params, world) | |
| torch.nn.utils.clip_grad_norm_(params, 1.0) | |
| opt.step() | |
| sched.step() | |
| opt.zero_grad(set_to_none=True) | |
| # a mid-run checkpoint (a crash loses at most N rows; learning curves come for free) | |
| if main0 and args.save_every and done < total and done // args.save_every > saved_at // args.save_every: | |
| save_adapter(model, out / f"rows_{done}") | |
| saved_at = done | |
| if main0 and (done % (50 * world) == 0 or done >= total): | |
| el = time.time() - t0 | |
| rec = {"rows": done, "of": total, "epoch": epoch, | |
| "ce": sum(acc["ce"]) / max(1, len(acc["ce"])), "kl": sum(acc["kl"]) / max(1, len(acc["kl"])), | |
| "lr": sched.get_last_lr()[0], "rows_per_s": done / el, "tokens_per_s": tokens / el, | |
| "gpu_gb": torch.cuda.max_memory_allocated() / 2**30} | |
| log.write(json.dumps(rec) + "\n") | |
| log.flush() | |
| print(json.dumps(rec), flush=True) | |
| if wb: | |
| wb.log(rec, step=done) | |
| acc = {"ce": [], "kl": []} | |
| epoch += 1 | |
| if done < total: # a checkpoint + held-out CE at each epoch boundary | |
| if main0: | |
| save_adapter(model, out / f"epoch_{epoch}") | |
| v = evaluate(model, processor, val_fam[:100], device) | |
| print(json.dumps({"epoch": epoch, "val_ce": v}), flush=True) | |
| if wb: | |
| wb.log({"val_ce": v, "epoch": epoch}, step=done) | |
| _barrier(world) | |
| if not main0: | |
| _barrier(world) | |
| return | |
| save_adapter(model, out) | |
| val_loss = evaluate(model, processor, val_fam[:200], device) | |
| meta = {"adapter": args.adapter, "families": sorted(fams), "base_model": MODEL, "ce_rows": len(ce_rows), | |
| "anchor_pool": len(kl_pool), "rows_trained": done, "val_ce": val_loss, "args": vars(args), | |
| "minutes": round((time.time() - t0) / 60, 1)} | |
| (out / "webmix_meta.json").write_text(json.dumps(meta, indent=1)) | |
| print(json.dumps(meta), flush=True) | |
| if wb: | |
| wb.log({"val_ce": val_loss, "epoch": epoch}, step=done) | |
| wb.summary.update({"val_ce_final": val_loss, "minutes": meta["minutes"]}) | |
| wb.finish() | |
| _barrier(world) | |
| def _barrier(world): | |
| if world > 1: | |
| import torch.distributed as dist | |
| dist.barrier() | |
| def _base_task(task_id): | |
| """A variant's held-out split follows its (first) source trajectory, so no trajectory leaks across it.""" | |
| return task_id.split("~")[0] | |
| def _variant(task_id): | |
| return "recovery" if task_id.endswith("~rec") else "chain" if task_id.startswith("chain-") else "plain" | |
| def evaluate(model, processor, rows, device): | |
| """Mean target-token CE on held-out rows of the adapter's families (None without any).""" | |
| import torch | |
| if not rows: | |
| return None | |
| model.eval() | |
| tot = 0.0 | |
| with torch.no_grad(): | |
| for r in rows: | |
| tot += -target_logprobs(model, encode(processor, r), device).mean().item() | |
| model.train() | |
| return tot / len(rows) | |
| if __name__ == "__main__": | |
| main() | |