#!/usr/bin/env python3 """ColQwen3.5 late-interaction eval: ViDoRe V1/V2/V3 + JinaVDR (paired + BEIR).""" from __future__ import annotations import argparse import glob import hashlib import io import json import math import os import time from collections import Counter, defaultdict from concurrent.futures import ThreadPoolExecutor from datetime import timedelta from pathlib import Path import pyarrow.parquet as pq import torch import torch.nn.functional as F from PIL import Image from peft import PeftModel from torch.utils.data import DataLoader, IterableDataset from transformers.models.qwen3_5 import Qwen3_5Config from colpali_engine.models import ColQwen3_5, ColQwen3_5Processor from colpali_engine.models.qwen3_5.colqwen3_5.modeling_colqwen3_5 import set_active_head from colpali_engine.utils.maxsim import maxsim_inbatch from paths import forbid_venv_path # Bump when paired/BEIR/query encoding rules change so stale task caches are ignored. EVAL_PROTOCOL = "paired-all-pages-dedup+process_queries+ndcg2r-20260827" REQUIRED_KS = (1, 5, 10) REQUIRED_METRICS = tuple( f"{metric}@{k}" for metric in ("recall", "ndcg", "mrr", "map") for k in REQUIRED_KS ) # --------------------------------------------------------------------------- # # helpers # --------------------------------------------------------------------------- # def task_tag(task_or_rec): """Stable task id: ``dataset/subset`` or ``dataset/subset[lang]``.""" if "task" in task_or_rec and isinstance(task_or_rec["task"], str): return task_or_rec["task"] tag = f"{task_or_rec['dataset']}/{task_or_rec['subset']}" lang = task_or_rec.get("lang") return f"{tag}[{lang}]" if lang else tag def task_slug(tag: str) -> str: """Filesystem-safe name for per-task result files.""" return tag.replace("/", "__").replace("[", "__").replace("]", "") def tasks_dir(out_dir: Path) -> Path: return Path(out_dir) / "tasks" def _has_complete_metrics(metrics) -> bool: if not isinstance(metrics, dict): return False return all( isinstance(metrics.get(name), (int, float)) and math.isfinite(float(metrics[name])) for name in REQUIRED_METRICS ) def _is_ok_task_rec( rec, protocol: str | None = EVAL_PROTOCOL, contract_id: str | None = None, ) -> bool: if not isinstance(rec, dict) or "error" in rec or "metrics" not in rec: return False if not _has_complete_metrics(rec["metrics"]): return False if protocol is None: protocol_ok = True else: protocol_ok = rec.get("eval_protocol") == protocol contract_ok = contract_id is None or rec.get("eval_contract_id") == contract_id return protocol_ok and contract_ok def save_task_result( out_dir: Path, rec: dict, run_name: str, contract_id: str, ) -> Path: """Persist one successful task under ``eval/tasks/`` (resume unit).""" d = tasks_dir(out_dir) d.mkdir(parents=True, exist_ok=True) payload = dict(rec) payload["eval_protocol"] = EVAL_PROTOCOL payload["run_name"] = run_name payload["eval_contract_id"] = contract_id path = d / f"{task_slug(payload['task'])}.json" tmp = path.with_suffix(".json.tmp") tmp.write_text(json.dumps(payload, ensure_ascii=False, indent=2), encoding="utf-8") tmp.replace(path) return path def load_cached_task_map( out_dir: Path, protocol: str = EVAL_PROTOCOL, contract_id: str | None = None, ) -> dict[str, dict]: """Load successful per-task results keyed by task tag.""" d = tasks_dir(out_dir) out = {} if d.is_dir(): for p in sorted(d.glob("*.json")): try: rec = json.loads(p.read_text(encoding="utf-8")) except Exception: continue if _is_ok_task_rec(rec, protocol=protocol, contract_id=contract_id): out[rec["task"]] = rec return out def to_pil(x): if isinstance(x, Image.Image): return x if isinstance(x, (bytes, bytearray)): return Image.open(io.BytesIO(x)) if isinstance(x, dict) and "bytes" in x: return Image.open(io.BytesIO(x["bytes"])) raise TypeError(f"cannot decode image of type {type(x)}") def read_parquet_rows(paths): """Read small text-only parquet files (queries / qrels).""" rows = [] for p in paths: rows.extend(pq.read_table(p).to_pylist()) return rows # --------------------------------------------------------------------------- # # task discovery # --------------------------------------------------------------------------- # def _split_files(d): """Eval split only: some subsets ship train-*.parquet alongside test-*.parquet.""" files = sorted(glob.glob(str(Path(d) / "*.parquet"))) test = [f for f in files if Path(f).name.startswith("test-")] return test or files def discover_tasks(eval_root: str): tasks = [] root = Path(eval_root) for dataset in ("eval", "eval_v2", "eval_v3", "JinaVDR", "demo"): ddir = root / dataset if not ddir.is_dir(): continue for sub in sorted(p for p in ddir.iterdir() if p.is_dir()): corpus_dirs = sorted(glob.glob(str(sub / "*-corpus"))) if corpus_dirs: # BEIR multi-language (vidore v3) # V3's six languages share the same page corpus. Canonicalize to # English (the frozen reproduce.py protocol) so it is encoded once. canonical_dir = sub / "english-corpus" if not canonical_dir.is_dir(): canonical_dir = Path(corpus_dirs[0]) canonical_corpus = _split_files(canonical_dir) for cd in corpus_dirs: lang = Path(cd).name.replace("-corpus", "") qd = sub / f"{lang}-queries" qrd = sub / f"{lang}-qrels" if qd.is_dir() and qrd.is_dir(): tasks.append( dict( dataset=dataset, subset=sub.name, fmt="beir", lang=lang, corpus=canonical_corpus, queries=_split_files(qd), qrels=_split_files(qrd), ) ) continue if (sub / "corpus").is_dir(): # BEIR single-language (vidore v2) tasks.append( dict( dataset=dataset, subset=sub.name, fmt="beir", lang=None, corpus=_split_files(sub / "corpus"), queries=_split_files(sub / "queries"), qrels=_split_files(sub / "qrels") or _split_files(sub / "docs"), ) ) continue paired = sorted(glob.glob(str(sub / "data" / "test-*.parquet"))) if paired: # PAIRED (vidore v1, JinaVDR) tasks.append(dict(dataset=dataset, subset=sub.name, fmt="paired", lang=None, paired=paired)) continue for ld in sorted(p for p in sub.iterdir() if p.is_dir() and not p.name.startswith(".")): lang_files = sorted(glob.glob(str(ld / "test-*.parquet"))) if lang_files: # PAIRED per language (JinaVDR multilingual subsets) tasks.append( dict(dataset=dataset, subset=sub.name, fmt="paired", lang=ld.name, paired=lang_files) ) return tasks def _pad_cat(embs): """Right-pad variable-length (B, L, D) chunks and concat on batch dim.""" if not embs: return torch.zeros((0, 0, 0)) Lmax = max(e.shape[1] for e in embs) embs = [F.pad(e, (0, 0, 0, Lmax - e.shape[1])) for e in embs] return torch.cat(embs, 0) class _ImageParquetDataset(IterableDataset): """Stream parquet image rows; decode to RGB PIL in DataLoader workers.""" def __init__( self, paths, meta_keys=None, skip_empty_query=False, dedupe_key=None, ): self.paths = paths self.meta_keys = meta_keys or [] self.skip_empty_query = skip_empty_query self.dedupe_key = dedupe_key def __iter__(self): info = torch.utils.data.get_worker_info() wid, nw = (0, 1) if info is None else (info.id, info.num_workers) idx = 0 seen = set() want_cols = ["image"] + list(self.meta_keys) if self.skip_empty_query and "query" not in want_cols: want_cols.append("query") if self.dedupe_key and self.dedupe_key not in want_cols: want_cols.append(self.dedupe_key) for p in self.paths: pf = pq.ParquetFile(p) avail = set(pf.schema_arrow.names) cols = [c for c in want_cols if c in avail] has_query = "query" in cols for rg in range(pf.num_row_groups): tbl = pf.read_row_group(rg, columns=cols) img_col = tbl.column("image") meta_cols = {k: tbl.column(k) for k in self.meta_keys if k in avail} q_col = tbl.column("query") if has_query else None for j in range(tbl.num_rows): if self.skip_empty_query and q_col is not None: q = q_col[j].as_py() q = "" if q is None else str(q).strip() if not q or q.lower() == "none": continue dedupe_value = ( tbl.column(self.dedupe_key)[j].as_py() if self.dedupe_key and self.dedupe_key in avail else None ) if dedupe_value is not None: page_key = str(dedupe_value) owner = int.from_bytes( hashlib.blake2b( page_key.encode("utf-8"), digest_size=8 ).digest(), "little", ) % nw if owner != wid or page_key in seen: continue seen.add(page_key) elif idx % nw != wid: idx += 1 continue idx += 1 img = to_pil(img_col[j].as_py()).convert("RGB") meta = {k: meta_cols[k][j].as_py() if k in meta_cols else None for k in self.meta_keys} yield (img, meta) def _collate_pil(items): return [x[0] for x in items], [x[1] for x in items] @torch.no_grad() def embed_image_stream( processor, model, paths, device, batch_size, num_workers, meta_keys=None, max_docs=0, skip_empty_query=False, dedupe_key=None, ): """Embed images with CPU preprocess overlapped against GPU forward.""" ds = _ImageParquetDataset( paths, meta_keys=meta_keys, skip_empty_query=skip_empty_query, dedupe_key=dedupe_key, ) loader = DataLoader( ds, batch_size=batch_size, num_workers=num_workers, collate_fn=_collate_pil, prefetch_factor=6 if num_workers > 0 else None, persistent_workers=False, ) out, metas, total = [], [], 0 with ThreadPoolExecutor(max_workers=1) as pool: pending = None it = iter(loader) def submit_next(): try: imgs, bm = next(it) except StopIteration: return None return (pool.submit(processor.process_images, imgs), bm, len(imgs)) pending = submit_next() while pending is not None: fut, bm, n = pending batch = fut.result() nxt = submit_next() batch = {k: v.to(device, non_blocking=True) for k, v in batch.items()} out.append(model(**batch).to("cpu")) metas.extend(bm) total += n if max_docs and total >= max_docs: break pending = nxt doc_emb = _pad_cat(out) if max_docs and doc_emb.shape[0] > max_docs: doc_emb = doc_emb[:max_docs] metas = metas[:max_docs] return doc_emb, metas @torch.no_grad() def embed_queries(processor, model, texts, device, batch_size): """Encode queries with ColQwen query prefix + 10 augmentation tokens. Must use ``process_queries`` (not ``process_texts``): training collator and standard ViDoRe eval append ``query_augmentation_token * 10``. """ out = [] for i in range(0, len(texts), batch_size): chunk = texts[i : i + batch_size] chunk = [t if (t and t.strip()) else " " for t in chunk] batch = processor.process_queries(chunk) batch = {k: v.to(device) for k, v in batch.items()} # Clear rope cache before query forward (ColQwen hybrid-attn note). for obj in (model, getattr(model, "get_base_model", lambda: None)()): if obj is not None and hasattr(obj, "rope_deltas"): obj.rope_deltas = None out.append(model(**batch)) return _pad_cat(out) def chunked_maxsim(q_emb, doc_emb, chunk_q=64, chunk_d=256): """(Nq, Nd) MaxSim. Docs may live on CPU; each chunk is moved to the query device once.""" dev = q_emb.device Nq, Nd = q_emb.shape[0], doc_emb.shape[0] scores = torch.zeros(Nq, Nd, device=dev, dtype=torch.float32) for di in range(0, Nd, chunk_d): d = doc_emb[di : di + chunk_d].to(dev, non_blocking=True).contiguous() for qi in range(0, Nq, chunk_q): q = q_emb[qi : qi + chunk_q].contiguous() scores[qi : qi + chunk_q, di : di + chunk_d] = maxsim_inbatch(q, d) return scores METRIC_NAMES = ("recall", "ndcg", "mrr", "map") def metrics_from_scores(scores, relevant, graded, ks): acc = {f"{m}@{k}": [] for k in ks for m in METRIC_NAMES} for q in range(scores.shape[0]): rel_set = relevant[q] if not rel_set: continue rel_grad = graded[q] order = scores[q].argsort(descending=True).tolist() for k in ks: topk = order[:k] acc[f"recall@{k}"].append(len(set(topk) & rel_set) / len(rel_set)) dcg = sum( ((2 ** rel_grad.get(c, 0.0)) - 1) / math.log2(i + 2) for i, c in enumerate(topk) if rel_grad.get(c, 0.0) > 0 ) ideal = sorted( ((2 ** g) - 1 for g in rel_grad.values() if g > 0), reverse=True )[:k] idcg = sum(g / math.log2(i + 2) for i, g in enumerate(ideal)) acc[f"ndcg@{k}"].append(dcg / idcg if idcg > 0 else 0.0) acc[f"mrr@{k}"].append( next((1.0 / (i + 1) for i, c in enumerate(topk) if c in rel_set), 0.0) ) hits, ap = 0, 0.0 for i, c in enumerate(topk): if c in rel_set: hits += 1 ap += hits / (i + 1) denom = min(len(rel_set), k) acc[f"map@{k}"].append(ap / denom if denom else 0.0) return {name: (sum(v) / len(v) if v else 0.0) for name, v in acc.items()} # --------------------------------------------------------------------------- # # runners # --------------------------------------------------------------------------- # def _paired_query_rows(paths, page_to_idx): """Read paired query metadata without loading the image column.""" queries, golds = [], [] for path in paths: pf = pq.ParquetFile(path) required = {"query", "image_filename"} missing = required - set(pf.schema_arrow.names) if missing: raise ValueError(f"paired parquet missing {sorted(missing)}: {path}") for rg in range(pf.num_row_groups): table = pf.read_row_group(rg, columns=["query", "image_filename"]) for query, page in zip( table.column("query").to_pylist(), table.column("image_filename").to_pylist(), ): text = "" if query is None else str(query).strip() if not text or text.lower() == "none": continue doc = page_to_idx.get(str(page)) if doc is not None: queries.append(text) golds.append(doc) return queries, golds def run_paired(task, processor, model, device, ks, embed_bs, num_workers, max_queries, max_docs): """ViDoRe V1 / JinaVDR all-page candidate protocol. Empty-query rows do not become evaluation queries, but their unique pages remain retrieval candidates. Duplicate page images are encoded only once. """ doc_emb, metas = embed_image_stream( processor, model, task["paired"], device, embed_bs, num_workers, meta_keys=["image_filename"], max_docs=max_docs, dedupe_key="image_filename", ) if not metas: raise ValueError(f"no candidate pages for paired subset {task['subset']}") page_to_idx = {} for i, m in enumerate(metas): key = m.get("image_filename") if key is None: raise ValueError( f"paired candidate missing image_filename in {task['subset']}" ) key = str(key) if key in page_to_idx: raise ValueError(f"duplicate page was encoded twice: {key}") page_to_idx[key] = i queries, golds = _paired_query_rows(task["paired"], page_to_idx) if not queries: raise ValueError(f"no non-empty queries for paired subset {task['subset']}") if max_queries and max_queries < len(queries): queries = queries[:max_queries] golds = golds[:max_queries] q_emb = embed_queries(processor, model, queries, device, embed_bs) scores = chunked_maxsim(q_emb, doc_emb) relevant = [{g} for g in golds] graded = [{g: 1.0} for g in golds] if scores.shape[0] <= 16 and scores.shape[1] <= 16: idx_to_page = {i: k for k, i in page_to_idx.items()} print(f"[rank] {task_tag(task)}") for qi, query in enumerate(queries): order = scores[qi].argsort(descending=True).tolist() gold_page = idx_to_page[golds[qi]] hit = "HIT" if order[0] == golds[qi] else "MISS" top = [ f"{idx_to_page[j]}={float(scores[qi, j]):.3f}" for j in order[:3] ] print(f"[rank] {hit} gold={gold_page} {query}") print(f"[rank] top {', '.join(top)}") return ( metrics_from_scores(scores, relevant, graded, ks), dict(n_docs=doc_emb.shape[0], n_queries=q_emb.shape[0]), ) def run_beir( task, processor, model, device, ks, embed_bs, num_workers, max_queries, max_docs, corpus_cache=None, ): corpus_key = (tuple(task["corpus"]), int(max_docs)) if corpus_cache is not None and corpus_cache.get("key") == corpus_key: doc_emb, metas = corpus_cache["value"] print(f"[eval] reuse encoded corpus for {task_tag(task)} ({doc_emb.shape[0]} docs)") else: doc_emb, metas = embed_image_stream( processor, model, task["corpus"], device, embed_bs, num_workers, meta_keys=["corpus-id", "id"], max_docs=max_docs, ) if corpus_cache is not None: corpus_cache.clear() corpus_cache.update(key=corpus_key, value=(doc_emb, metas)) def _first(m, *keys, default): for k in keys: v = m.get(k) if v is not None: return v return default corpus_ids = [str(_first(m, "corpus-id", "id", default=f"c{i}")) for i, m in enumerate(metas)] cid_to_idx = {cid: i for i, cid in enumerate(corpus_ids)} if len(cid_to_idx) != len(corpus_ids): raise ValueError(f"duplicate corpus ids in {task['subset']}") q_rows = read_parquet_rows(task["queries"]) q_texts = [str(r["text"] if "text" in r else r["query"]) for r in q_rows] q_ids = [str(r.get("id", r.get("query-id", f"q{i}"))) for i, r in enumerate(q_rows)] rel_by_q, grad_by_q = defaultdict(set), defaultdict(dict) for r in read_parquet_rows(task["qrels"]): qid = str(r.get("query-id", r.get("id", ""))) cid = str(r.get("corpus-id", r.get("id", ""))) s = float(r.get("score", 1.0) or 1.0) if s > 0 and cid in cid_to_idx: rel_by_q[qid].add(cid_to_idx[cid]) grad_by_q[qid][cid_to_idx[cid]] = s keep = [ i for i in range(len(q_ids)) if q_texts[i].strip() and q_texts[i].lower() != "none" and q_ids[i] in rel_by_q and rel_by_q[q_ids[i]] ] if max_queries and max_queries < len(keep): keep = keep[:max_queries] if not keep: raise ValueError(f"no non-empty queries with valid qrels for {task['subset']}") sub_emb = embed_queries(processor, model, [q_texts[i] for i in keep], device, embed_bs) sub_ids = [q_ids[i] for i in keep] relevant = [rel_by_q[q] for q in sub_ids] graded = [grad_by_q[q] for q in sub_ids] scores = chunked_maxsim(sub_emb, doc_emb) return ( metrics_from_scores(scores, relevant, graded, ks), dict(n_docs=doc_emb.shape[0], n_queries=len(sub_ids)), ) # --------------------------------------------------------------------------- # # main # --------------------------------------------------------------------------- # def parse_args(): ap = argparse.ArgumentParser(description=__doc__) ap.add_argument("--base-model", default="") ap.add_argument("--adapter-dir", default=None) ap.add_argument( "--skip-adapter-weights", action="store_true", help="keep adapter-dir for run_config/provenance but load only --base-model", ) ap.add_argument("--eval-root", default="") ap.add_argument("--output-dir", required=True) ap.add_argument("--max-visual-tokens", type=int, default=None) ap.add_argument("--embed-batch", type=int, default=16) ap.add_argument("--ks", default="1,5,10") ap.add_argument("--max-queries", type=int, default=0) ap.add_argument("--max-docs", type=int, default=0) ap.add_argument("--datasets", default="") ap.add_argument("--num-workers", type=int, default=8) ap.add_argument("--run-name", default="eval") ap.add_argument("--attn", default="flash_attention_2") ap.add_argument("--bidirectional-attention", choices=("auto", "on", "off"), default="auto") ap.add_argument( "--head-dim", type=int, default=None, help=" Matryoshka: which head to score with. Required when the run trained " "several heads, since a multi-head forward returns a dict, not one tensor. " "Must be one of the checkpoint head_dims.", ) ap.add_argument("--allow-config-override", action="store_true") ap.add_argument( "--overwrite-output", action="store_true", help="Archive existing eval dir and start clean (ignores --resume).", ) ap.add_argument( "--resume", action="store_true", help="Keep good tasks under eval/tasks/; only run missing ones.", ) return ap.parse_args() def resolve_eval_contract(args): """Resolve and validate evaluation settings against the training manifest.""" config = {} config_path = None if args.adapter_dir: candidate = Path(args.adapter_dir) / "run_config.json" if candidate.is_file(): config_path = candidate config = json.loads(candidate.read_text(encoding="utf-8")) model_dir = Path(args.adapter_dir or args.base_model or ".") model_cfg = _read_json(model_dir / "config.json") if not config.get("head_dims") and model_cfg.get("head_dims"): config["head_dims"] = model_cfg["head_dims"] trained_mvt = config.get("max_visual_tokens") trained_bidir = config.get("bidirectional_attention") if args.max_visual_tokens is None: args.max_visual_tokens = int(trained_mvt) if trained_mvt is not None else 1024 elif trained_mvt is not None and int(args.max_visual_tokens) != int(trained_mvt): if not args.allow_config_override: raise ValueError( f"MVT mismatch: eval={args.max_visual_tokens}, trained={trained_mvt}; " "pass --allow-config-override to force" ) if args.bidirectional_attention == "auto": args.bidirectional_attention = ( trained_bidir if trained_bidir in ("on", "off") else "on" ) elif trained_bidir in ("on", "off") and args.bidirectional_attention != trained_bidir: if not args.allow_config_override: raise ValueError( f"attention mismatch: eval={args.bidirectional_attention}, " f"trained={trained_bidir}; pass --allow-config-override to force" ) # guess: refuse rather than silently evaluate an arbitrary width. trained_heads = config.get("head_dims") if trained_heads: if args.head_dim is None: raise ValueError( f"this run trained heads {trained_heads}; pass --head-dim to choose one " "(eval_run.sh loops over all of them)" ) if int(args.head_dim) not in [int(d) for d in trained_heads]: raise ValueError( f"--head-dim {args.head_dim} is not among the trained heads {trained_heads}" ) elif args.head_dim is not None: raise ValueError( "--head-dim given but this run is single-head " f"(col_dim={config.get('col_dim') or model_cfg.get('dim')}); drop --head-dim" ) return config_path, config def _read_json(path): p = Path(path) if not p.is_file(): return {} try: return json.loads(p.read_text(encoding="utf-8")) except Exception: return {} def _file_identity(path, hash_content=False): p = Path(path) stat = p.stat() rec = { "path": str(p.resolve()), "size": stat.st_size, "mtime_ns": stat.st_mtime_ns, } if hash_content: digest = hashlib.sha256() with p.open("rb") as handle: for chunk in iter(lambda: handle.read(8 * 1024 * 1024), b""): digest.update(chunk) rec["sha256"] = digest.hexdigest() return rec def build_eval_contract(args, train_config, tasks, ks): """Fingerprint model, data and scoring settings used by resumable tasks.""" model_files = [] for root in (args.base_model, args.adapter_dir): if not root: continue directory = Path(root) if not directory.is_dir(): continue candidates = set(directory.glob("*.json")) candidates.update(directory.glob("*.safetensors")) for path in sorted(candidates): hash_content = ( path.stat().st_size <= 16 * 1024 * 1024 or path.name == "adapter_model.safetensors" ) model_files.append(_file_identity(path, hash_content=hash_content)) data_paths = set() for task in tasks: for field in ("paired", "corpus", "queries", "qrels"): data_paths.update(task.get(field) or []) data_files = [_file_identity(path) for path in sorted(data_paths)] payload = { "protocol": EVAL_PROTOCOL, "run_name": args.run_name, "base_model": str(Path(args.base_model).resolve()), "adapter_dir": ( str(Path(args.adapter_dir).resolve()) if args.adapter_dir else None ), "skip_adapter_weights": args.skip_adapter_weights, "model_files": model_files, "training": { "col_dim": train_config.get("col_dim"), "head_dims": train_config.get("head_dims"), "teacher_md5": train_config.get("teacher_md5"), "max_visual_tokens": train_config.get("max_visual_tokens"), "bidirectional_attention": train_config.get("bidirectional_attention"), "attn": train_config.get("attn"), }, "evaluation": { "max_visual_tokens": args.max_visual_tokens, "bidirectional_attention": args.bidirectional_attention, # Part of the contract, not cosmetic: without it the six per-head passes # of one run would share a contract id and resume into each other's cache. "head_dim": args.head_dim, "attn": args.attn, "ks": list(ks), "max_queries": args.max_queries, "max_docs": args.max_docs, "datasets": args.datasets, "eval_root": str(Path(args.eval_root).resolve()), }, "data_files": data_files, "tasks": [task_tag(task) for task in tasks], } encoded = json.dumps( payload, ensure_ascii=False, sort_keys=True, separators=(",", ":") ).encode("utf-8") return hashlib.sha256(encoded).hexdigest() def collect_run_metadata(args, train_config, model, world): """Provenance for the results table: model / data / training / eval settings.""" adapter_dir = Path(args.adapter_dir) if args.adapter_dir else None adapter_cfg = ( _read_json(adapter_dir / "adapter_config.json") if adapter_dir and not args.skip_adapter_weights else {} ) base_cfg = _read_json(Path(args.base_model) / "config.json") data_root = train_config.get("data_root") export = _read_json(Path(data_root) / "export_stats.json") if data_root else {} steps = train_config.get("steps") epoch_done = None world_train = train_config.get("world_size") if adapter_dir: ckpts = sorted( adapter_dir.glob("checkpoint-*"), key=lambda p: int(p.name.rsplit("-", 1)[-1]) ) if ckpts: state = _read_json(ckpts[-1] / "trainer_state.json") steps = steps or state.get("global_step") epoch_done = state.get("epoch") world_train = world_train or len(list(ckpts[-1].glob("rng_state_*.pth"))) or None per_device = train_config.get("per_device_batch_size") accum = train_config.get("grad_accum") eff_batch = ( per_device * accum * world_train if None not in (per_device, accum, world_train) else None ) total_params = sum(p.numel() for p in model.parameters()) lora_params = sum(p.numel() for n, p in model.named_parameters() if "lora_" in n) text_cfg = base_cfg.get("text_config") or {} hardneg_loss = train_config.get("hardneg_loss", "negative_ce") listwise = hardneg_loss == "listwise" hardneg = bool(train_config.get("use_hardnegatives")) or listwise git_hash = "" if adapter_dir and (adapter_dir / "git_hash.txt").is_file(): git_hash = (adapter_dir / "git_hash.txt").read_text(encoding="utf-8").strip() return dict( model=dict( run_name=args.run_name, adapter_dir=str(adapter_dir) if adapter_dir else None, base_model=train_config.get("base_model") or args.base_model, experiment_type="/".join( [ "EVIE" if args.skip_adapter_weights else ("LoRA" if adapter_cfg else "full-ft"), "bidir" if args.bidirectional_attention == "on" else "causal", "listwise" if listwise else ("hardneg" if hardneg else "in-batch"), ] ), total_params=total_params, trainable_params=lora_params, trainable_ratio=round(lora_params / total_params, 6) if total_params else None, # head's, not the warm start's config.json dim (Preview says 128). emb_dim=args.head_dim if args.head_dim is not None else base_cfg.get("dim"), head_dim=args.head_dim, head_dims=train_config.get("head_dims"), teacher_dir=train_config.get("teacher_dir"), teacher_md5=train_config.get("teacher_md5"), hidden_size=text_cfg.get("hidden_size"), num_hidden_layers=text_cfg.get("num_hidden_layers"), dtype="bfloat16", attn=args.attn, git_hash=git_hash, ), data=dict( data_root=data_root, version=export.get("version"), policy=export.get("policy"), train_pairs=export.get("written_rows"), by_source=export.get("by_source"), max_samples_per_source=train_config.get("max_samples_per_source"), hardneg_root=( train_config.get("listwise_root") if listwise else train_config.get("hardneg_root") ) if hardneg else None, num_hard_negs=train_config.get("num_hard_negs") if not listwise and hardneg else None, ), training=dict( seed=train_config.get("seed"), epochs=train_config.get("epochs"), epoch_done=epoch_done, steps=steps, world_size=world_train, per_device_batch_size=per_device, grad_accum=accum, effective_batch_size=eff_batch, consumed_samples=steps * eff_batch if (steps and eff_batch) else None, learning_rate=train_config.get("learning_rate"), lr_scheduler="cosine", warmup_ratio=train_config.get("warmup_ratio"), weight_decay=train_config.get("weight_decay"), lora_r=adapter_cfg.get("r", train_config.get("lora_r")), lora_alpha=adapter_cfg.get("lora_alpha", train_config.get("lora_alpha")), lora_dropout=adapter_cfg.get("lora_dropout", train_config.get("lora_dropout")), lora_target_modules=adapter_cfg.get("target_modules"), loss=( "ColbertListwiseKLLoss" if listwise else ("ColbertNegativeCELoss" if hardneg else "ColbertLoss") ), loss_temperature=train_config.get("loss_temperature"), in_batch_term_weight=train_config.get("hardneg_in_batch_weight") if hardneg else None, max_visual_tokens=train_config.get("max_visual_tokens"), grad_checkpointing=train_config.get("grad_checkpointing"), train_runtime_seconds=train_config.get("train_runtime_seconds"), ), eval=dict( eval_root=args.eval_root, eval_protocol=EVAL_PROTOCOL, max_visual_tokens=args.max_visual_tokens, bidirectional_attention=args.bidirectional_attention, embed_batch=args.embed_batch, num_workers=args.num_workers, world_size=world, datasets_filter=args.datasets or None, max_queries=args.max_queries or None, max_docs=args.max_docs or None, truncated=bool(args.max_queries or args.max_docs), date=time.strftime("%Y-%m-%d %H:%M:%S"), ), ) DS_LABEL = { "eval": "ViDoRe V1", "eval_v2": "ViDoRe V2", "eval_v3": "ViDoRe V3", "JinaVDR": "JinaVDR", "demo": "Bundled demo", } def write_summary_md(path, summary, ks): md = summary["metadata"] model, data, train, ev = md["model"], md["data"], md["training"], md["eval"] head, avg, per_ds = summary["headline"], summary["averages"], summary["per_dataset"] def pct(v): return "—" if v is None else f"{v * 100:.2f}" def num(v): return "—" if v is None else f"{v:,}" def params(v): return "—" if v is None else (f"{v / 1e9:.2f}B" if v >= 1e9 else f"{v / 1e6:.1f}M") lines = [ f"# {model['run_name']} eval", "", f"status `{summary['status']}` · {summary['completed_tasks']}/{summary['expected_tasks']} complete · " f"{summary['n_failed']} failed · {ev['date']}", "", "## Headline", "", "| model | type | size | trainable | emb | train pairs | " "V1 nDCG@10 | V2 nDCG@10 | V3 nDCG@10 | Jina nDCG@10 | avg4 |", "| --- | --- | ---: | ---: | ---: | ---: | ---: | ---: | ---: | ---: | ---: |", f"| {model['run_name']} | {model['experiment_type']} | {params(model['total_params'])} | " f"{params(model['trainable_params'])} | {model['emb_dim']} | {num(data['train_pairs'])} | " f"{pct(head.get('V1 nDCG@10'))} | {pct(head.get('V2 nDCG@10'))} | " f"{pct(head.get('V3 nDCG@10'))} | {pct(head.get('JinaVDR nDCG@10'))} | " f"{pct(avg.get('avg4_ndcg@10'))} |", "", "## Metrics (four families × @1/@5/@10)", "", "| board | subsets | queries | docs | " + " | ".join(f"{m}@{k}" for m in METRIC_NAMES for k in ks) + " |", "| --- | ---: | ---: | ---: |" + " ---: |" * (len(METRIC_NAMES) * len(ks)), ] for ds, label in DS_LABEL.items(): e = per_ds.get(ds) if not e: continue cells = " | ".join(pct(e.get(f"{m}@{k}")) for m in METRIC_NAMES for k in ks) lines.append( f"| {label} | {e['n_subsets']}/{e['n_subsets_expected']} | " f"{num(e['n_queries'])} | {num(e['n_docs'])} | {cells} |" ) lines += [ "", "## Config", "", "| key | value |", "| --- | --- |", f"| base_model | `{model['base_model']}` |", f"| architecture | hidden {model['hidden_size']} × {model['num_hidden_layers']} · emb {model['emb_dim']} |", f"| dtype / attn | {model['dtype']} / {model['attn']} |", f"| data | {data['version'] or '—'} · {data['policy'] or '—'} · {num(data['train_pairs'])} pairs |", f"| data_root | `{data['data_root']}` |", f"| hardneg | {data['hardneg_root'] or '—'} · num_negs={data['num_hard_negs'] or '—'} |", f"| seed / epochs / steps | {train['seed']} / {train['epochs']} / {train['steps']} |", f"| batch | per_device {train['per_device_batch_size']} × accum {train['grad_accum']} " f"× world {train['world_size']} = {train['effective_batch_size']} |", f"| consumed samples | {num(train['consumed_samples'])} |", f"| lr / warmup / wd | {train['learning_rate']} ({train['lr_scheduler']}) / " f"{train['warmup_ratio']} / {train['weight_decay']} |", f"| LoRA | r={train['lora_r']} α={train['lora_alpha']} dropout={train['lora_dropout']} |", f"| loss | {train['loss']} · T={train['loss_temperature'] or '—'} · " f"in_batch_weight={train['in_batch_term_weight'] or '—'} |", f"| MVT train / eval | {train['max_visual_tokens']} / {ev['max_visual_tokens']} |", f"| bidir eval | {ev['bidirectional_attention']} |", f"| protocol | `{ev.get('eval_protocol') or '—'}` |", f"| eval data | `{ev['eval_root']}` · ks={ks} · truncated={ev['truncated']} |", ] if summary["failed_tasks"]: lines += ["", "## Failed tasks", ""] lines += [f"- `{r['task']}`: {r['error']}" for r in summary["failed_tasks"]] Path(path).write_text("\n".join(lines) + "\n", encoding="utf-8") def _task_corpus_size(task): """Approx embed workload (no image decode) for load balancing.""" paths = task.get("paired") or task.get("corpus") or [] if not paths: return 0 try: return sum(pq.ParquetFile(p).metadata.num_rows for p in paths) except Exception: return len(paths) def _balanced_shard(tasks, world, rank): """Greedy split while keeping tasks that share a corpus on the same rank.""" grouped = {} for task in tasks: key = ( ("beir-corpus", tuple(task["corpus"])) if task["fmt"] == "beir" else ("single-task", task_tag(task)) ) grouped.setdefault(key, []).append(task) sized = sorted( ( (max(_task_corpus_size(task) for task in group), i, group) for i, group in enumerate(grouped.values()) ), key=lambda item: -item[0], ) loads = [0] * world buckets = [[] for _ in range(world)] for sz, _, group in sized: r = min(range(world), key=lambda w: loads[w]) buckets[r].extend(group) loads[r] += sz return buckets[rank] def setup_ddp(): if "LOCAL_RANK" not in os.environ and "RANK" not in os.environ: return 0, 1, 0 rank = int(os.environ.get("RANK", os.environ.get("LOCAL_RANK", 0))) world = int(os.environ.get("WORLD_SIZE", 1)) local = int(os.environ.get("LOCAL_RANK", rank)) timeout = timedelta(seconds=int(os.environ.get("EVAL_DDP_TIMEOUT_S", "21600"))) # Bind this process to its GPU before the first collective, otherwise every # rank allocates NCCL buffers on cuda:0 ("Duplicate GPU detected"). if torch.cuda.is_available(): torch.cuda.set_device(local) torch.distributed.init_process_group( backend="nccl", device_id=torch.device(f"cuda:{local}"), timeout=timeout ) else: torch.distributed.init_process_group(backend="gloo", timeout=timeout) return rank, world, local def main(): args = parse_args() config_path, train_config = resolve_eval_contract(args) rank, world, local = setup_ddp() device = torch.device(f"cuda:{local}" if torch.cuda.is_available() else "cpu") ks = [int(k) for k in args.ks.split(",") if k.strip()] if tuple(ks) != REQUIRED_KS: raise ValueError( f"complete evaluation requires --ks 1,5,10; got {args.ks!r}" ) out_dir = forbid_venv_path(args.output_dir, "output-dir") resume = bool(args.resume) and not args.overwrite_output tasks = discover_tasks(args.eval_root) if args.datasets: want = {d.strip() for d in args.datasets.split(",") if d.strip()} tasks = [t for t in tasks if t["dataset"] in want] else: tasks = [t for t in tasks if t["dataset"] != "demo"] expected_tasks = len(tasks) expected_by_ds = Counter(t["dataset"] for t in tasks) if not args.datasets: want_full = {"eval": 10, "eval_v2": 4, "eval_v3": 48, "JinaVDR": 76} bad = { key: (want_full[key], expected_by_ds.get(key, 0)) for key in want_full if expected_by_ds.get(key, 0) != want_full[key] } if bad: raise RuntimeError(f"eval discovery incomplete: {bad}; refuse to start") # Only rank 0 mutates the shared output path. prep_error = [None] cached_payload = [None] contract_payload = [None] if rank == 0: try: contract_payload[0] = build_eval_contract(args, train_config, tasks, ks) if args.overwrite_output and out_dir.exists() and any(out_dir.iterdir()): stamp = time.strftime("%Y%m%d_%H%M%S") archive = out_dir.with_name(f"{out_dir.name}_archive_{stamp}") out_dir.rename(archive) print(f"[eval] archived previous output -> {archive}") elif out_dir.exists() and any(out_dir.iterdir()) and not resume: raise FileExistsError( f"non-empty eval output: {out_dir}; pass --resume to continue " "or --overwrite-output to archive it" ) out_dir.mkdir(parents=True, exist_ok=True) # Drop stale rank shards from prior world-size / crashed runs so merge # never re-ingests old errors or metrics from leftover rank_*.json. for stale in out_dir.glob("rank_*.json"): stale.unlink() for stale in (out_dir / "summary.json", out_dir / "summary.md"): stale.unlink(missing_ok=True) cached_payload[0] = load_cached_task_map( out_dir, contract_id=contract_payload[0] ) print( f"[eval] run={args.run_name} protocol={EVAL_PROTOCOL} " f"contract={contract_payload[0][:12]} " f"mode={'resume' if resume else 'fresh'} " f"cached_ok={len(cached_payload[0])} out={out_dir}" ) except Exception as exc: prep_error[0] = repr(exc) if world > 1: torch.distributed.broadcast_object_list(prep_error, src=0) torch.distributed.broadcast_object_list(cached_payload, src=0) torch.distributed.broadcast_object_list(contract_payload, src=0) if prep_error[0]: if world > 1: torch.distributed.destroy_process_group() raise RuntimeError(f"eval output preflight failed: {prep_error[0]}") if world > 1: torch.distributed.barrier() cached = cached_payload[0] or {} eval_contract_id = contract_payload[0] model_source = args.base_model is_adapter = bool( args.adapter_dir and not args.skip_adapter_weights and (Path(args.adapter_dir) / "adapter_config.json").is_file() ) if args.adapter_dir and not is_adapter and not args.skip_adapter_weights: model_source = args.adapter_dir print( f"[eval rank{rank}] loading model on {device} " f"(source={model_source}, adapter={args.adapter_dir if is_adapter else None})" ) processor = ColQwen3_5Processor.from_pretrained( model_source, max_num_visual_tokens=args.max_visual_tokens ) model_config = Qwen3_5Config.from_pretrained(model_source) trained_col_dim = train_config.get("col_dim") trained_heads = train_config.get("head_dims") or getattr( model_config, "head_dims", None ) if trained_heads: model_config.head_dims = [int(d) for d in trained_heads] model_config.dim = max(model_config.head_dims) mrl_prefix = bool(train_config.get("mrl_prefix", getattr(model_config, "mrl_prefix", False))) model_config.mrl_prefix = mrl_prefix layout = "prefix-MRL" if mrl_prefix else "ModuleDict" print( f"[eval rank{rank}] custom_text_proj = {layout}{model_config.head_dims}, " f"scoring head d={args.head_dim}" ) elif is_adapter: if trained_col_dim is None: raise ValueError("adapter evaluation requires col_dim in run_config.json") model_config.dim = int(trained_col_dim) print(f"[eval rank{rank}] custom_text_proj dim={model_config.dim} from run_config.json") model = ColQwen3_5.from_pretrained( model_source, config=model_config, torch_dtype=torch.bfloat16, attn_implementation=args.attn, ) if args.bidirectional_attention == "on": model.enable_bidirectional_attention() print("[model] bidirectional attention ON") else: print("[model] bidirectional attention OFF (causal)") if is_adapter: model = PeftModel.from_pretrained(model, args.adapter_dir) if trained_heads: # After PEFT wrapping, so the selection reaches the ColQwen3_5 underneath. # Every downstream scorer expects one tensor per forward. set_active_head(model, int(args.head_dim)) model = model.to(device).eval() pending = [ t for t in tasks if not (resume and task_tag(t) in cached) ] if world > 1: pending = _balanced_shard(pending, world, rank) print( f"[eval rank{rank}] pending {len(pending)} " f"(expected_total={expected_tasks}, cached_ok={len(cached)}) " f"(~{sum(_task_corpus_size(t) for t in pending)} corpus imgs) " f"counts={dict(expected_by_ds)}" ) results = [] # One-entry GPU cache. _balanced_shard keeps shared V3 language tasks # contiguous on the same rank, so the corpus tensor is reused six times # without accumulating multiple domains in VRAM. beir_corpus_cache = {} for t in pending: tag = task_tag(t) try: t0 = time.time() if t["fmt"] == "paired": metrics, info = run_paired( t, processor, model, device, ks, args.embed_batch, args.num_workers, args.max_queries, args.max_docs, ) else: metrics, info = run_beir( t, processor, model, device, ks, args.embed_batch, args.num_workers, args.max_queries, args.max_docs, corpus_cache=beir_corpus_cache, ) dt = time.time() - t0 rec = dict( task=tag, dataset=t["dataset"], subset=t["subset"], lang=t.get("lang"), fmt=t["fmt"], metrics=metrics, info=info, seconds=round(dt, 1), eval_protocol=EVAL_PROTOCOL, eval_contract_id=eval_contract_id, run_name=args.run_name, ) results.append(rec) save_task_result( out_dir, rec, run_name=args.run_name, contract_id=eval_contract_id, ) print( f"[eval rank{rank}] {tag}: {metrics} " f"({info['n_queries']}q / {info['n_docs']}d, {dt:.1f}s)" ) except Exception as e: import traceback print(f"[eval rank{rank}] {tag} FAILED: {e!r}") traceback.print_exc() results.append( dict( task=tag, dataset=t["dataset"], subset=t["subset"], lang=t.get("lang"), fmt=t["fmt"], error=str(e), eval_protocol=EVAL_PROTOCOL, eval_contract_id=eval_contract_id, run_name=args.run_name, ) ) rank_path = out_dir / f"rank_{rank}.json" rank_tmp = out_dir / f".rank_{rank}.json.tmp.{os.getpid()}" rank_tmp.write_text( json.dumps(results, ensure_ascii=False, indent=2), encoding="utf-8" ) rank_tmp.replace(rank_path) if torch.cuda.is_available(): peak = torch.cuda.max_memory_allocated(device) / 1024**3 print(f"[eval rank{rank}] peak VRAM: {peak:.2f} GB on {device}") if world > 1: # Evaluation ranks have highly uneven runtimes (large V3 corpora versus # tiny paired sets). Do not hold finished ranks in a NCCL collective: # the old final barrier deterministically timed out while slow ranks # were still encoding. Rank 0 finalizes through atomic shared-FS shards. torch.distributed.destroy_process_group() if rank != 0: return finalize_timeout = int(os.environ.get("EVAL_FINALIZE_TIMEOUT_S", "21600")) deadline = time.monotonic() + finalize_timeout last_report = 0.0 missing_ranks = list(range(world)) while missing_ranks: missing_ranks = [] for r in range(world): p = out_dir / f"rank_{r}.json" try: json.loads(p.read_text(encoding="utf-8")) except (FileNotFoundError, json.JSONDecodeError, OSError): missing_ranks.append(r) if not missing_ranks: break now = time.monotonic() if now >= deadline: print( f"[eval] finalize timeout after {finalize_timeout}s; " f"missing rank shards: {missing_ranks}" ) break if now - last_report >= 60: print( f"[eval] waiting for {len(missing_ranks)}/{world} rank shards: " f"{missing_ranks}" ) last_report = now time.sleep(2) incomplete = [0] if rank == 0: # Prefer durable tasks/ cache; fold in this-run failures from rank_*.json. # Ignore leftover rank_N.json from a previous larger world size. merged_map = load_cached_task_map( out_dir, contract_id=eval_contract_id ) failed = [] for r in range(world): p = out_dir / f"rank_{r}.json" if not p.exists(): continue for rec in json.loads(p.read_text(encoding="utf-8")): if "error" in rec: failed.append(rec) elif _is_ok_task_rec(rec, contract_id=eval_contract_id): merged_map[rec["task"]] = rec # Ignore unrecognized rows (e.g. protocol-mismatched leftovers). # Drop failures that were successfully produced this session (or still cached). failed = [f for f in failed if f["task"] not in merged_map] order = {task_tag(t): i for i, t in enumerate(tasks)} expected_tags = set(order) missing_tags = [task_tag(t) for t in tasks if task_tag(t) not in merged_map] merged = sorted( (rec for rec in merged_map.values() if rec["task"] in expected_tags), key=lambda rec: order.get(rec["task"], 10**9), ) failed_count = len(failed) by_ds = defaultdict(list) for rec in merged: by_ds[rec["dataset"]].append(rec) def avg_metric(recs, name): vals = [r["metrics"].get(name) for r in recs] vals = [v for v in vals if v is not None] return (sum(vals) / len(vals)) if vals else None per_ds = {} for ds, recs in by_ds.items(): entry = {f"{m}@{k}": avg_metric(recs, f"{m}@{k}") for m in METRIC_NAMES for k in ks} entry["n_subsets"] = len(recs) entry["n_subsets_expected"] = expected_by_ds.get(ds, len(recs)) entry["n_queries"] = sum(r["info"]["n_queries"] for r in recs) entry["n_docs"] = sum(r["info"]["n_docs"] for r in recs) per_ds[ds] = entry hdr = " ".join(f"{m + '@' + str(k):>9}" for m in METRIC_NAMES for k in ks) def fmt_row(get): return " ".join(f"{(get(f'{m}@{k}') or 0) * 100:8.2f}%" for m in METRIC_NAMES for k in ks) print("\n================ EVAL SUMMARY ================") for ds in DS_LABEL: if ds not in by_ds: continue m = per_ds[ds] print( f"\n-- {DS_LABEL[ds]} ({m['n_subsets']}/{m['n_subsets_expected']} subsets, " f"{m['n_queries']}q / {m['n_docs']}d) --" ) print(f" {hdr}") print(f" {fmt_row(m.get)}") all_recs = [r for ds in by_ds for r in by_ds[ds]] print(f"\n-- ALL DATASETS ({len(all_recs)}/{expected_tasks} subsets) --") print(f" {hdr}") print(f" {fmt_row(lambda name: avg_metric(all_recs, name))}") headline = { "V1 nDCG@10": per_ds.get("eval", {}).get("ndcg@10"), "V2 nDCG@10": per_ds.get("eval_v2", {}).get("ndcg@10"), "V3 nDCG@10": per_ds.get("eval_v3", {}).get("ndcg@10"), "JinaVDR nDCG@10": per_ds.get("JinaVDR", {}).get("ndcg@10"), } four = [headline[k] for k in ( "V1 nDCG@10", "V2 nDCG@10", "V3 nDCG@10", "JinaVDR nDCG@10" )] averages = { "avg4_ndcg@10": ( sum(four) / 4 if all(v is not None for v in four) else None ), } print("\n================ HEADLINE (nDCG@10) ================") for name, v in headline.items(): print(f" {name:<18}: {'n/a' if v is None else f'{v * 100:.2f}%'}") avg4 = averages["avg4_ndcg@10"] print(f" {'Avg4 nDCG@10':<18}: {'n/a' if avg4 is None else f'{avg4 * 100:.2f}%'}") invalid_metric_tasks = [ rec["task"] for rec in merged if not _has_complete_metrics(rec["metrics"]) ] # Complete iff every task belongs to this exact contract and contains all # four metric families at @1/@5/@10. complete = not failed and not missing_tags and not invalid_metric_tasks metadata = collect_run_metadata(args, train_config, model, world) summary = dict( status="complete" if complete else "partial", run_name=args.run_name, adapter_dir=args.adapter_dir, eval_protocol=EVAL_PROTOCOL, eval_contract_id=eval_contract_id, training_config=str(config_path) if config_path else None, max_visual_tokens=args.max_visual_tokens, bidirectional_attention=args.bidirectional_attention, ks=ks, expected_tasks=expected_tasks, completed_tasks=expected_tasks - len(missing_tags), n_failed=failed_count, n_cached=len(cached), n_ran=sum(1 for r in results if "metrics" in r) if world == 1 else None, failed_tasks=failed, missing_tasks=missing_tags, invalid_metric_tasks=invalid_metric_tasks, headline=headline, averages=averages, metadata=metadata, per_dataset=per_ds, per_task=merged, ) summary_path = out_dir / "summary.json" summary_tmp = out_dir / f".summary.json.tmp.{os.getpid()}" summary_tmp.write_text( json.dumps(summary, ensure_ascii=False, indent=2), encoding="utf-8" ) summary_tmp.replace(summary_path) write_summary_md(out_dir / "summary.md", summary, ks) print(f"\n[eval] summary written to {out_dir / 'summary.json'} and summary.md") if not complete: print( f"\n[eval] INCOMPLETE: {len(merged)}/{expected_tasks} ok, " f"{failed_count} failed this run, {len(missing_tags)} missing" ) for rec in failed: print(f" - FAIL {rec['task']}: {rec['error']}") for tag in missing_tags[:20]: print(f" - MISS {tag}") if len(missing_tags) > 20: print(f" - ... {len(missing_tags) - 20} more missing") print( f"[eval] tip: delete bad files under {tasks_dir(out_dir)} then " "re-run with --resume to patch only those tasks" ) incomplete[0] = 0 if complete else 1 if incomplete[0]: raise SystemExit(2) if __name__ == "__main__": main()