# -*- coding: utf-8 -*- import logging import json import random import time import numpy as np import os import pickle import sys import torch import torch.distributed as dist import torch.nn.functional as F import yaml from torch.utils.data import DataLoader from tqdm import tqdm from transformers import HfArgumentParser, AutoConfig from datasets import concatenate_datasets from datasets.distributed import split_dataset_by_node from src.arguments import ModelArguments, DataArguments, TrainingArguments from src.data.collator.eval_collator import MultimodalEvalDataCollator from src.data.eval_dataset.base_eval_dataset import AutoEvalPairDataset, generate_cand_dataset from src.eval_utils.metrics import RankingMetrics from src.model.model_cut_layer_AOP import MMEBModel # NOTE: 使用你的 AOP 版本（支持 cut_layer + mask 透传） from src.model.processor import get_backbone_name, load_processor, COLPALI from src.utils import batch_to_device, print_rank, print_master logging.basicConfig(level=logging.INFO, format='[%(asctime)s] %(levelname)s [%(name)s:%(lineno)s] %(message)s') logger = logging.getLogger(__name__) # ----------------- 环境变量解析 ----------------- def _parse_bool(v: str, default=False): if v is None: return default v = v.strip().lower() return v in {"1","true","yes","y","t","on"} def _parse_float(v: str, default=None): try: return float(v) if v is not None else default except: return default def _parse_int(v: str, default=None): try: return int(v) if v is not None else default except: return default def get_env_eval_layers(): """ LM_LAYERS: "4,8,12,last"（last/none/-1 -> None 代表最后一层） 未设置则默认 [None] """ v = os.environ.get("LM_LAYERS", None) if v: toks = [t.strip() for t in v.split(',') if t.strip() != ""] layers = [] for tok in toks: tl = tok.lower() if tl in {"last","none","null","-1"}: layers.append(None) else: try: val = int(tok) if val > 0: layers.append(val) except: logger.warning(f"Invalid token '{tok}' in LM_LAYERS; ignored.") # 去重保序 seen=set(); uniq=[] for l in layers: key = -1 if l is None else l if key in seen: continue seen.add(key); uniq.append(l) return uniq or [None] return [None] def get_env_zip_config(): """ VisionZip（输入侧 token 压缩）配置： - ZIP_ENABLED=1 开启 - ZIP_APPLY=qry|cand|both - ZIP_METHOD=visionzip|none - ZIP_KEEP_DOM / ZIP_KEEP_CTX = dominant/context 保留份额 """ cfg = { "enabled": _parse_bool(os.environ.get("ZIP_ENABLED"), False), "apply_to": (os.environ.get("ZIP_APPLY","both").strip().lower()), "method": os.environ.get("ZIP_METHOD","visionzip").strip().lower(), "keep_dom": _parse_float(os.environ.get("ZIP_KEEP_DOM"), 0.45), "keep_ctx": _parse_float(os.environ.get("ZIP_KEEP_CTX"), 0.10), } if cfg["method"] == "none": cfg["enabled"] = False return cfg def get_env_aop_config(): """ AOP（层内剪裁）配置： - AOP_ENABLED=1 - AOP_APPLY=qry|cand|both - AOP_LAYER=N（1-based，在进入该层前剪裁一次） - AOP_MODE=delta|ratio - AOP_KEEP_RATIO / AOP_DELTA / AOP_KHAT / AOP_MIN_KEEP / AOP_USE_BIAS - AOP_ATTN_IMPL=sdpa 可选覆盖注意力实现（便于取权重或稳定） """ enabled = _parse_bool(os.environ.get("AOP_ENABLED"), False) apply_to = os.environ.get("AOP_APPLY", "qry").strip().lower() layer_idx = _parse_int(os.environ.get("AOP_LAYER"), None) mode = os.environ.get("AOP_MODE", "delta").strip().lower() delta = _parse_float(os.environ.get("AOP_DELTA"), 0.10) khat = _parse_float(os.environ.get("AOP_KHAT"), 1.0) keep_ratio = _parse_float(os.environ.get("AOP_KEEP_RATIO"), 1.0) min_keep = _parse_int(os.environ.get("AOP_MIN_KEEP"), 64) use_bias = _parse_bool(os.environ.get("AOP_USE_BIAS"), True) attn_impl = os.environ.get("AOP_ATTN_IMPL","").strip().lower() if layer_idx is None and enabled: logger.warning("AOP_ENABLED=1 但未设置 AOP_LAYER，关闭 AOP。"); enabled=False return { "enabled": enabled, "apply_to": apply_to, "layer_idx": layer_idx, "mode": mode, "delta": delta, "K_hat": khat, "keep_ratio": keep_ratio, "min_keep": min_keep, "use_bias": use_bias, "eps": 1e-6, "attn_impl_override": attn_impl if attn_impl in {"sdpa"} else "", } # ----------------- Hook & utils ----------------- timing_info = {} token_info = {"vision_tokens":0,"text_input_tokens":0,"text_output_tokens":0,"total_llm_input_tokens":0} def timing_pre_hook(module, input): module_id = id(module) if module_id not in timing_info: timing_info[module_id] = [] timing_info[module_id].append((time.time(), 'pre', module.__class__.__name__)) def timing_post_hook(module, input, output): module_id = id(module) if module_id not in timing_info: return timing_info[module_id].append((time.time(), 'post', module.__class__.__name__)) module_name = module.__class__.__name__ if "vision" in module_name.lower() and "transformer" in module_name.lower(): if isinstance(output, torch.Tensor): token_info["vision_tokens"] = output.shape[0] elif hasattr(output, 'last_hidden_state'): token_info["vision_tokens"] = output.last_hidden_state.shape[1] def register_model_hooks(model): # 在 encoder 的 visual/merger/LLM/lm_head 上打点 regs=[] core = model if hasattr(model,'encoder') and model.encoder is not None: core = model.encoder if hasattr(core,'visual') and core.visual is not None: core.visual.register_forward_pre_hook(timing_pre_hook) core.visual.register_forward_hook(timing_post_hook); regs.append(core.visual) if hasattr(core,'visual') and hasattr(core.visual,'merger') and core.visual.merger is not None: core.visual.merger.register_forward_pre_hook(timing_pre_hook) core.visual.merger.register_forward_hook(timing_post_hook); regs.append(core.visual.merger) if hasattr(core,'model') and core.model is not None: core.model.register_forward_pre_hook(timing_pre_hook) core.model.register_forward_hook(timing_post_hook); regs.append(core.model) if hasattr(core,'lm_head') and core.lm_head is not None: core.lm_head.register_forward_pre_hook(timing_pre_hook) core.lm_head.register_forward_hook(timing_post_hook); regs.append(core.lm_head) return regs def pad_dataset_to_divisible(dataset, world_size): n = len(dataset) if n % world_size == 0: return dataset, n m = world_size - (n % world_size) pad = dataset.select([i % len(dataset) for i in range(m)]) return concatenate_datasets([dataset, pad]), n + m # ----------------- 编码函数（合并 AOP + VisionZip 注入 + cut_layer） ----------------- def encode_embeddings( model: MMEBModel, loader: DataLoader, training_args: TrainingArguments, model_args: ModelArguments, full_dataset, encode_side: str, zip_cfg: dict, aop_cfg: dict, description: str = "Encoding" ): local_rank = dist.get_rank() if dist.is_initialized() else 0 is_late_interaction = (model_args.model_backbone == COLPALI) embeds, infos, batch_stats_list, img_masks_all = [], [], [], [] local_max_len = 0 model.eval() regs = register_model_hooks(model) with torch.no_grad(): for inputs, dataset_info in tqdm(loader, desc=f"{description} (rank {local_rank})", disable=local_rank > 0): # reset stats timing_info.clear() token_info.update({"vision_tokens":0,"text_input_tokens":0,"text_output_tokens":0,"total_llm_input_tokens":0}) inputs = batch_to_device(inputs, training_args.device) B = inputs.get('input_ids', torch.empty(1,1)).shape[0] # VisionZip：按侧注入 zip_runtime_cfg if zip_cfg and zip_cfg.get("enabled", False): apply_to = zip_cfg.get("apply_to","both") side_enable = (apply_to == "both") or (apply_to == encode_side) or (encode_side=="cand" and apply_to=="tgt") if side_enable: inputs["zip_runtime_cfg"] = { "enable": True, "method": zip_cfg.get("method","visionzip"), "keep_dominant_ratio": float(zip_cfg.get("keep_dom", 0.45)), "keep_context_ratio": float(zip_cfg.get("keep_ctx", 0.10)), } # AOP：按侧临时开启 enabled（encoder.aop_prune_config 由 main 挂载） _orig_enabled = None enc = model.encoder if isinstance(aop_cfg, dict) and aop_cfg: _orig_enabled = aop_cfg.get("enabled", False) apply_to = aop_cfg.get("apply_to","qry") side_enable = (apply_to == "both") or (apply_to == encode_side) or (encode_side=="cand" and apply_to=="tgt") aop_cfg["enabled"] = bool(side_enable and _orig_enabled) setattr(enc, "aop_prune_config", aop_cfg) with torch.autocast(enabled=True, dtype=torch.bfloat16, device_type="cuda"): t0 = time.time() if encode_side == "qry": out = model(qry=inputs) reps = out["qry_reps"].detach() infos.extend(dataset_info) else: out = model(tgt=inputs) reps = out["tgt_reps"].detach() infos.extend([x["cand_name"] for x in dataset_info]) t1 = time.time() if isinstance(aop_cfg, dict) and _orig_enabled is not None: aop_cfg["enabled"] = _orig_enabled setattr(enc, "aop_prune_config", aop_cfg) # image token masks（优先从 out 取；你的 model_cut_layer_AOP 已返回） img_masks = None if isinstance(out, dict): img_masks = out.get("image_token_bool_masks", None) if img_masks is None and hasattr(model, "_last_image_token_bool_masks"): img_masks = getattr(model, "_last_image_token_bool_masks") # 统一序列化：list[None|list[bool]] * B if img_masks is None: img_masks_list = [None] * B elif torch.is_tensor(img_masks): if img_masks.dim() == 2: img_masks_list = img_masks.detach().cpu().tolist() else: img_masks_list = [None] * B elif isinstance(img_masks, list): # 尝试转 list tmp=[] for m in img_masks: if torch.is_tensor(m): tmp.append(m.detach().cpu().tolist()) else: tmp.append(m) img_masks_list = tmp if len(img_masks_list) < B: img_masks_list += [None]*(B-len(img_masks_list)) if len(img_masks_list) > B: img_masks_list = img_masks_list[:B] else: img_masks_list = [None] * B img_masks_all.extend(img_masks_list) # 统计 if 'input_ids' in inputs and inputs['input_ids'] is not None: token_info["total_llm_input_tokens"] = int(inputs['input_ids'].shape[1]) token_info["text_input_tokens"] = max(0, token_info["total_llm_input_tokens"] - token_info["vision_tokens"]) batch_stats = { "batch_size": B, "total_inference_time_seconds": float(t1 - t0), "module_inference_times": {}, "token_counts": { "visual_tokens": token_info["vision_tokens"], "language_input_tokens_raw": token_info["text_input_tokens"], "llm_total_input_tokens": token_info["total_llm_input_tokens"], "language_output_tokens": token_info["text_output_tokens"], } } # 模块时间 for m in regs: mid = id(m) name = m.__class__.__name__ times = timing_info.get(mid, []) durations=[]; pre=None for (ts, tp, _) in times: if tp == 'pre': pre = ts elif tp == 'post' and pre is not None: durations.append(ts - pre); pre=None if durations: batch_stats["module_inference_times"][name] = { "total": sum(durations), "count": len(durations), "avg": sum(durations)/len(durations) } else: batch_stats["module_inference_times"][name] = {"total":0.0,"count":0,"avg":0.0} batch_stats_list.append(batch_stats) print_rank(f"[{encode_side}] time={t1-t0:.4f}s, vis_tokens={token_info['vision_tokens']}") if is_late_interaction and reps.dim()==3: local_max_len = max(local_max_len, reps.shape[1]) embeds.append(reps) if not embeds: return np.array([]), [], [], [] # Late-interaction padding if is_late_interaction: if dist.is_initialized(): lm = torch.tensor(local_max_len, device=training_args.device) dist.all_reduce(lm, op=dist.ReduceOp.MAX) global_max_len = int(lm.item()) else: global_max_len = local_max_len padded=[] for e in embeds: if e.dim()==3: B, L, H = e.shape pad = global_max_len - L e = F.pad(e, (0,0,0,pad), "constant", 0) padded.append(e) embeds_tensor = torch.cat(padded, dim=0).contiguous() else: embeds_tensor = torch.cat(embeds, dim=0).contiguous() # DDP gather if dist.is_initialized() and len(full_dataset) >= dist.get_world_size(): print_master(f"Gathering {encode_side} embeddings across ranks...") output_shape = list(embeds_tensor.shape); output_shape[0] = len(full_dataset) embeds_tensor = embeds_tensor.to(training_args.device) gathered = torch.empty(output_shape, dtype=embeds_tensor.dtype, device=training_args.device) dist.all_gather_into_tensor(gathered, embeds_tensor) final_embeddings = gathered.cpu().float().numpy() gathered_infos=[None for _ in range(dist.get_world_size())] dist.all_gather_object(gathered_infos, infos) all_infos=[x for r in gathered_infos for x in r] gathered_stats=[None for _ in range(dist.get_world_size())] dist.all_gather_object(gathered_stats, batch_stats_list) all_stats=[s for r in gathered_stats for s in r] gathered_masks=[None for _ in range(dist.get_world_size())] dist.all_gather_object(gathered_masks, img_masks_all) all_masks=[m for r in gathered_masks for m in r] else: final_embeddings = embeds_tensor.cpu().float().numpy() all_infos = infos all_stats = batch_stats_list all_masks = img_masks_all return final_embeddings, all_infos, all_stats, all_masks # ----------------- 主入口 ----------------- def main(): # DDP init（与 torchrun 兼容） if "RANK" in os.environ and dist.is_available() and not dist.is_initialized(): timeout = int(os.environ.get("DDP_TIMEOUT_MIN", "60")) dist.init_process_group(backend="nccl", timeout=torch.distributed.timedelta(minutes=timeout)) for arg in sys.argv: if arg.startswith("--local-rank="): rank = arg.split("=")[1] sys.argv.remove(arg); sys.argv += ['--local_rank', rank] parser = HfArgumentParser((ModelArguments, DataArguments, TrainingArguments)) model_args, data_args, training_args = parser.parse_args_into_dataclasses() os.makedirs(data_args.encode_output_path, exist_ok=True) layers_to_eval = get_env_eval_layers() zip_cfg = get_env_zip_config() aop_cfg = get_env_aop_config() print_master(f"Eval layers: {layers_to_eval}") print_master(f"[ZIP] enabled={zip_cfg.get('enabled',False)}, apply_to={zip_cfg.get('apply_to')}, method={zip_cfg.get('method')}") print_master(f"[AOP] enabled={aop_cfg.get('enabled',False)}, apply_to={aop_cfg.get('apply_to')}, layer={aop_cfg.get('layer_idx')}") # 加载模型 hf_config = AutoConfig.from_pretrained(model_args.model_name, trust_remote_code=True) if not getattr(model_args, "model_backbone", None): model_backbone = get_backbone_name(hf_config=hf_config, model_type=model_args.model_type) setattr(model_args, 'model_backbone', model_backbone) setattr(training_args, 'model_backbone', model_backbone) print_master(f"Backbone: {model_args.model_backbone}") # 仅 rank0 先下载，其他等待缓存 local_rank = dist.get_rank() if dist.is_initialized() else 0 if local_rank == 0: processor = load_processor(model_args, data_args) model = MMEBModel.load(model_args, is_trainable=False, processor=processor) print_master(f"[rank=0] loaded {model_args.model_name}") if dist.is_initialized(): dist.barrier() if local_rank != 0: processor = load_processor(model_args, data_args) time.sleep(random.randint(2*local_rank, 3*local_rank)) model = MMEBModel.load(model_args, is_trainable=False, processor=processor) model.eval() model = model.to(training_args.device, dtype=torch.bfloat16) # 默认“最后一层”不裁层；每次循环会按需覆盖 model.set_inference_layers(qry_layers=None, tgt_layers=None) # 注入 AOP 底模配置（实例属性） if aop_cfg.get("enabled", False): setattr(model.encoder, "aop_prune_config", aop_cfg) attn_override = aop_cfg.get("attn_impl_override","") if attn_override and hasattr(model.encoder,"model") and hasattr(model.encoder.model,"config"): prev = model.encoder.model.config._attn_implementation model.encoder.model.config._attn_implementation = attn_override print_master(f"[AOP] attn impl override: {prev} -> {attn_override}") else: print_master("[AOP] disabled") with open(data_args.dataset_config, 'r') as f: dataset_configs = yaml.safe_load(f) # ----------------------- 主评测循环 ----------------------- for dataset_name, task_cfg in dataset_configs.items(): if dist.is_initialized(): dist.barrier() print_master(f"\n--- Evaluating {dataset_name} ---") # 基目录修正 if data_args.data_basedir: for k in ["image_root","video_root","frame_root","clip_root","data_path"]: if task_cfg.get(k): task_cfg[k] = os.path.join(data_args.data_basedir, task_cfg[k]) full_qry_dataset, corpus = AutoEvalPairDataset.instantiate(model_args=model_args, data_args=data_args, **task_cfg) full_cand_dataset = generate_cand_dataset(full_qry_dataset, corpus) eval_qry_dataset, eval_cand_dataset = full_qry_dataset, full_cand_dataset if dist.is_initialized(): ws = dist.get_world_size() padded_qry, _ = pad_dataset_to_divisible(full_qry_dataset, ws) padded_cand, _ = pad_dataset_to_divisible(full_cand_dataset, ws) eval_qry_dataset = split_dataset_by_node(padded_qry, rank=local_rank, world_size=ws) eval_cand_dataset = split_dataset_by_node(padded_cand, rank=local_rank, world_size=ws) for keep_layers in layers_to_eval: tag = f"layer{keep_layers}" if keep_layers else "layerlast" print_master(f"[{dataset_name}] tag={tag}, keep_layers={keep_layers}") model.set_inference_layers(qry_layers=keep_layers, tgt_layers=keep_layers) # 输出路径 q_path = os.path.join(data_args.encode_output_path, f"{dataset_name}_qry_{tag}") c_path = os.path.join(data_args.encode_output_path, f"{dataset_name}_tgt_{tag}") info_path = os.path.join(data_args.encode_output_path, f"{dataset_name}_info.jsonl") q_stats_path = os.path.join(data_args.encode_output_path, f"{dataset_name}_qry_inference_stats_{tag}.json") c_stats_path = os.path.join(data_args.encode_output_path, f"{dataset_name}_cand_inference_stats_{tag}.json") q_masks_path = os.path.join(data_args.encode_output_path, f"{dataset_name}_qry_img_token_masks_{tag}.jsonl") c_masks_path = os.path.join(data_args.encode_output_path, f"{dataset_name}_cand_img_token_masks_{tag}.jsonl") # 统计累积器 def _init_total(): return { "total_inference_time_seconds": 0.0, "module_inference_times": {}, "token_counts": {"visual_tokens":0,"language_input_tokens_raw":0,"llm_total_input_tokens":0,"language_output_tokens":0}, "data_point_count": 0 } def _acc(total, bs): bsz = bs["batch_size"] total["total_inference_time_seconds"] += bs["total_inference_time_seconds"] for m, ms in bs["module_inference_times"].items(): if m not in total["module_inference_times"]: total["module_inference_times"][m] = {"total":0.0,"count":0} total["module_inference_times"][m]["total"] += ms.get("total",0.0) total["module_inference_times"][m]["count"] += ms.get("count",0) for k in total["token_counts"]: total["token_counts"][k] += bs["token_counts"][k] * bsz total["data_point_count"] += bsz def _finalize(total, out_path, task_name, side_name): if local_rank != 0: return n = max(1, total["data_point_count"]) final = { "task_name": task_name, "encode_side": side_name, "data_point_count": total["data_point_count"], "inference_times":{ "total_inference_time_seconds": total["total_inference_time_seconds"], "avg_inference_time_per_item_seconds": total["total_inference_time_seconds"]/n, "module_average_times_per_call": {}, "module_total_times_seconds": {}, "module_calls_count": {}, }, "token_counts":{ "total_visual_tokens": total["token_counts"]["visual_tokens"], "avg_visual_tokens_per_item": total["token_counts"]["visual_tokens"]/n, "total_language_input_tokens_raw": total["token_counts"]["language_input_tokens_raw"], "avg_language_input_tokens_raw_per_item": total["token_counts"]["language_input_tokens_raw"]/n, "total_llm_total_input_tokens": total["token_counts"]["llm_total_input_tokens"], "avg_llm_total_input_tokens_per_item": total["token_counts"]["llm_total_input_tokens"]/n, "total_language_output_tokens": total["token_counts"]["language_output_tokens"], "avg_language_output_tokens_per_item": total["token_counts"]["language_output_tokens"]/n, } } for m, ms in total["module_inference_times"].items(): final["inference_times"]["module_total_times_seconds"][m] = ms["total"] final["inference_times"]["module_calls_count"][m] = ms["count"] final["inference_times"]["module_average_times_per_call"][m] = (ms["total"]/ms["count"]) if ms["count"]>0 else 0.0 with open(out_path, 'w', encoding='utf-8') as f: json.dump(final, f, ensure_ascii=False, indent=4) print_master(f"[{task_name}] {side_name} stats saved: {out_path}") # 编码 QRY if (not os.path.exists(q_path)) or (not os.path.exists(info_path)): print_master(f"[{tag}] Encode queries...") coll_q = MultimodalEvalDataCollator(processor, model_args, data_args, "qry") loader_q = DataLoader(eval_qry_dataset, batch_size=training_args.per_device_eval_batch_size, collate_fn=coll_q, num_workers=training_args.dataloader_num_workers) q_embeds, q_infos, q_stats, q_masks = encode_embeddings( model, loader_q, training_args, model_args, full_qry_dataset, encode_side="qry", zip_cfg=zip_cfg, aop_cfg=aop_cfg, description=f"Queries[{tag}] {dataset_name}" ) q_embeds = q_embeds[:len(full_qry_dataset)] q_infos = q_infos[:len(full_qry_dataset)] q_masks = q_masks[:len(full_qry_dataset)] q_total = _init_total() for bs in q_stats: _acc(q_total, bs) if local_rank == 0: with open(q_path, 'wb') as f: pickle.dump(q_embeds, f) if not os.path.exists(info_path): with open(info_path, 'w') as f: for info in q_infos: f.write(json.dumps(info) + '\n') with open(q_masks_path, 'w', encoding='utf-8') as f: for i, m in enumerate(q_masks): f.write(json.dumps({"index": i, "mask": m}, ensure_ascii=False) + "\n") _finalize(q_total, q_stats_path, dataset_name, f"query[{tag}]") if dist.is_initialized(): dist.barrier() # 编码 CAND if not os.path.exists(c_path): print_master(f"[{tag}] Encode candidates...") coll_c = MultimodalEvalDataCollator(processor, model_args, data_args, "cand") loader_c = DataLoader(eval_cand_dataset, batch_size=training_args.per_device_eval_batch_size, collate_fn=coll_c, num_workers=training_args.dataloader_num_workers) c_embeds, c_ids, c_stats, c_masks = encode_embeddings( model, loader_c, training_args, model_args, full_cand_dataset, encode_side="cand", zip_cfg=zip_cfg, aop_cfg=aop_cfg, description=f"Cands[{tag}] {dataset_name}" ) c_embeds = c_embeds[:len(full_cand_dataset)] c_ids = c_ids[:len(full_cand_dataset)] c_masks = c_masks[:len(full_cand_dataset)] c_total = _init_total() for bs in c_stats: _acc(c_total, bs) if local_rank == 0: cdict = {cid: emb for cid, emb in zip(c_ids, c_embeds)} with open(c_path, 'wb') as f: pickle.dump(cdict, f) with open(c_masks_path, 'w', encoding='utf-8') as f: for cid, m in zip(c_ids, c_masks): f.write(json.dumps({"cand_id": str(cid), "mask": m}, ensure_ascii=False) + "\n") _finalize(c_total, c_stats_path, dataset_name, f"cand[{tag}]") if dist.is_initialized(): dist.barrier() # 评分（同你原逻辑） if local_rank == 0: info_path = os.path.join(data_args.encode_output_path, f"{dataset_name}_info.jsonl") gt_infos = [json.loads(l) for l in open(info_path)] rank_global = (dataset_configs[dataset_name].get("eval_type", "global") == "global") metrics_to_report = dataset_configs[dataset_name].get("metrics", ["hit","ndcg","precision","recall","f1","map","mrr"]) for keep_layers in layers_to_eval: tag = f"layer{keep_layers}" if keep_layers else "layerlast" q_path = os.path.join(data_args.encode_output_path, f"{dataset_name}_qry_{tag}") c_path = os.path.join(data_args.encode_output_path, f"{dataset_name}_tgt_{tag}") with open(q_path, 'rb') as f: qry_embeds = pickle.load(f) with open(c_path, 'rb') as f: cand_embed_dict = pickle.load(f) pred_dicts=[] if rank_global: ck = list(cand_embed_dict.keys()) ce = np.stack([cand_embed_dict[k] for k in ck]) if isinstance(qry_embeds, np.ndarray) and qry_embeds.ndim==3: # late-interaction # 这里按需调用 processor.score（略），若需要可补上 sim = qry_embeds @ ce.T # 占位：如有自定义 late score 可以在此替换 else: sim = qry_embeds @ ce.T ranked = np.argsort(-sim, axis=1) for qid, gi in tqdm(enumerate(gt_infos), total=len(gt_infos), desc=f"[{tag}] scoring(all) {dataset_name}"): rid = ranked[qid] label = gi["label_name"] if isinstance(gi["label_name"], list) else [gi["label_name"]] pred_dicts.append({"prediction":[ck[i] for i in rid],"label":label,"rel_scores":gi.get("rel_scores")}) else: for qid, (qe, gi) in tqdm(enumerate(zip(qry_embeds, gt_infos)), total=len(gt_infos), desc=f"[{tag}] scoring(local) {dataset_name}"): cand_ids = gi["cand_names"] ce = np.stack([cand_embed_dict[k] for k in cand_ids]) if isinstance(qry_embeds, np.ndarray) and qry_embeds.ndim==3: sim_vec = qe @ ce.T else: sim_vec = qe @ ce.T rid = np.argsort(-sim_vec) label = gi["label_name"] if isinstance(gi["label_name"], list) else [gi["label_name"]] pred_dicts.append({"prediction":[cand_ids[i] for i in rid],"label":label,"rel_scores":gi.get("rel_scores")}) metrics = RankingMetrics(metrics_to_report) score = metrics.evaluate(pred_dicts) score["num_pred"] = len(pred_dicts); score["num_data"] = len(gt_infos) layer_score_path = os.path.join(data_args.encode_output_path, f"{dataset_name}_score_{tag}.json") layer_pred_path = os.path.join(data_args.encode_output_path, f"{dataset_name}_pred_{tag}.jsonl") with open(layer_score_path, "w") as f: json.dump(score, f, indent=4) with open(layer_pred_path, "w") as f: for p in pred_dicts: f.write(json.dumps(p) + '\n') print_master(f"[{dataset_name}] {tag} score: " + json.dumps({k: (f'{v:.4f}' if isinstance(v,(int,float)) else v) for k,v in score.items()})) if dist.is_initialized(): dist.barrier() if __name__ == '__main__': main()