eval_test_time_early_exit.py

# import datetime
# 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
# import transformers
# import math

# from torch.utils.data import DataLoader
# from tqdm import tqdm
# from transformers import HfArgumentParser, AutoConfig, AutoTokenizer
# from datasets import Dataset, concatenate_datasets
# from datasets.distributed import split_dataset_by_node
# from src.model.vlm_backbone.qwen2_vl.modeling_qwen2_vl_train_tokrnpooling import Qwen2VLForConditionalGeneration as _Qwen2VLForConditionalGeneration_src

# 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_add_text_cut import MMEBModel
# from src.model.processor import get_backbone_name, load_processor, COLPALI
# from src.utils import batch_to_device, print_rank, print_master

# from dataclasses import dataclass

# def get_env_mid_layer():
#     v = os.environ.get("MID_LM_LAYER", "").strip()
#     if v == "" or v.lower() in {"none", "null"}:
#         return None
#     try:
#         return int(v)
#     except:
#         logger.warning(f"Invalid MID_LM_LAYER={v}, ignore.")
#         return None

# # ------------- AOP-Prune config parsing -------------
# 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_aop_config():
#     """
#     从环境变量读取 AOP 剪裁配置。仅作为“驱动层”的简要测试开关；
#     实际剪裁逻辑在底模里（Qwen2-VLModel.forward）实现。
#     """
#     enabled     = _parse_bool(os.environ.get("AOP_ENABLED"), False)
#     apply_to    = os.environ.get("AOP_APPLY", "qry").strip().lower()  # qry|cand|both
#     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)

#     # 按类型控制
#     prune_vision = _parse_bool(os.environ.get("AOP_PRUNE_VISION"), True)
#     prune_text   = _parse_bool(os.environ.get("AOP_PRUNE_TEXT"), False)

#     delta_v     = _parse_float(os.environ.get("AOP_DELTA_VISION"), None)
#     khat_v      = _parse_float(os.environ.get("AOP_KHAT_VISION"), None)
#     keep_ratio_v= _parse_float(os.environ.get("AOP_KEEP_RATIO_VISION"), None)
#     min_keep_v  = _parse_int(os.environ.get("AOP_MIN_KEEP_VISION"), None)

#     delta_t     = _parse_float(os.environ.get("AOP_DELTA_TEXT"), None)
#     khat_t      = _parse_float(os.environ.get("AOP_KHAT_TEXT"), None)
#     keep_ratio_t= _parse_float(os.environ.get("AOP_KEEP_RATIO_TEXT"), None)
#     min_keep_t  = _parse_int(os.environ.get("AOP_MIN_KEEP_TEXT"), 32)

#     protect_text_last = _parse_int(os.environ.get("AOP_PROTECT_TEXT_LAST"), 16)
#     protect_special   = _parse_bool(os.environ.get("AOP_PROTECT_SPECIAL"), True)

#     margin_src  = os.environ.get("AOP_MARGIN", "").strip().lower()    # "" or "mid"
#     attn_impl   = os.environ.get("AOP_ATTN_IMPL", "").strip().lower() # "" or "sdpa"

#     if layer_idx is None and enabled:
#         logger.warning("AOP_ENABLED=1 但未设置 AOP_LAYER，关闭 AOP。"); enabled=False

#     # 新增：选择策略（aop | random）
#     selection = os.environ.get("AOP_SELECTION", "aop").strip().lower()
#     if _parse_bool(os.environ.get("AOP_RANDOM"), False):
#         selection = "random"
#     random_seed = _parse_int(os.environ.get("AOP_RANDOM_SEED"), None)

#     # 选择策略：aop | random | attention
#     selection = os.environ.get("AOP_SELECTION", "aop").strip().lower()
#     if _parse_bool(os.environ.get("AOP_RANDOM"), False):
#         selection = "random"
#     random_seed = _parse_int(os.environ.get("AOP_RANDOM_SEED"), None)
#     attn_agg = os.environ.get("AOP_ATTENTION_AGG", "mean").strip().lower()  # mean|max|sum

#     cfg = {
#         "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,

#         # 类型开关
#         "prune_vision": prune_vision,
#         "prune_text": prune_text,

#         # 视觉桶
#         "delta_vision": delta_v,
#         "K_hat_vision": khat_v,
#         "keep_ratio_vision": keep_ratio_v,
#         "min_keep_vision": min_keep_v,

#         # 文本桶
#         "delta_text": delta_t,
#         "K_hat_text": khat_t,
#         "keep_ratio_text": keep_ratio_t,
#         "min_keep_text": min_keep_t,

#         # 文本保护
#         "protect_text_last": protect_text_last,
#         "protect_special": protect_special,

#         # 可选：排名安全预算
#         "margin_mid": None if margin_src != "mid" else "USE_MID_MARGIN",
#         "epsilon_hat": None,
#         "attn_impl_override": attn_impl if attn_impl in {"sdpa"} else "",

#         # NEW: 选择策略
#         "selection": selection,          # "aop" 或 "random"
#         "random_seed": random_seed,      # 可选
#         "attn_agg": attn_agg,
#     }
#     return cfg

# def get_env_eval_layers():
#     """
#     解析环境变量 LM_LAYERS（优先）或兼容旧的 MID_LM_LAYER。
#     - LM_LAYERS 示例："4,8,12,last"；可包含 'last'/'none'/'null'/'-1' 表示最后一层(None)。
#     - 若未设置 LM_LAYERS，则回落到旧逻辑：MID_LM_LAYER=None -> [None]；否则 [mid, None]
#     返回: list[ int | None ]，例如 [4, 8, 12, None]；None 代表最后一层。
#     """
#     v = os.environ.get("LM_LAYERS", None)
#     if v is not None:
#         v = v.strip()

#     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)
#                     else:
#                         logger.warning(f"Ignoring non-positive layer '{tok}' in LM_LAYERS.")
#                 except Exception:
#                     logger.warning(f"Invalid token '{tok}' in LM_LAYERS; must be int or 'last'/'none'.")
#         # 去重但保持顺序
#         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)
#         if not uniq:
#             return [None]
#         return uniq
#     else:
#         # 兼容旧逻辑
#         mid = get_env_mid_layer()
#         return [None] if mid is None else [mid, None]

# # === Early-Exit config & helpers ===
# def get_env_ee_config():
#     ee_enabled = os.environ.get("EE_ENABLED", "0").strip().lower() in {"1","true","yes","on","y","t"}
#     layer = int(os.environ.get("EE_LAYER", os.environ.get("AOP_LAYER", "12")))  # 默认用 AOP_LAYER
#     method = os.environ.get("EE_METHOD", "margin").strip().lower()  # margin|p1p2|entropy|gini|combined
#     tau = float(os.environ.get("EE_TAU", "0.2"))
#     topk = int(os.environ.get("EE_TOPK", "1024"))
#     temp = float(os.environ.get("EE_TEMP", "0.05"))
#     save = os.environ.get("EE_SAVE", "1").strip().lower() in {"1","true","yes","on","y","t"}
#     combw = os.environ.get("EE_COMB_WEIGHTS", "1.0,0.5,0.5")
#     try:
#         w_margin, w_conf, w_sq = [float(x) for x in combw.split(",")]
#     except Exception:
#         w_margin, w_conf, w_sq = 1.0, 0.5, 0.5
#     return dict(
#         enabled=ee_enabled, layer=layer, method=method, tau=tau,
#         topk=topk, temp=temp, save=save,
#         w_margin=w_margin, w_conf=w_conf, w_sq=w_sq
#     )

# def _softmax_np(x: np.ndarray, temp: float = 1.0) -> np.ndarray:
#     x = x - np.max(x)
#     ex = np.exp(x / max(1e-6, temp))
#     s = np.sum(ex)
#     return ex / max(s, 1e-12)

# def confidence_from_topk(scores: np.ndarray, method="margin", temp=0.05, w_margin=1.0, w_conf=0.5, w_sq=0.5) -> float:
#     # scores: 已按降序排列（topK）
#     if scores.size == 0:
#         return 0.0
#     if scores.size == 1:
#         return 1e9
#     margin = float(scores[0] - scores[1])
#     p = _softmax_np(scores, temp=temp)
#     p1p2 = float(p[0] - p[1])
#     H = - float(np.sum(p * np.log(p + 1e-12))) / np.log(len(p))   # 归一化熵 ∈ [0,1]
#     conf = 1.0 - H
#     sqsum = float(np.sum(p**2))  # Gini 的等价度量（越大越集中）
#     if method == "margin":   return margin
#     if method == "p1p2":     return p1p2
#     if method == "entropy":  return conf
#     if method == "gini":     return sqsum
#     # combined
#     return w_margin*margin + w_conf*conf + w_sq*sqsum

# def run_early_exit_queries(
#     model: MMEBModel,
#     processor,
#     model_args: ModelArguments,
#     data_args: DataArguments,
#     training_args: TrainingArguments,
#     qry_dataset: Dataset,
#     cand_mid_dict: dict,
#     cand_last_dict: dict,
#     ee_cfg: dict,
#     dataset_name: str,
#     out_dir: str,
#     global_ranking: bool = True,
# ):
#     device = training_args.device
#     local_rank = dist.get_rank() if dist.is_initialized() else 0
#     is_main = (not dist.is_initialized()) or (local_rank == 0)

#     # 候选矩阵 -> GPU（bfloat16）
#     cand_ids = list(cand_mid_dict.keys())
#     cand_id2row = {str(cid): i for i, cid in enumerate(cand_ids)}
#     # cand_mid = np.stack([cand_mid_dict[c] for c in cand_ids]).astype(np.float32)
#     # cand_last = np.stack([cand_last_dict[c] for c in cand_ids]).astype(np.float32)
#     # cand_mid_t  = torch.from_numpy(cand_mid).to(device=device, dtype=torch.bfloat16)
#     # cand_last_t = torch.from_numpy(cand_last).to(device=device, dtype=torch.bfloat16)
#     cand_mid = np.stack([cand_mid_dict[c] for c in cand_ids]).astype(np.float32)
#     cand_last = np.stack([cand_last_dict[c] for c in cand_ids]).astype(np.float32)

#     # 先搬 cand_mid 到 GPU；cand_last 延迟到真的需要续跑时再搬
#     cand_mid_t  = torch.from_numpy(cand_mid).to(device=device, dtype=torch.bfloat16)
#     cand_last_t = None  # NEW: 延迟到 need_idx>0 分支内首次使用时再构造

#     # DataLoader（仅 query）
#     collator = MultimodalEvalDataCollator(processor, model_args, data_args, "qry")
#     loader = DataLoader(
#         qry_dataset,
#         batch_size=training_args.per_device_eval_batch_size,
#         collate_fn=collator,
#         num_workers=training_args.dataloader_num_workers
#     )

#     pred_dicts = []
#     details = []

#     # AOP 按侧门控（仅对 query 生效）
#     aop_cfg = getattr(model.encoder, "aop_prune_config", None)
#     _orig_enabled = None
#     side_enable = True
#     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 == "qry")

#     # 门控用的 k（margin/p1p2 只需要 top2）
#     k_conf = 2
#     tau   = float(ee_cfg["tau"])
#     method= ee_cfg["method"]
#     temp  = float(ee_cfg["temp"])

#     start_time = time.time()
#     idx_global = 0

#     for inputs, infos in tqdm(loader, desc=f"[EE] {dataset_name}@L{ee_cfg['layer']} (rank {local_rank})", disable=local_rank>0):
#         inputs = batch_to_device(inputs, device)

#         # if isinstance(aop_cfg, dict) and aop_cfg:
#         #     aop_cfg["enabled"] = bool(_orig_enabled and side_enable)
#         #     setattr(model.encoder, "aop_prune_config", aop_cfg)

#         # # 1) 前半程：跑到中间层（stop_at_layer），跳过 logits
#         # with torch.no_grad(), torch.autocast(device_type="cuda", dtype=torch.bfloat16, enabled=True):
#         #     out_mid = model.encoder(
#         #         **inputs,
#         #         return_dict=True,
#         #         output_hidden_states=False,
#         #         stop_at_layer=int(ee_cfg["layer"]),
#         #         compute_lm_head=False,   # 关键：不算 logits
#         #     )
#         orig_cfg = None
#         if isinstance(aop_cfg, dict) and aop_cfg:
#             orig_cfg   = dict(aop_cfg)  # 备份原配置，mid 后恢复
#             aop_layer  = aop_cfg.get("layer_idx", None)
#             ee_layer   = int(ee_cfg["layer"])
#             apply_to   = aop_cfg.get("apply_to", "qry").strip().lower()
#             # 新规则：mid 阶段是否启用 AOP
#             aop_on_mid = bool(
#                 _orig_enabled and side_enable and
#                 (aop_layer is not None) and (aop_layer < ee_layer) and
#                 (apply_to in {"qry", "both"})
#             )
#             aop_cfg_mid = dict(aop_cfg)
#             aop_cfg_mid["enabled"] = aop_on_mid
#             setattr(model.encoder, "aop_prune_config", aop_cfg_mid)

#         # 1) 前半程：跑到中间层（stop_at_layer），跳过 logits
#         with torch.no_grad(), torch.autocast(device_type="cuda", dtype=torch.bfloat16, enabled=True):
#             out_mid = model.encoder(
#                 **inputs,
#                 return_dict=True,
#                 output_hidden_states=False,
#                 stop_at_layer=int(ee_cfg["layer"]),
#                 compute_lm_head=False,   # 不算 logits
#             )

#         # 恢复原始 AOP 配置（避免影响后续续跑逻辑）
#         if isinstance(orig_cfg, dict):
#             setattr(model.encoder, "aop_prune_config", orig_cfg)

#         # EOS 池化 -> GPU
#         hs_mid = getattr(out_mid, "last_hidden_state", None)
#         if hs_mid is None:
#             assert out_mid.hidden_states is not None and len(out_mid.hidden_states) > 0
#             hs_mid = out_mid.hidden_states[-1]
#         am_mid = getattr(out_mid, "attention_mask", None)
#         if am_mid is None:
#             am_mid = inputs.get("attention_mask", None)
#         if hasattr(am_mid, "device") and am_mid.device != hs_mid.device:
#             am_mid = am_mid.to(hs_mid.device)
#         reps_mid_t = model._pooling(hs_mid, am_mid).detach().to(device=device, dtype=torch.bfloat16)  # [B,D]

#         B = reps_mid_t.size(0)
#         use_local = (not global_ranking)

#         # 2) 门控：GPU上 top2 + 置信度
#         if not use_local:
#             # 全库：top2 即可
#             scores_t = reps_mid_t @ cand_mid_t.T                # [B, Nc]
#             vals_t, idxs_t = torch.topk(scores_t, k=min(k_conf, scores_t.size(1)), dim=1)  # [B,2]
#             p_t = torch.softmax(vals_t / max(temp, 1e-6), dim=1)
#             if vals_t.size(1) >= 2:
#                 margin_t = vals_t[:, 0] - vals_t[:, 1]
#                 p1p2_t   = p_t[:, 0] - p_t[:, 1]
#             else:
#                 margin_t = torch.full((B,), float("inf"), device=device, dtype=vals_t.dtype)
#                 p1p2_t   = torch.ones(B, device=device, dtype=vals_t.dtype)
#             H_t = -(p_t * (torch.log(p_t + 1e-12))).sum(dim=1) / math.log(max(vals_t.size(1),1))
#             conf_map = {"margin": margin_t, "p1p2": p1p2_t, "entropy": 1.0 - H_t, "gini": (p_t ** 2).sum(dim=1)}
#             confs_t = conf_map.get(method, margin_t)  # [B]
#             exit_mask = (confs_t >= tau).detach().cpu().numpy().astype(bool)
#         else:
#             confs = []
#             for b in range(B):
#                 cand_local = infos[b]["cand_names"]
#                 rows = [cand_id2row.get(str(cid), -1) for cid in cand_local]
#                 rows = [r for r in rows if r >= 0]
#                 if len(rows) == 0:
#                     confs.append(0.0); continue
#                 cmat_t = cand_mid_t[rows]                # [Nl, D]
#                 sv_t   = (reps_mid_t[b:b+1] @ cmat_t.T)[0]  # [Nl]
#                 k = 2 if sv_t.size(0) >= 2 else 1
#                 vals_t, _ = torch.topk(sv_t, k=k, dim=0)
#                 p_t = torch.softmax(vals_t / max(temp, 1e-6), dim=0)
#                 if k >= 2:
#                     margin = (vals_t[0] - vals_t[1]).item()
#                     p1p2   = (p_t[0] - p_t[1]).item()
#                 else:
#                     margin, p1p2 = float("inf"), 1.0
#                 H = (-(p_t * (torch.log(p_t + 1e-12))).sum() / math.log(max(k,1))).item()
#                 gini = ((p_t ** 2).sum()).item()
#                 d = {"margin": margin, "p1p2": p1p2, "entropy": 1.0 - H, "gini": gini}
#                 confs.append(d.get(method, margin))
#             exit_mask = (np.array(confs) >= tau)

#         # 3) 早停：直接 mid 排序（只在需要保存 details 时构建 topk 列表）
#         for j in np.where(exit_mask)[0].tolist():
#             if not use_local:
#                 scores_j = (reps_mid_t[j:j+1] @ cand_mid_t.T)[0]  # [Nc]
#                 order = torch.argsort(scores_j, dim=0, descending=True).detach().cpu().numpy()
#                 cids = [cand_ids[i] for i in order]
#             else:
#                 cand_local = infos[j]["cand_names"]
#                 rows = [cand_id2row.get(str(cid), -1) for cid in cand_local]
#                 rows = [r for r in rows if r >= 0]
#                 if len(rows) == 0:
#                     cids = []
#                 else:
#                     cmat_t = cand_mid_t[rows]
#                     vec = (reps_mid_t[j:j+1] @ cmat_t.T)[0]
#                     order_local = torch.argsort(vec, dim=0, descending=True).detach().cpu().numpy()
#                     cids = [str(cand_local[i]) for i in order_local]
#             rel_docids = infos[j]["label_name"] if isinstance(infos[j]["label_name"], list) else [infos[j]["label_name"]]
#             pred_dicts.append({"prediction": cids, "label": rel_docids, "rel_scores": None})

#         # 4) 续跑：仅对未早停子集，从中间态继续到 last（跳过 logits）
#         need_idx = np.where(~exit_mask)[0].tolist()
#         if len(need_idx) > 0:
#             if cand_last_t is None:
#                 cand_last_t = torch.from_numpy(cand_last).to(device=device, dtype=torch.bfloat16)
#             if isinstance(aop_cfg, dict) and aop_cfg:
#                 aop_resume = dict(aop_cfg)
#                 aop_resume["enabled"] = bool(_orig_enabled and side_enable)
#                 setattr(model.encoder, "aop_prune_config", aop_resume)
#             interm = getattr(out_mid, "intermediate_state", None)
#             assert interm is not None, "Model must return intermediate_state when stop_at_layer is set."
#             hs = interm["hidden_states"].detach()
#             am = interm["attention_mask"].detach()
#             pos = interm["position_ids"].detach()
#             vm = interm.get("vision_mask", None)
#             tm = interm.get("text_mask", None)
#             next_layer = int(interm["next_layer_idx"])

#             hs_sub = hs[need_idx]
#             am_sub = am[need_idx]
#             pos_sub = pos[:, need_idx, :]
#             vm_sub = vm[need_idx] if vm is not None else None
#             tm_sub = tm[need_idx] if tm is not None else None
#             resume_state = {
#                 "hidden_states": hs_sub,
#                 "attention_mask": am_sub,
#                 "position_ids": pos_sub,
#                 "vision_mask": vm_sub,
#                 "text_mask": tm_sub,
#                 "next_layer_idx": next_layer,
#             }

#             with torch.no_grad(), torch.autocast(device_type="cuda", dtype=torch.bfloat16, enabled=True):
#                 out_last = model.encoder(
#                     return_dict=True,
#                     output_hidden_states=False,
#                     stop_at_layer=None,
#                     resume_state=resume_state,
#                     compute_lm_head=False,  # 关键：不算 logits
#                 )

#             hs_last = getattr(out_last, "last_hidden_state", None)
#             if hs_last is None:
#                 assert out_last.hidden_states is not None and len(out_last.hidden_states) > 0
#                 hs_last = out_last.hidden_states[-1]
#             am_last = getattr(out_last, "attention_mask", None)
#             if am_last is None:
#                 am_last = am_sub
#             if hasattr(am_last, "device") and am_last.device != hs_last.device:
#                 am_last = am_last.to(hs_last.device)
#             reps_last_t = model._pooling(hs_last, am_last).detach().to(device=device, dtype=torch.bfloat16)

#             if not use_local:
#                 scores_last_t = reps_last_t @ cand_last_t.T
#                 order_t = torch.argsort(scores_last_t, dim=1, descending=True)
#                 for k, j in enumerate(need_idx):
#                     order = order_t[k].detach().cpu().numpy()
#                     cids = [cand_ids[i] for i in order]
#                     rel_docids = infos[j]["label_name"] if isinstance(infos[j]["label_name"], list) else [infos[j]["label_name"]]
#                     pred_dicts.append({"prediction": cids, "label": rel_docids, "rel_scores": None})
#             else:
#                 for k, j in enumerate(need_idx):
#                     cand_local = infos[j]["cand_names"]
#                     rows = [cand_id2row.get(str(cid), -1) for cid in cand_local]
#                     rows = [r for r in rows if r >= 0]
#                     if len(rows) == 0:
#                         cids = []
#                     else:
#                         cmat_last_t = cand_last_t[rows]
#                         vec_t = (reps_last_t[k:k+1] @ cmat_last_t.T)[0]
#                         order_local = torch.argsort(vec_t, dim=0, descending=True).detach().cpu().numpy()
#                         cids = [str(cand_local[i]) for i in order_local]
#                     rel_docids = infos[j]["label_name"] if isinstance(infos[j]["label_name"], list) else [infos[j]["label_name"]]
#                     pred_dicts.append({"prediction": cids, "label": rel_docids, "rel_scores": None})

#         idx_global += B

#     # 评测并保存
#     metrics_to_report = ["hit", "ndcg", "precision", "recall", "f1", "map", "mrr"]
#     score = RankingMetrics(metrics_to_report).evaluate(pred_dicts)

#     if is_main:
#         os.makedirs(out_dir, exist_ok=True)
#         with open(os.path.join(out_dir, f"{dataset_name}_score_earlyexit.json"), "w") as f:
#             json.dump(score, f, indent=4)
#         # 建议测速时 EE_SAVE=0，不写 details
#         if ee_cfg.get("save", False):
#             with open(os.path.join(out_dir, f"{dataset_name}_pred_earlyexit.jsonl"), "w", encoding="utf-8") as f:
#                 for row in pred_dicts: f.write(json.dumps(row, ensure_ascii=False) + "\n")

#     elapsed = time.time() - start_time
#     return score, elapsed

# def make_layer_tag(keep_layers: int | None):
#     return f"layer{keep_layers}" if keep_layers and keep_layers > 0 else "layerlast"

# def dot_sim(a: np.ndarray, b: np.ndarray) -> np.ndarray:
#     # a: [Nq, D], b: [Nc, D], both L2-normalized already if normalize=true
#     return a @ b.T

# def build_score_details(qid: int, cand_ids: list, score_vec: np.ndarray, ranked_indices: np.ndarray):
#     return {
#         "qid": int(qid),
#         "cand_scores": [
#             {"cand_id": str(cand_ids[i]), "score": float(score_vec[i])}
#             for i in ranked_indices
#         ]
#     }

# def top1_top2_margin(score_vec: np.ndarray) -> float:
#     if len(score_vec) < 2:
#         return float("inf")  # 只有一个候选时视作极大margin
#     top2 = np.partition(score_vec, -2)[-2:]
#     top2.sort()
#     return float(top2[-1] - top2[-2])

# def simulate_early_exit_by_margin(
#     sims_mid: list[dict], sims_last: list[dict], labels: list[list[str]], metrics_to_report: list[str],
#     taus: list[float], rank_global: bool
# ):
#     """
#     sims_mid / sims_last: 每个query一个dict: {cand_id: score}
#     labels: 每个query的正样本cand_id列表
#     返回：不同tau下的覆盖率、指标
#     """
#     assert len(sims_mid) == len(sims_last) == len(labels)
#     N = len(labels)
#     results = []

#     from src.eval_utils.metrics import RankingMetrics
#     metrics = RankingMetrics(metrics_to_report)

#     # 预构造 用于metrics.evaluate 的pred_dict
#     def to_pred_dicts(use_mid_mask: list[bool]) -> list[dict]:
#         pred_dicts = []
#         for qid in range(N):
#             sims_use = sims_mid[qid] if use_mid_mask[qid] else sims_last[qid]
#             # 排序
#             ranked = sorted(sims_use.items(), key=lambda x: -x[1])
#             pred_dicts.append({
#                 "prediction": [cid for cid, _ in ranked],
#                 "label": labels[qid],
#                 "rel_scores": None
#             })
#         return pred_dicts

#     # 计算中间层margin
#     margins = []
#     for qid in range(N):
#         # 取前两大分数的margin
#         if len(sims_mid[qid]) == 0:
#             margins.append(0.0)
#             continue
#         scores = np.array(list(sims_mid[qid].values()), dtype=np.float32)
#         margins.append(top1_top2_margin(scores))

#     margins = np.array(margins, dtype=np.float32)

#     for tau in taus:
#         use_mid_mask = (margins >= tau).tolist()
#         pred_dicts = to_pred_dicts(use_mid_mask)
#         score_dict = metrics.evaluate(pred_dicts)
#         coverage = float(np.mean(use_mid_mask))  # 早停覆盖率
#         results.append({
#             "tau": tau,
#             "coverage": coverage,
#             **score_dict
#         })
#     return results

# def top1_top2_margin_from_array(score_vec: np.ndarray) -> float:
#     if score_vec is None or len(score_vec) == 0:
#         return 0.0
#     if len(score_vec) == 1:
#         return float('inf')
#     # 取前两大
#     top2 = np.partition(score_vec, -2)[-2:]
#     top2.sort()
#     return float(top2[-1] - top2[-2])

# logging.basicConfig(level=logging.INFO, format='[%(asctime)s] %(levelname)s [%(name)s:%(lineno)s] %(message)s')
# logger = logging.getLogger(__name__)

# # --- Global Dictionaries for Hooks (will be cleared before each encode_embeddings call) ---
# timing_info = {}
# token_info = {
#     "vision_tokens": 0,
#     "text_input_tokens": 0,  # Refers to the original text token count
#     "text_output_tokens": 0, # Not directly applicable here as we are encoding, not generating. Will be 0.
#     "total_llm_input_tokens": 0, # Refers to the total tokens LLM receives (visual + formatted text)
# }

# # --- Hook Functions Definition ---
# 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:
#         # print(f"Warning: No pre-hook data for module {module.__class__.__name__} ({module_id})")
#         return

#     timing_info[module_id].append((time.time(), 'post', module.__class__.__name__))

#     # Collect vision token count (only from Vision Transformer module's post hook)
#     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] # For visual features, usually (batch_size, num_tokens, hidden_dim)
#         elif hasattr(output, 'last_hidden_state'):
#             token_info["vision_tokens"] = output.last_hidden_state.shape[1]


# def register_model_hooks(model):
#     registered_modules = []
    
#     core_model = model.encoder if hasattr(model, "encoder") and model.encoder is not None else model
        
#     # Vision module
#     if hasattr(core_model, 'visual') and core_model.visual is not None:
#         vision_module = core_model.visual
#         vision_module.register_forward_pre_hook(timing_pre_hook)
#         vision_module.register_forward_hook(timing_post_hook)
#         registered_modules.append(vision_module)
#         print_master(f"Registered hooks for vision module: {vision_module.__class__.__name__}")
#     else:
#         print_master(f"WARNING: No 'visual' attribute found on core_model ({type(core_model)}).")


#     # Merger module (if inside visual) - it's part of the vision component
#     if hasattr(core_model, 'visual') and hasattr(core_model.visual, 'merger') and core_model.visual.merger is not None:
#         merger_module = core_model.visual.merger
#         merger_module.register_forward_pre_hook(timing_pre_hook)
#         merger_module.register_forward_hook(timing_post_hook)
#         registered_modules.append(merger_module)
#         print_master(f"Registered hooks for merger module: {merger_module.__class__.__name__}")
#     else:
#         print_master(f"WARNING: No 'merger' attribute found on core_model.visual ({type(getattr(core_model, 'visual', 'N/A'))}).")

#     # Language model body
#     if hasattr(core_model, 'model') and core_model.model is not None:
#         llm_main_module = core_model.model
#         llm_main_module.register_forward_pre_hook(timing_pre_hook)
#         llm_main_module.register_forward_hook(timing_post_hook)
#         registered_modules.append(llm_main_module)
#         print_master(f"Registered hooks for LLM main module: {llm_main_module.__class__.__name__}")
#     else:
#         print_master(f"WARNING: No 'model' attribute found on core_model ({type(core_model)}).")


#     # LM Head
#     if hasattr(core_model, 'lm_head') and core_model.lm_head is not None:
#         lm_head_module = core_model.lm_head
#         lm_head_module.register_forward_pre_hook(timing_pre_hook)
#         lm_head_module.register_forward_hook(timing_post_hook)
#         registered_modules.append(lm_head_module)
#         print_master(f"Registered hooks for LM head module: {lm_head_module.__class__.__name__}")
#     else:
#         print_master(f"WARNING: No 'lm_head' attribute found on core_model ({type(core_model)}).")


#     if not registered_modules:
#         print_master("Warning: No major modules found for hook registration. Check model architecture.")
#     return registered_modules


# def pad_dataset_to_divisible(dataset, world_size):
#     num_samples = len(dataset)
#     if num_samples % world_size == 0:
#         return dataset, num_samples

#     num_to_add = world_size - (num_samples % world_size)
#     padded_size = num_samples + num_to_add

#     padding_data = dataset.select([i % len(dataset) for i in range(num_to_add)])
#     padded_dataset = concatenate_datasets([dataset, padding_data])
#     return padded_dataset, padded_size

# def encode_embeddings(
#     model: MMEBModel,
#     loader: DataLoader,
#     training_args: TrainingArguments,
#     model_args: ModelArguments,
#     full_dataset: Dataset,
#     encode_side: str,
#     description: str = "Encoding"
# ) -> tuple[np.ndarray, list, list, list]:  # CHANGED: + list for img_token_masks
#     """
#     Encodes embeddings for a given dataset using the model, handling both standard and
#     late-interaction models in a DDP-safe manner.
#     Returns:
#         - embeddings: np.ndarray
#         - infos_or_ids: list
#         - batch_stats_list: list
#         - img_token_masks: list[None | list[bool]]  # NEW
#     """
#     local_rank = dist.get_rank() if dist.is_initialized() else 0
#     world_size = dist.get_world_size() if dist.is_initialized() else 1

#     # Check if the model is a late-interaction type
#     is_late_interaction = (model_args.model_backbone == COLPALI)

#     local_embeds = []
#     local_gt_infos = []
#     local_max_len = 0
    
#     # --- New: List to store statistics for each batch ---
#     batch_stats_list = []

#     # --- NEW: Collect masks ---
#     local_img_token_masks = []   # post image mask per sample
#     local_txt_token_masks = []   # NEW: post text mask per sample
#     local_post_attn_masks = []   # NEW: post attention_mask per sample (after prune, 1/0)

#     # --- NEW: per-sample token reduction records ---
#     local_token_records = []     # 每条样本一个 dict，含 pre/post/delta 数量

#     model.eval()

#     # Register hooks for the model once per encode_embeddings call
#     registered_hooks = register_model_hooks(model)

#     # --- NEW: helpers to取mask并序列化 ---
#     def _search_key(obj, key: str):
#         # 递归搜索 dict/list/tuple，找到指定 key
#         if isinstance(obj, dict):
#             if key in obj:
#                 return obj[key]
#             for v in obj.values():
#                 r = _search_key(v, key)
#                 if r is not None:
#                     return r
#         elif isinstance(obj, (list, tuple)):
#             for v in obj:
#                 r = _search_key(v, key)
#                 if r is not None:
#                     return r
#         return None

#     def _to_serializable_mask_list(mask_list, batch_size: int):
#         # 将模型返回的 mask（list/tensor/ndarray/None）转成 [None | list[bool]] * B
#         if mask_list is None:
#             return [None] * batch_size

#         out = []
#         if isinstance(mask_list, (list, tuple)):
#             for m in mask_list:
#                 if m is None:
#                     out.append(None)
#                 elif torch.is_tensor(m):
#                     out.append(m.detach().cpu().tolist())
#                 elif isinstance(m, np.ndarray):
#                     out.append(m.tolist())
#                 else:
#                     # already python list/bool
#                     out.append(m)
#         elif torch.is_tensor(mask_list):
#             # 若是 2D 张量（B, L），直接 tolist() -> list[list[bool/int]]
#             out = mask_list.detach().cpu().tolist()
#         elif isinstance(mask_list, np.ndarray):
#             out = mask_list.tolist()
#         else:
#             # 未知类型，保守返回 None 占位
#             out = [None] * batch_size

#         # 长度对齐 batch_size
#         if isinstance(out, list):
#             if len(out) < batch_size:
#                 out = out + [None] * (batch_size - len(out))
#             elif len(out) > batch_size:
#                 out = out[:batch_size]
#         return out
    
#     def _to_bool_lists(m, batch_size: int):
#         lst = _to_serializable_mask_list(m, batch_size)
#         # 归一化成 list[ list[bool] | None ]
#         out = []
#         for x in lst:
#             if x is None:
#                 out.append(None)
#             else:
#                 # x 可能是 list[int] 或 list[bool]
#                 out.append([bool(int(v)) for v in x])
#         return out

#     with torch.no_grad():
#         for inputs, dataset_info in tqdm(loader, desc=f"{description} (rank {local_rank})", disable=local_rank > 0):
#             # --- Reset statistics for each inference pass ---
#             timing_info.clear()
#             token_info["vision_tokens"] = 0
#             token_info["text_input_tokens"] = 0
#             token_info["text_output_tokens"] = 0
#             token_info["total_llm_input_tokens"] = 0

#             inputs = batch_to_device(inputs, training_args.device)
#             current_batch_size = inputs['input_ids'].shape[0] if 'input_ids' in inputs and inputs['input_ids'] is not None else 1

#             with torch.autocast(enabled=True, dtype=torch.bfloat16, device_type="cuda"):
#                 start_inference_time = time.time()
#                 # ---- NEW: 按侧开/关 AOP ----
#                 aop_cfg = getattr(model.encoder, "aop_prune_config", None)
#                 _orig_enabled = None
#                 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)
#                     aop_cfg["enabled"] = bool(side_enable and _orig_enabled)
#                     setattr(model.encoder, "aop_prune_config", aop_cfg)
#                 if encode_side == "qry":
#                     output = model(qry=inputs)
#                     reps = output["qry_reps"].detach()
#                     local_gt_infos.extend(dataset_info)
#                 else:
#                     output = model(tgt=inputs)
#                     reps = output["tgt_reps"].detach()
#                     local_gt_infos.extend([info["cand_name"] for info in dataset_info])
#                 # ---- NEW: 恢复 enabled（避免影响下个 encode_side）----
#                 if isinstance(aop_cfg, dict) and _orig_enabled is not None:
#                     aop_cfg["enabled"] = _orig_enabled
#                     setattr(model.encoder, "aop_prune_config", aop_cfg)
#                 end_inference_time = time.time()

#             # --- NEW: 提取 post-prune 的 image/text 掩码 与 post attention_mask ---
#             img_masks_raw = None
#             txt_masks_raw = None
#             post_attn_raw = None

#             if isinstance(output, dict):
#                 img_masks_raw = _search_key(output, "image_token_bool_masks")
#                 txt_masks_raw = _search_key(output, "text_token_bool_masks")     # NEW
#                 post_attn_raw = _search_key(output, "post_attention_mask")       # NEW（我们的 MMEBModel.forward 里带了这个键）

#             # 兼容：若挂在 model 上
#             if img_masks_raw is None and hasattr(model, "image_token_bool_masks"):
#                 img_masks_raw = getattr(model, "image_token_bool_masks")
#             if txt_masks_raw is None and hasattr(model, "text_token_bool_masks"):
#                 txt_masks_raw = getattr(model, "text_token_bool_masks")
#             if post_attn_raw is None and hasattr(model, "post_attention_mask"):
#                 post_attn_raw = getattr(model, "post_attention_mask")

#             img_masks_serializable = _to_serializable_mask_list(img_masks_raw, current_batch_size)
#             txt_masks_serializable = _to_serializable_mask_list(txt_masks_raw, current_batch_size)     # NEW
#             post_attn_serializable = _to_serializable_mask_list(post_attn_raw, current_batch_size)     # NEW

#             local_img_token_masks.extend(img_masks_serializable)
#             local_txt_token_masks.extend(txt_masks_serializable)        # NEW
#             local_post_attn_masks.extend(post_attn_serializable)        # NEW

#             # --- NEW: 计算本 batch 的 pre/post/delta 数量并累计 ---
#             cfg = getattr(model.encoder, "config", None)
#             # pre masks 来自 inputs（删前）
#             input_ids = inputs.get("input_ids", None)
#             attn2d_pre = inputs.get("attention_mask", None)
#             if input_ids is None or attn2d_pre is None or cfg is None:
#                 # 无法统计，留空
#                 pre_vis_counts = [0] * current_batch_size
#                 pre_txt_counts = [0] * current_batch_size
#                 pre_tot_counts = [0] * current_batch_size
#             else:
#                 iid = input_ids
#                 am = attn2d_pre.to(torch.bool)
#                 image_token_id = getattr(cfg, "image_token_id", None)
#                 video_token_id = getattr(cfg, "video_token_id", None)
#                 bos_id = getattr(cfg, "bos_token_id", None)
#                 eos_id = getattr(cfg, "eos_token_id", None)
#                 pad_id = getattr(cfg, "pad_token_id", None)

#                 is_image = (iid == image_token_id) if (image_token_id is not None and image_token_id >= 0) else torch.zeros_like(iid, dtype=torch.bool)
#                 is_video = (iid == video_token_id) if (video_token_id is not None and video_token_id >= 0) else torch.zeros_like(iid, dtype=torch.bool)
#                 is_vision = is_image | is_video

#                 is_special = torch.zeros_like(iid, dtype=torch.bool)
#                 for tid in [bos_id, eos_id, pad_id]:
#                     if tid is not None and tid >= 0:
#                         is_special |= (iid == tid)

#                 pre_txt_mask = am & (~is_vision) & (~is_special)
#                 pre_vis_mask = am & is_vision

#                 pre_vis_counts = pre_vis_mask.sum(dim=1).tolist()
#                 pre_txt_counts = pre_txt_mask.sum(dim=1).tolist()
#                 pre_tot_counts = am.sum(dim=1).tolist()

#             # post masks（删后）来自模型输出；与 post_attn 做与运算
#             post_text_masks = _to_bool_lists(txt_masks_raw, current_batch_size)      # list[ list[bool] | None ]
#             post_image_masks = _to_bool_lists(img_masks_raw, current_batch_size)
#             post_attn_masks  = _to_bool_lists(post_attn_raw, current_batch_size)

#             sum_pre_text = 0; sum_post_text = 0
#             sum_pre_vis  = 0; sum_post_vis  = 0
#             sum_pre_tot  = 0; sum_post_tot  = 0

#             for i in range(current_batch_size):
#                 pre_text = int(pre_txt_counts[i]) if i < len(pre_txt_counts) else 0
#                 pre_vis  = int(pre_vis_counts[i]) if i < len(pre_vis_counts) else 0
#                 pre_tot  = int(pre_tot_counts[i]) if i < len(pre_tot_counts) else 0

#                 # post 计数：mask 可能为 None
#                 m_text = post_text_masks[i] if post_text_masks is not None and i < len(post_text_masks) else None
#                 m_img  = post_image_masks[i] if post_image_masks is not None and i < len(post_image_masks) else None
#                 m_attn = post_attn_masks[i]  if post_attn_masks is not None  and i < len(post_attn_masks)  else None

#                 if m_attn is None:
#                     post_text = 0; post_vis = 0; post_tot = 0
#                 else:
#                     # 与 attention_mask 后统计 True 的数
#                     if m_text is not None:
#                         post_text = sum(1 for a, t in zip(m_attn, m_text) if a and t)
#                     else:
#                         post_text = 0
#                     if m_img is not None:
#                         post_vis = sum(1 for a, v in zip(m_attn, m_img) if a and v)
#                     else:
#                         post_vis = 0
#                     post_tot = sum(1 for a in m_attn if a)

#                 # 累计 batch 级
#                 sum_pre_text += pre_text; sum_post_text += post_text
#                 sum_pre_vis  += pre_vis;  sum_post_vis  += post_vis
#                 sum_pre_tot  += pre_tot;  sum_post_tot  += post_tot

#                 # 保存 per-sample 记录（用于 JSONL）
#                 local_token_records.append({
#                     "side": encode_side,
#                     "pre":  {"text": pre_text,  "vision": pre_vis,  "total": pre_tot},
#                     "post": {"text": post_text, "vision": post_vis, "total": post_tot},
#                     "delta":{"text": pre_text - post_text, "vision": pre_vis - post_vis, "total": pre_tot - post_tot},
#                 })

#             # --- Update total LLM input tokens after the model call ---
#             if 'input_ids' in inputs and inputs['input_ids'] is not None:
#                 token_info["total_llm_input_tokens"] = inputs['input_ids'].shape[1]
#                 token_info["text_input_tokens"] = token_info["total_llm_input_tokens"] - token_info["vision_tokens"]
#                 token_info["text_input_tokens"] = max(0, token_info["text_input_tokens"])

#             # --- Collect and Store Batch Statistics ---
#             batch_inference_time = end_inference_time - start_inference_time
            
#             current_batch_stats = {
#                 "batch_size": current_batch_size,
#                 "total_inference_time_seconds": batch_inference_time,
#                 "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"],
#                 }
#             }
#             current_batch_stats["token_reduction"] = {
#                 "sum_pre_text":  sum_pre_text,
#                 "sum_post_text": sum_post_text,
#                 "sum_pre_vision":  sum_pre_vis,
#                 "sum_post_vision": sum_post_vis,
#                 "sum_pre_total": sum_pre_tot,
#                 "sum_post_total": sum_post_tot,
#             }

#             # Calculate and store module timings for the current batch
#             for module_obj in registered_hooks:
#                 module_id = id(module_obj)
#                 module_name = module_obj.__class__.__name__
#                 times = timing_info.get(module_id, [])
#                 durations = []
#                 pre_times = {} 
#                 for t, event_type, _ in times:
#                     if event_type == 'pre':
#                         pre_times[module_id] = t
#                     elif event_type == 'post' and module_id in pre_times:
#                         duration = t - pre_times.pop(module_id)
#                         durations.append(duration)
                
#                 if durations:
#                     current_batch_stats["module_inference_times"][module_name] = {
#                         "total": sum(durations),
#                         "count": len(durations),
#                         "avg": sum(durations) / len(durations)
#                     }
#                 else:
#                     current_batch_stats["module_inference_times"][module_name] = {
#                         "total": 0.0,
#                         "count": 0,
#                         "avg": 0.0
#                     }
            
#             batch_stats_list.append(current_batch_stats)

#             # --- Debug prints (optional) ---
#             print_rank(f"\n--- Inference Statistics for {encode_side} batch (Rank {local_rank}) ---")
#             print_rank(f"Batch Inference took: {batch_inference_time:.4f} seconds")
#             print_rank("--- Module Inference Timing Statistics ---")
#             for module_name, stats in current_batch_stats["module_inference_times"].items():
#                 print_rank(f"**{module_name}**: Total: {stats['total']:.6f}s, Count: {stats['count']}, Avg: {stats['avg']:.6f}s")
#             print_rank("--- Token Count Statistics ---")
#             print_rank(f"**视觉 token 数量**: {current_batch_stats['token_counts']['visual_tokens']}")
#             print_rank(f"**语言输入 token 数量 (仅原始文本)**: {current_batch_stats['token_counts']['language_input_tokens_raw']}")
#             print_rank(f"**LLM总输入 token 数量 (包含视觉 + 格式化文本)**: {current_batch_stats['token_counts']['llm_total_input_tokens']}")
#             print_rank(f"**语言输出 token 数量**: {current_batch_stats['token_counts']['language_output_tokens']}")

#             if is_late_interaction and reps.dim() == 3:
#                 local_max_len = max(local_max_len, reps.shape[1])

#             local_embeds.append(reps)

#     if not local_embeds:
#         # Handle cases where a rank gets no data
#         return np.array([]), [], [], []  # CHANGED: 4个返回值

#     # === DDP Synchronization and Padding for Late-Interaction Models ===
#     if is_late_interaction:
#         if dist.is_initialized():
#             # 1: global max length
#             local_max_len_tensor = torch.tensor(local_max_len, device=training_args.device)
#             dist.all_reduce(local_max_len_tensor, op=dist.ReduceOp.MAX)
#             global_max_len = local_max_len_tensor.item()
#         else:
#             global_max_len = local_max_len

#         # 2: pad to global max length
#         padded_embeds = []
#         for reps_batch in local_embeds:
#             if reps_batch.dim() == 3:
#                 B, L, H = reps_batch.shape
#                 padding_size = global_max_len - L
#                 padded_batch = F.pad(reps_batch, (0, 0, 0, padding_size), "constant", 0)
#                 padded_embeds.append(padded_batch)
#             else:
#                 padded_embeds.append(reps_batch)

#         embeds_tensor = torch.cat(padded_embeds, dim=0).contiguous()
#     else:
#         embeds_tensor = torch.cat(local_embeds, dim=0).contiguous()

#     # === Gather embeddings and keys from all ranks ===
#     if dist.is_initialized() and full_dataset.num_rows >= world_size:
#         print_master(f"Gathering {encode_side} embeddings across all ranks...")

#         # tensor gather
#         output_shape = list(embeds_tensor.shape)
#         output_shape[0] = full_dataset.num_rows
#         embeds_tensor = embeds_tensor.to(training_args.device)
#         gathered_embeds_tensor = torch.empty(output_shape, dtype=embeds_tensor.dtype, device=training_args.device)
#         dist.all_gather_into_tensor(gathered_embeds_tensor, embeds_tensor)
#         final_embeddings = gathered_embeds_tensor.cpu().float().numpy()

#         # object gather for infos and stats
#         gathered_gt_infos = [None for _ in range(world_size)]
#         dist.all_gather_object(gathered_gt_infos, local_gt_infos)
#         all_gt_infos = [key for rank_keys in gathered_gt_infos for key in rank_keys]

#         gathered_batch_stats = [None for _ in range(world_size)]
#         dist.all_gather_object(gathered_batch_stats, batch_stats_list)
#         all_batch_stats = [stats for rank_stats in gathered_batch_stats for stats in rank_stats]

#         # --- NEW: gather masks ---
#         gathered_masks = [None for _ in range(world_size)]
#         dist.all_gather_object(gathered_masks, local_img_token_masks)
#         all_img_token_masks = [m for rank_list in gathered_masks for m in rank_list]
#         # NEW: gather text masks
#         gathered_txt_masks = [None for _ in range(world_size)]
#         dist.all_gather_object(gathered_txt_masks, local_txt_token_masks)
#         all_txt_token_masks = [m for rank_list in gathered_txt_masks for m in rank_list]

#         # NEW: gather post attention masks（如需）
#         gathered_post_attn = [None for _ in range(world_size)]
#         dist.all_gather_object(gathered_post_attn, local_post_attn_masks)
#         all_post_attn_masks = [m for rank_list in gathered_post_attn for m in rank_list]

#         # NEW: gather token records
#         gathered_token_recs = [None for _ in range(world_size)]
#         dist.all_gather_object(gathered_token_recs, local_token_records)
#         all_token_records = [r for rank_list in gathered_token_recs for r in rank_list]
#     else:
#         all_gt_infos = local_gt_infos
#         final_embeddings = embeds_tensor.cpu().float().numpy()
#         all_batch_stats = batch_stats_list
#         all_img_token_masks = local_img_token_masks  # NEW
#         all_txt_token_masks = local_txt_token_masks
#         all_post_attn_masks = local_post_attn_masks
#         all_token_records = local_token_records

#     return final_embeddings, all_gt_infos, all_batch_stats, all_img_token_masks, all_txt_token_masks, all_token_records

# # === NEW: 一次前向同时导出 cand 的中间层和最后一层向量 ===
# def encode_candidates_both_layers(
#     model: MMEBModel,
#     loader: DataLoader,
#     training_args: TrainingArguments,
#     model_args: ModelArguments,
#     full_dataset: Dataset,
#     mid_layer: int,
# ) -> tuple[np.ndarray, np.ndarray, list]:
#     """
#     单次forward到最后一层，直接从 hidden_states 取：
#       - mid_hidden = hidden_states[mid_layer]   # 表示经过 mid_layer 层后的状态（见Qwen2_5_VLModel的all_hidden_states定义）
#       - last_hidden = hidden_states[-1]         # 最后一层norm后的状态
#     然后用 _pooling(attention_mask) 取句向量，返回：
#       - cand_mid_embeds: np.ndarray [Nc, D]
#       - cand_last_embeds: np.ndarray [Nc, D]
#       - cand_ids: list[str]
#     说明：
#       - cand 侧默认不做 AOP 剪枝（AOP_APPLY=qry 时天然关闭），因此 mid/last 的序列长度一致，可直接用原 attention_mask 做池化。
#     """
#     local_rank = dist.get_rank() if dist.is_initialized() else 0
#     model.eval()

#     all_mid = []
#     all_last = []
#     all_ids = []

#     with torch.no_grad():
#         for inputs, dataset_info in tqdm(loader, desc=f"Candidates[BOTH] (rank {local_rank})", disable=local_rank > 0):
#             inputs = batch_to_device(inputs, training_args.device)
#             # cand 侧确保不触发 AOP（如果你的 AOP_APPLY=qry/both，会在底模按侧门控；此处再做一次保险）
#             aop_cfg = getattr(model.encoder, "aop_prune_config", None)
#             _orig_enabled = None
#             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 == "cand")
#                 aop_cfg["enabled"] = bool(side_enable and _orig_enabled)
#                 setattr(model.encoder, "aop_prune_config", aop_cfg)

#             with torch.autocast(enabled=True, dtype=torch.bfloat16, device_type="cuda"):
#                 # 关键：一次forward拿全层的hidden_states
#                 out = model.encoder(
#                     **inputs,
#                     return_dict=True,
#                     output_hidden_states=True,  # 必须
#                     stop_at_layer=None,         # 走全层
#                 )

#             # 取 hidden_states 并索引中间层/最后一层
#             hs_list = out.hidden_states
#             assert hs_list is not None and len(hs_list) > mid_layer, \
#                 f"hidden_states is None or too short. Need index {mid_layer}, got len={0 if hs_list is None else len(hs_list)}"
#             mid_hs  = hs_list[mid_layer]      # [B, L, D]：等价“经过 mid_layer 层后的状态”（即 pre-layer(mid_layer+1)）
#             last_hs = hs_list[-1]             # [B, L, D]：最终norm后的状态

#             # 用原 attention_mask 池化（cand侧未剪枝）
#             am = inputs.get("attention_mask", None)
#             if am is not None and hasattr(am, "device"):
#                 if am.device != mid_hs.device:
#                     am = am.to(mid_hs.device)

#             reps_mid  = model._pooling(mid_hs, am)   # [B, D]
#             reps_last = model._pooling(last_hs, am)  # [B, D]

#             all_mid.append(reps_mid.detach().float().cpu())
#             all_last.append(reps_last.detach().float().cpu())
#             all_ids.extend([info["cand_name"] for info in dataset_info])

#             # 恢复 AOP 开关（避免影响其它侧）
#             if isinstance(aop_cfg, dict) and _orig_enabled is not None:
#                 aop_cfg["enabled"] = _orig_enabled
#                 setattr(model.encoder, "aop_prune_config", aop_cfg)

#     if not all_mid:
#         return np.array([]), np.array([]), []

#     cand_mid_embeds  = torch.cat(all_mid,  dim=0).numpy()
#     cand_last_embeds = torch.cat(all_last, dim=0).numpy()
#     return cand_mid_embeds, cand_last_embeds, all_ids

# def main():
#     # ----------------------- Distributed init -----------------------
#     if "RANK" in os.environ and dist.is_available() and not dist.is_initialized():
#         dist.init_process_group(backend="nccl", timeout=datetime.timedelta(minutes=60))
#     local_rank = dist.get_rank() if dist.is_initialized() else 0
#     world_size = dist.get_world_size() if dist.is_initialized() else 1

#     print_master("Distributed init debug info:")
#     print_master(f"RANK: {os.environ.get('RANK')}")
#     print_master(f"LOCAL_RANK: {os.environ.get('LOCAL_RANK')}")
#     print_master(f"WORLD_SIZE: {os.environ.get('WORLD_SIZE')}")
#     print_master(f"MASTER_ADDR: {os.environ.get('MASTER_ADDR')}")
#     print_master(f"MASTER_PORT: {os.environ.get('MASTER_PORT')}")
#     if dist.is_initialized():
#         print_rank(f"dist.get_rank(): {dist.get_rank()}")
#         print_rank(f"dist.get_world_size(): {dist.get_world_size()}")

#     # 兼容 torchrun 参数
#     for arg in sys.argv:
#         if arg.startswith("--local-rank="):
#             rank = arg.split("=")[1]
#             sys.argv.remove(arg)
#             sys.argv.append('--local_rank')
#             sys.argv.append(rank)

#     # ----------------------- Parse args -----------------------
#     parser = HfArgumentParser((ModelArguments, DataArguments, TrainingArguments))
#     model_args, data_args, training_args = parser.parse_args_into_dataclasses()
#     model_args: ModelArguments
#     data_args: DataArguments
#     training_args: TrainingArguments
#     os.makedirs(data_args.encode_output_path, exist_ok=True)

#     # 支持多层评测（优先 LM_LAYERS，兼容 MID_LM_LAYER）
#     layers_to_eval = get_env_eval_layers()
#     print_master(f"Eval layers (qry/tgt): {layers_to_eval}")

#     # ----------------------- Model loading -----------------------
#     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'Model Backbone: {model_args.model_backbone}')

#     # 仅 rank0 下载，其他rank等待缓存
#     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] Loading the model from Huggingface: {model_args.model_name}...")
#     if torch.distributed.is_initialized():
#         torch.distributed.barrier()
#     if local_rank != 0:
#         print_rank(f"Loading the model from cache...")
#         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)

#     # ---- NEW: AOP 剪裁配置注入（驱动底模里已实现的 AOP 逻辑）----
#     aop_cfg = get_env_aop_config()
#     if aop_cfg["enabled"]:
#         # 把配置塞到底模；底模 forward 中读取该 dict 并执行剪裁
#         setattr(model.encoder, "aop_prune_config", aop_cfg)
#         # 可选：为了便于在判定层取注意力或手算 qk，覆盖注意力实现
#         attn_override = aop_cfg.get("attn_impl_override", "")
#         if attn_override:
#             try:
#                 if 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] override attn impl: {prev} -> {attn_override} (仅测试建议)")
#             except Exception as e:
#                 print_master(f"[AOP] try override attn impl failed: {e}")

#         print_master("[AOP] AOP-Prune enabled with config: " + json.dumps({
#             "apply_to": aop_cfg["apply_to"],
#             "layer_idx": aop_cfg["layer_idx"],
#             "mode": aop_cfg["mode"],
#             "delta": aop_cfg["delta"],
#             "K_hat": aop_cfg["K_hat"],
#             "keep_ratio": aop_cfg["keep_ratio"],
#             "min_keep": aop_cfg["min_keep"],
#             "use_bias": aop_cfg["use_bias"],
#             "margin_mid?": (aop_cfg["margin_mid"] is not None),
#             "prune_text": aop_cfg.get("prune_text", False),
#             "keep_ratio_text": aop_cfg.get("keep_ratio_text", None),
#             "keep_ratio_vision": aop_cfg.get("keep_ratio_vision", None),
#             "selection": aop_cfg.get("selection", "aop"),
#             "attn_agg": aop_cfg.get("attn_agg", "mean"),
#         }))
#     else:
#         print_master("[AOP] disabled (set AOP_ENABLED=1 to enable)")
        
#     # 确保“最后一层”时不裁层（避免类里默认20层的坑）
#     model.set_inference_layers(qry_layers=None, tgt_layers=None)

#     with open(data_args.dataset_config, 'r') as yaml_file:
#         dataset_configs = yaml.safe_load(yaml_file)

#     # ----------------------- Main evaluation loop -----------------------
#     for dataset_idx, (dataset_name, task_config) in enumerate(dataset_configs.items()):
#         if dist.is_initialized():
#             dist.barrier()
#         print_master(f"\n--- Evaluating {dataset_name} ---")

#         # 根据 data_basedir 修正路径
#         if data_args.data_basedir is not None:
#             for key in ["image_root", "video_root", "frame_root", "clip_root", "data_path"]:
#                 if data_args.data_basedir and task_config.get(key):
#                     task_config[key] = os.path.join(data_args.data_basedir, task_config[key])

#         # 构建数据集
#         full_eval_qry_dataset, corpus = AutoEvalPairDataset.instantiate(model_args=model_args, data_args=data_args, **task_config)
#         full_eval_cand_dataset = generate_cand_dataset(full_eval_qry_dataset, corpus)
#         eval_qry_dataset, eval_cand_dataset = full_eval_qry_dataset, full_eval_cand_dataset

#         if dist.is_initialized():
#             world_size = dist.get_world_size()
#             padded_qry_dataset, _ = pad_dataset_to_divisible(full_eval_qry_dataset, world_size)
#             padded_cand_dataset, _ = pad_dataset_to_divisible(full_eval_cand_dataset, world_size)
#             eval_qry_dataset = split_dataset_by_node(padded_qry_dataset, rank=local_rank, world_size=world_size)
#             eval_cand_dataset = split_dataset_by_node(padded_cand_dataset, rank=local_rank, world_size=world_size)
#         else:
#             padded_qry_dataset, padded_cand_dataset = full_eval_qry_dataset, full_eval_cand_dataset


#         # === EE-only: 仅在线早停推理（先确保两份 candidate 向量）===
#         ee_cfg = get_env_ee_config()
#         assert ee_cfg["enabled"], "EE_ENABLED must be 1 for EE-only pipeline."

#         # 依据 EE_LAYER 构造 tag
#         mid_layer = int(ee_cfg["layer"])
#         mid_tag = make_layer_tag(mid_layer)     # e.g., layer12
#         last_tag = "layerlast"

#         # 准备路径
#         cand_mid_path  = os.path.join(data_args.encode_output_path, f"{dataset_name}_tgt_{mid_tag}")
#         cand_last_path = os.path.join(data_args.encode_output_path, f"{dataset_name}_tgt_{last_tag}")

#         # 构造 cand DataLoader（一次性，不切分）
#         eval_cand_collator = MultimodalEvalDataCollator(processor, model_args, data_args, "cand")
#         eval_cand_loader = DataLoader(
#             full_eval_cand_dataset,
#             batch_size=training_args.per_device_eval_batch_size,
#             collate_fn=eval_cand_collator,
#             num_workers=training_args.dataloader_num_workers
#         )

#         # === 替换为：一次前向，导出 cand 的 mid/last 两份向量 ===
#         need_mid  = (not os.path.exists(cand_mid_path))
#         need_last = (not os.path.exists(cand_last_path))

#         if need_mid or need_last:
#             print_master(f"[{dataset_name}] EE-only: encoding candidates BOTH layers in one pass (mid={mid_tag}, last={last_tag}) ...")
#             # 走全层（不截层）
#             model.set_inference_layers(qry_layers=None, tgt_layers=None)

#             cand_embeds_mid, cand_embeds_last, all_cand_ids = encode_candidates_both_layers(
#                 model=model,
#                 loader=eval_cand_loader,
#                 training_args=training_args,
#                 model_args=model_args,
#                 full_dataset=full_eval_cand_dataset,
#                 mid_layer=mid_layer,
#             )
#             if local_rank == 0:
#                 if need_mid:
#                     cand_embed_dict_mid = {cid: emb for cid, emb in zip(all_cand_ids, cand_embeds_mid)}
#                     with open(cand_mid_path, "wb") as f:
#                         pickle.dump(cand_embed_dict_mid, f)
#                     print_master(f"[{dataset_name}] EE-only: saved {mid_tag} candidate embeddings -> {cand_mid_path}")
#                 if need_last:
#                     cand_embed_dict_last = {cid: emb for cid, emb in zip(all_cand_ids, cand_embeds_last)}
#                     with open(cand_last_path, "wb") as f:
#                         pickle.dump(cand_embed_dict_last, f)
#                     print_master(f"[{dataset_name}] EE-only: saved {last_tag} candidate embeddings -> {cand_last_path}")
#         else:
#             print_master(f"[{dataset_name}] EE-only: reuse existing candidates (mid={cand_mid_path}, last={cand_last_path})")

#         if dist.is_initialized():
#             dist.barrier()

#         # 3) 在线早停门控 + 子集续跑（不做离线分层评分/曲线）
#         if local_rank == 0:
#             with open(cand_mid_path, "rb") as f:
#                 cand_mid_dict = pickle.load(f)
#             with open(cand_last_path, "rb") as f:
#                 cand_last_dict = pickle.load(f)

#             rank_global = task_config.get("eval_type", "global") == "global"
#             print_master(f"[{dataset_name}] Run ONLINE early-exit at layer={ee_cfg['layer']}, method={ee_cfg['method']}, tau={ee_cfg['tau']}, topk={ee_cfg['topk']}, global={rank_global}")
#             run_early_exit_queries(
#                 model=model,
#                 processor=processor,
#                 model_args=model_args,
#                 data_args=data_args,
#                 training_args=training_args,
#                 qry_dataset=full_eval_qry_dataset,   # 全量 query
#                 cand_mid_dict=cand_mid_dict,
#                 cand_last_dict=cand_last_dict,
#                 ee_cfg=ee_cfg,
#                 dataset_name=dataset_name,
#                 out_dir=data_args.encode_output_path,
#                 global_ranking=rank_global,
#             )

#         if dist.is_initialized():
#             dist.barrier()
#         # === EE-only 结束；直接进入下一个数据集 ===
#         continue


# if __name__ == '__main__':
#     main()

import datetime
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
import transformers
import math

from torch.utils.data import DataLoader
from tqdm import tqdm
from transformers import HfArgumentParser, AutoConfig, AutoTokenizer
from datasets import Dataset, concatenate_datasets
from datasets.distributed import split_dataset_by_node
from src.model.vlm_backbone.qwen2_vl.modeling_qwen2_vl_train_tokrnpooling import Qwen2VLForConditionalGeneration as _Qwen2VLForConditionalGeneration_src

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_add_text_cut import MMEBModel
from src.model.processor import get_backbone_name, load_processor, COLPALI
from src.utils import batch_to_device, print_rank, print_master

from dataclasses import dataclass

def get_env_mid_layer():
    v = os.environ.get("MID_LM_LAYER", "").strip()
    if v == "" or v.lower() in {"none", "null"}:
        return None
    try:
        return int(v)
    except:
        logger.warning(f"Invalid MID_LM_LAYER={v}, ignore.")
        return None

# ------------- AOP-Prune config parsing -------------
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_aop_config():
    """
    从环境变量读取 AOP 剪裁配置。仅作为“驱动层”的简要测试开关；
    实际剪裁逻辑在底模里（Qwen2-VLModel.forward）实现。
    """
    enabled     = _parse_bool(os.environ.get("AOP_ENABLED"), False)
    apply_to    = os.environ.get("AOP_APPLY", "qry").strip().lower()  # qry|cand|both
    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)

    # 按类型控制
    prune_vision = _parse_bool(os.environ.get("AOP_PRUNE_VISION"), True)
    prune_text   = _parse_bool(os.environ.get("AOP_PRUNE_TEXT"), False)

    delta_v     = _parse_float(os.environ.get("AOP_DELTA_VISION"), None)
    khat_v      = _parse_float(os.environ.get("AOP_KHAT_VISION"), None)
    keep_ratio_v= _parse_float(os.environ.get("AOP_KEEP_RATIO_VISION"), None)
    min_keep_v  = _parse_int(os.environ.get("AOP_MIN_KEEP_VISION"), None)

    delta_t     = _parse_float(os.environ.get("AOP_DELTA_TEXT"), None)
    khat_t      = _parse_float(os.environ.get("AOP_KHAT_TEXT"), None)
    keep_ratio_t= _parse_float(os.environ.get("AOP_KEEP_RATIO_TEXT"), None)
    min_keep_t  = _parse_int(os.environ.get("AOP_MIN_KEEP_TEXT"), 32)

    protect_text_last = _parse_int(os.environ.get("AOP_PROTECT_TEXT_LAST"), 16)
    protect_special   = _parse_bool(os.environ.get("AOP_PROTECT_SPECIAL"), True)

    margin_src  = os.environ.get("AOP_MARGIN", "").strip().lower()    # "" or "mid"
    attn_impl   = os.environ.get("AOP_ATTN_IMPL", "").strip().lower() # "" or "sdpa"

    if layer_idx is None and enabled:
        logger.warning("AOP_ENABLED=1 但未设置 AOP_LAYER，关闭 AOP。"); enabled=False

    # 新增：选择策略（aop | random）
    selection = os.environ.get("AOP_SELECTION", "aop").strip().lower()
    if _parse_bool(os.environ.get("AOP_RANDOM"), False):
        selection = "random"
    random_seed = _parse_int(os.environ.get("AOP_RANDOM_SEED"), None)

    # 选择策略：aop | random | attention
    selection = os.environ.get("AOP_SELECTION", "aop").strip().lower()
    if _parse_bool(os.environ.get("AOP_RANDOM"), False):
        selection = "random"
    random_seed = _parse_int(os.environ.get("AOP_RANDOM_SEED"), None)
    attn_agg = os.environ.get("AOP_ATTENTION_AGG", "mean").strip().lower()  # mean|max|sum

    cfg = {
        "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,

        # 类型开关
        "prune_vision": prune_vision,
        "prune_text": prune_text,

        # 视觉桶
        "delta_vision": delta_v,
        "K_hat_vision": khat_v,
        "keep_ratio_vision": keep_ratio_v,
        "min_keep_vision": min_keep_v,

        # 文本桶
        "delta_text": delta_t,
        "K_hat_text": khat_t,
        "keep_ratio_text": keep_ratio_t,
        "min_keep_text": min_keep_t,

        # 文本保护
        "protect_text_last": protect_text_last,
        "protect_special": protect_special,

        # 可选：排名安全预算
        "margin_mid": None if margin_src != "mid" else "USE_MID_MARGIN",
        "epsilon_hat": None,
        "attn_impl_override": attn_impl if attn_impl in {"sdpa"} else "",

        # NEW: 选择策略
        "selection": selection,          # "aop" 或 "random"
        "random_seed": random_seed,      # 可选
        "attn_agg": attn_agg,
    }
    return cfg

def get_env_eval_layers():
    """
    解析环境变量 LM_LAYERS（优先）或兼容旧的 MID_LM_LAYER。
    - LM_LAYERS 示例："4,8,12,last"；可包含 'last'/'none'/'null'/'-1' 表示最后一层(None)。
    - 若未设置 LM_LAYERS，则回落到旧逻辑：MID_LM_LAYER=None -> [None]；否则 [mid, None]
    返回: list[ int | None ]，例如 [4, 8, 12, None]；None 代表最后一层。
    """
    v = os.environ.get("LM_LAYERS", None)
    if v is not None:
        v = v.strip()

    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)
                    else:
                        logger.warning(f"Ignoring non-positive layer '{tok}' in LM_LAYERS.")
                except Exception:
                    logger.warning(f"Invalid token '{tok}' in LM_LAYERS; must be int or 'last'/'none'.")
        # 去重但保持顺序
        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)
        if not uniq:
            return [None]
        return uniq
    else:
        # 兼容旧逻辑
        mid = get_env_mid_layer()
        return [None] if mid is None else [mid, None]

# === Early-Exit config & helpers ===
def get_env_ee_config():
    ee_enabled = os.environ.get("EE_ENABLED", "0").strip().lower() in {"1","true","yes","on","y","t"}
    layer = int(os.environ.get("EE_LAYER", os.environ.get("AOP_LAYER", "12")))  # 默认用 AOP_LAYER
    method = os.environ.get("EE_METHOD", "margin").strip().lower()  # margin|p1p2|entropy|gini|combined
    tau = float(os.environ.get("EE_TAU", "0.2"))
    topk = int(os.environ.get("EE_TOPK", "1024"))
    temp = float(os.environ.get("EE_TEMP", "0.05"))
    save = os.environ.get("EE_SAVE", "1").strip().lower() in {"1","true","yes","on","y","t"}
    combw = os.environ.get("EE_COMB_WEIGHTS", "1.0,0.5,0.5")
    try:
        w_margin, w_conf, w_sq = [float(x) for x in combw.split(",")]
    except Exception:
        w_margin, w_conf, w_sq = 1.0, 0.5, 0.5
    return dict(
        enabled=ee_enabled, layer=layer, method=method, tau=tau,
        topk=topk, temp=temp, save=save,
        w_margin=w_margin, w_conf=w_conf, w_sq=w_sq
    )

def _softmax_np(x: np.ndarray, temp: float = 1.0) -> np.ndarray:
    x = x - np.max(x)
    ex = np.exp(x / max(1e-6, temp))
    s = np.sum(ex)
    return ex / max(s, 1e-12)

def confidence_from_topk(scores: np.ndarray, method="margin", temp=0.05, w_margin=1.0, w_conf=0.5, w_sq=0.5) -> float:
    # scores: 已按降序排列（topK）
    if scores.size == 0:
        return 0.0
    if scores.size == 1:
        return 1e9
    margin = float(scores[0] - scores[1])
    p = _softmax_np(scores, temp=temp)
    p1p2 = float(p[0] - p[1])
    H = - float(np.sum(p * np.log(p + 1e-12))) / np.log(len(p))   # 归一化熵 ∈ [0,1]
    conf = 1.0 - H
    sqsum = float(np.sum(p**2))  # Gini 的等价度量（越大越集中）
    if method == "margin":   return margin
    if method == "p1p2":     return p1p2
    if method == "entropy":  return conf
    if method == "gini":     return sqsum
    # combined
    return w_margin*margin + w_conf*conf + w_sq*sqsum

def run_early_exit_queries(
    model: MMEBModel,
    processor,
    model_args: ModelArguments,
    data_args: DataArguments,
    training_args: TrainingArguments,
    qry_dataset: Dataset,
    cand_mid_dict: dict,
    cand_last_dict: dict,
    ee_cfg: dict,
    dataset_name: str,
    out_dir: str,
    global_ranking: bool = True,
):
    """
    仅在线早停推理（不画曲线），并额外输出：
      - 每个 query 的中间层 vs cand_last 的 top-K 相似度 (mid2last)
      - 对未早停的 query，再输出最后一层 vs cand_last 的 top-K 相似度 (last2last)

    输出文件：
      {out_dir}/{dataset}_score_earlyexit.json    - 检索指标
      {out_dir}/{dataset}_pred_earlyexit.jsonl    - 预测列表（原有）
      {out_dir}/{dataset}_sim_earlyexit.jsonl     - 本函数新增的相似度信息（需 EE_SAVE=1）
    """
    device = training_args.device
    local_rank = dist.get_rank() if dist.is_initialized() else 0
    is_main = (not dist.is_initialized()) or (local_rank == 0)

    # 候选矩阵
    cand_ids = list(cand_mid_dict.keys())
    cand_id2row = {str(cid): i for i, cid in enumerate(cand_ids)}

    cand_mid = np.stack([cand_mid_dict[c] for c in cand_ids]).astype(np.float32)
    cand_last = np.stack([cand_last_dict[c] for c in cand_ids]).astype(np.float32)

    cand_mid_t  = torch.from_numpy(cand_mid).to(device=device, dtype=torch.bfloat16)
    cand_last_t = torch.from_numpy(cand_last).to(device=device, dtype=torch.bfloat16)

    # query DataLoader
    collator = MultimodalEvalDataCollator(processor, model_args, data_args, "qry")
    loader = DataLoader(
        qry_dataset,
        batch_size=training_args.per_device_eval_batch_size,
        collate_fn=collator,
        num_workers=training_args.dataloader_num_workers
    )

    pred_dicts = []

    # 是否保存相似度（沿用 EE_SAVE）
    save_scores = ee_cfg.get("save", False)
    topk_sim = int(ee_cfg.get("topk", 1024))
    sim_records = [] if (save_scores and is_main) else None

    # AOP 按侧开启
    aop_cfg = getattr(model.encoder, "aop_prune_config", None)
    _orig_enabled = None
    side_enable = True
    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 == "qry")

    # 门控相关
    k_conf = 2
    tau   = float(ee_cfg["tau"])
    method= ee_cfg["method"]
    temp  = float(ee_cfg["temp"])

    idx_global = 0
    start_time = time.time()

    for inputs, infos in tqdm(
        loader,
        desc=f"[EE] {dataset_name}@L{ee_cfg['layer']} (rank {local_rank})",
        disable=local_rank > 0,
    ):
        inputs = batch_to_device(inputs, device)

        # -------- 1) 跑到中间层（stop_at_layer），不算 logits --------
        orig_cfg = None
        if isinstance(aop_cfg, dict) and aop_cfg:
            orig_cfg   = dict(aop_cfg)
            aop_layer  = aop_cfg.get("layer_idx", None)
            ee_layer   = int(ee_cfg["layer"])
            apply_to   = aop_cfg.get("apply_to", "qry").strip().lower()
            aop_on_mid = bool(
                _orig_enabled and side_enable and
                (aop_layer is not None) and (aop_layer < ee_layer) and
                (apply_to in {"qry", "both"})
            )
            aop_cfg_mid = dict(aop_cfg)
            aop_cfg_mid["enabled"] = aop_on_mid
            setattr(model.encoder, "aop_prune_config", aop_cfg_mid)

        with torch.no_grad(), torch.autocast(
            device_type="cuda", dtype=torch.bfloat16, enabled=True
        ):
            out_mid = model.encoder(
                **inputs,
                return_dict=True,
                output_hidden_states=False,
                stop_at_layer=int(ee_cfg["layer"]),
                compute_lm_head=False,
            )

        if isinstance(orig_cfg, dict):
            setattr(model.encoder, "aop_prune_config", orig_cfg)

        # EOS pooling 得到中间层表征
        hs_mid = getattr(out_mid, "last_hidden_state", None)
        if hs_mid is None:
            assert out_mid.hidden_states is not None and len(out_mid.hidden_states) > 0
            hs_mid = out_mid.hidden_states[-1]
        am_mid = getattr(out_mid, "attention_mask", None)
        if am_mid is None:
            am_mid = inputs.get("attention_mask", None)
        if hasattr(am_mid, "device") and am_mid.device != hs_mid.device:
            am_mid = am_mid.to(hs_mid.device)
        reps_mid_t = model._pooling(hs_mid, am_mid).detach().to(device=device, dtype=torch.bfloat16)  # [B,D]

        B = reps_mid_t.size(0)
        use_local = (not global_ranking)

        # -------- 2) 门控：基于 mid→mid 的 top2 分数 --------
        if not use_local:
            # 全库：cand_mid_t
            scores_t = reps_mid_t @ cand_mid_t.T                # [B, Nc]
            vals_t, idxs_t = torch.topk(
                scores_t, k=min(k_conf, scores_t.size(1)), dim=1
            )                                                   # [B,2]
            p_t = torch.softmax(vals_t / max(temp, 1e-6), dim=1)
            if vals_t.size(1) >= 2:
                margin_t = vals_t[:, 0] - vals_t[:, 1]
                p1p2_t   = p_t[:, 0] - p_t[:, 1]
            else:
                margin_t = torch.full((B,), float("inf"), device=device, dtype=vals_t.dtype)
                p1p2_t   = torch.ones(B, device=device, dtype=vals_t.dtype)
            H_t = -(p_t * (torch.log(p_t + 1e-12))).sum(dim=1) / math.log(max(vals_t.size(1),1))
            conf_map = {
                "margin":  margin_t,
                "p1p2":    p1p2_t,
                "entropy": 1.0 - H_t,
                "gini":    (p_t ** 2).sum(dim=1),
            }
            confs_t = conf_map.get(method, margin_t)
            exit_mask = (confs_t >= tau).detach().cpu().numpy().astype(bool)
        else:
            # local：对每个 query 单独用 cand_mid_t[rows]
            confs = []
            for b in range(B):
                cand_local = infos[b]["cand_names"]
                rows = [cand_id2row.get(str(cid), -1) for cid in cand_local]
                rows = [r for r in rows if r >= 0]
                if len(rows) == 0:
                    confs.append(0.0)
                    continue
                cmat_t = cand_mid_t[rows]                # [Nl, D]
                sv_t   = (reps_mid_t[b:b+1] @ cmat_t.T)[0]  # [Nl]
                k = 2 if sv_t.size(0) >= 2 else 1
                vals_t, _ = torch.topk(sv_t, k=k, dim=0)
                p_t = torch.softmax(vals_t / max(temp, 1e-6), dim=0)
                if k >= 2:
                    margin = (vals_t[0] - vals_t[1]).item()
                    p1p2   = (p_t[0] - p_t[1]).item()
                else:
                    margin, p1p2 = float("inf"), 1.0
                H = (-(p_t * (torch.log(p_t + 1e-12))).sum() / math.log(max(k,1))).item()
                gini = ((p_t ** 2).sum()).item()
                d = {"margin": margin, "p1p2": p1p2, "entropy": 1.0 - H, "gini": gini}
                confs.append(d.get(method, margin))
            exit_mask = (np.array(confs) >= tau)

        # -------- 3) 检索 + 相似度记录 --------
        # 早停样本
        exit_indices = np.where(exit_mask)[0].tolist()
        # 续跑样本
        need_indices = np.where(~exit_mask)[0].tolist()

        # A. 早停：直接用 mid→mid 排序，但我们额外算 mid→last 的 top-K 相似度
        for j in exit_indices:
            # 1) 排序（pred_dicts）
            if not use_local:
                scores_mid_mid = (reps_mid_t[j:j+1] @ cand_mid_t.T)[0]  # [Nc]
                order = torch.argsort(scores_mid_mid, dim=0, descending=True).detach().cpu().numpy()
                cids = [cand_ids[i] for i in order]
            else:
                cand_local = infos[j]["cand_names"]
                rows = [cand_id2row.get(str(cid), -1) for cid in cand_local]
                rows = [r for r in rows if r >= 0]
                if len(rows) == 0:
                    cids = []
                else:
                    cmat_t = cand_mid_t[rows]
                    sv = (reps_mid_t[j:j+1] @ cmat_t.T)[0]
                    order_local = torch.argsort(sv, dim=0, descending=True).detach().cpu().numpy()
                    cids = [str(cand_local[i]) for i in order_local]
            rel_docids = infos[j]["label_name"]
            if not isinstance(rel_docids, list):
                rel_docids = [rel_docids]
            pred_dicts.append({"prediction": cids, "label": rel_docids, "rel_scores": None})

            # 2) 相似度记录：mid→last
            if save_scores and is_main:
                if not use_local:
                    scores_mid_last = (reps_mid_t[j:j+1] @ cand_last_t.T)[0].detach().float().cpu()  # [Nc]
                    Nc = scores_mid_last.size(0)
                    K = min(topk_sim, Nc)
                    mid_vals, mid_inds = torch.topk(scores_mid_last, k=K, dim=0)
                    mid_ids = [cand_ids[i] for i in mid_inds.tolist()]
                else:
                    cand_local = infos[j]["cand_names"]
                    rows = [cand_id2row.get(str(cid), -1) for cid in cand_local]
                    rows = [r for r in rows if r >= 0]
                    if len(rows) == 0:
                        mid_vals = torch.empty(0)
                        mid_ids = []
                    else:
                        cmat_last = cand_last_t[rows]   # [Nl, D]
                        sv_last = (reps_mid_t[j:j+1] @ cmat_last.T)[0].detach().float().cpu()
                        K = min(topk_sim, sv_last.size(0))
                        mid_vals, mid_inds = torch.topk(sv_last, k=K, dim=0)
                        mid_ids = [str(cand_local[i]) for i in mid_inds.tolist()]

                rec = {
                    "qid": int(idx_global + j),
                    "early_exit": True,
                    "mid_topk_scores": mid_vals.tolist() if mid_vals.numel() > 0 else [],
                    "mid_topk_cand_ids": mid_ids,
                    "last_topk_scores": None,
                    "last_topk_cand_ids": None,
                }
                sim_records.append(rec)

        # B. 续跑：mid->last，再用 last→last 排序；同时记录 mid→last & last→last 相似度
        if len(need_indices) > 0:
            # 从中间态恢复
            if isinstance(aop_cfg, dict) and aop_cfg:
                aop_resume = dict(aop_cfg)
                aop_resume["enabled"] = bool(_orig_enabled and side_enable)
                setattr(model.encoder, "aop_prune_config", aop_resume)

            interm = getattr(out_mid, "intermediate_state", None)
            assert interm is not None, "Model must return intermediate_state when stop_at_layer is set."

            hs = interm["hidden_states"].detach()
            am = interm["attention_mask"].detach()
            pos = interm["position_ids"].detach()
            vm = interm.get("vision_mask", None)
            tm = interm.get("text_mask", None)
            next_layer = int(interm["next_layer_idx"])

            hs_sub  = hs[need_indices]
            am_sub  = am[need_indices]
            pos_sub = pos[:, need_indices, :]
            vm_sub  = vm[need_indices] if vm is not None else None
            tm_sub  = tm[need_indices] if tm is not None else None
            resume_state = {
                "hidden_states": hs_sub,
                "attention_mask": am_sub,
                "position_ids": pos_sub,
                "vision_mask": vm_sub,
                "text_mask": tm_sub,
                "next_layer_idx": next_layer,
            }

            with torch.no_grad(), torch.autocast(
                device_type="cuda", dtype=torch.bfloat16, enabled=True
            ):
                out_last = model.encoder(
                    return_dict=True,
                    output_hidden_states=False,
                    stop_at_layer=None,
                    resume_state=resume_state,
                    compute_lm_head=False,
                )

            hs_last = getattr(out_last, "last_hidden_state", None)
            if hs_last is None:
                assert out_last.hidden_states is not None and len(out_last.hidden_states) > 0
                hs_last = out_last.hidden_states[-1]
            am_last = getattr(out_last, "attention_mask", None)
            if am_last is None:
                am_last = am_sub
            if hasattr(am_last, "device") and am_last.device != hs_last.device:
                am_last = am_last.to(hs_last.device)

            reps_last_t = model._pooling(hs_last, am_last).detach().to(device=device, dtype=torch.bfloat16)

            if not use_local:
                scores_last_all = (reps_last_t @ cand_last_t.T).detach().float().cpu()  # [N_need, Nc]
                for k, j in enumerate(need_indices):
                    # 1) 排序预测
                    row = scores_last_all[k]
                    order = torch.argsort(row, dim=0, descending=True).tolist()
                    cids = [cand_ids[i] for i in order]
                    rel_docids = infos[j]["label_name"]
                    if not isinstance(rel_docids, list):
                        rel_docids = [rel_docids]
                    pred_dicts.append({"prediction": cids, "label": rel_docids, "rel_scores": None})

                    # 2) mid→last & last→last 相似度
                    if save_scores and is_main:
                        # mid2last
                        scores_mid_last = (reps_mid_t[j:j+1] @ cand_last_t.T)[0].detach().float().cpu()
                        Nc = scores_mid_last.size(0)
                        K = min(topk_sim, Nc)
                        mid_vals, mid_inds = torch.topk(scores_mid_last, k=K, dim=0)
                        mid_ids = [cand_ids[i] for i in mid_inds.tolist()]

                        # last2last
                        last_row = row
                        last_vals, last_inds = torch.topk(last_row, k=K, dim=0)
                        last_ids = [cand_ids[i] for i in last_inds.tolist()]

                        rec = {
                            "qid": int(idx_global + j),
                            "early_exit": False,
                            "mid_topk_scores": mid_vals.tolist(),
                            "mid_topk_cand_ids": mid_ids,
                            "last_topk_scores": last_vals.tolist(),
                            "last_topk_cand_ids": last_ids,
                        }
                        sim_records.append(rec)
            else:
                # local ranking
                for k, j in enumerate(need_indices):
                    cand_local = infos[j]["cand_names"]
                    rows = [cand_id2row.get(str(cid), -1) for cid in cand_local]
                    rows = [r for r in rows if r >= 0]
                    if len(rows) == 0:
                        cids = []
                        rel_docids = infos[j]["label_name"]
                        if not isinstance(rel_docids, list):
                            rel_docids = [rel_docids]
                        pred_dicts.append({"prediction": cids, "label": rel_docids, "rel_scores": None})
                        if save_scores and is_main:
                            rec = {
                                "qid": int(idx_global + j),
                                "early_exit": False,
                                "mid_topk_scores": [],
                                "mid_topk_cand_ids": [],
                                "last_topk_scores": [],
                                "last_topk_cand_ids": [],
                            }
                            sim_records.append(rec)
                        continue

                    # 1) 排序（last→last）
                    cmat_last = cand_last_t[rows]
                    sv_last = (reps_last_t[k:k+1] @ cmat_last.T)[0].detach().float().cpu()
                    order_local = torch.argsort(sv_last, dim=0, descending=True).tolist()
                    cids = [str(cand_local[i]) for i in order_local]
                    rel_docids = infos[j]["label_name"]
                    if not isinstance(rel_docids, list):
                        rel_docids = [rel_docids]
                    pred_dicts.append({"prediction": cids, "label": rel_docids, "rel_scores": None})

                    # 2) mid→last & last→last 相似度
                    if save_scores and is_main:
                        cmat_last = cand_last_t[rows]
                        # mid2last
                        sv_mid_last = (reps_mid_t[j:j+1] @ cmat_last.T)[0].detach().float().cpu()
                        K = min(topk_sim, sv_mid_last.size(0))
                        mid_vals, mid_inds = torch.topk(sv_mid_last, k=K, dim=0)
                        mid_ids = [str(cand_local[i]) for i in mid_inds.tolist()]
                        # last2last
                        sv_last_row = sv_last
                        last_vals, last_inds = torch.topk(sv_last_row, k=K, dim=0)
                        last_ids = [str(cand_local[i]) for i in last_inds.tolist()]

                        rec = {
                            "qid": int(idx_global + j),
                            "early_exit": False,
                            "mid_topk_scores": mid_vals.tolist(),
                            "mid_topk_cand_ids": mid_ids,
                            "last_topk_scores": last_vals.tolist(),
                            "last_topk_cand_ids": last_ids,
                        }
                        sim_records.append(rec)

        idx_global += B

    # -------- 4) 评测 + 写出 --------
    metrics_to_report = ["hit", "ndcg", "precision", "recall", "f1", "map", "mrr"]
    score = RankingMetrics(metrics_to_report).evaluate(pred_dicts)

    if is_main:
        os.makedirs(out_dir, exist_ok=True)
        # 原有的 early-exit 检索结果
        with open(os.path.join(out_dir, f"{dataset_name}_score_earlyexit.json"), "w") as f:
            json.dump(score, f, indent=4)
        if ee_cfg.get("save", False):
            with open(os.path.join(out_dir, f"{dataset_name}_pred_earlyexit.jsonl"), "w", encoding="utf-8") as f:
                for row in pred_dicts:
                    f.write(json.dumps(row, ensure_ascii=False) + "\n")
        # 新增的 mid/last 相似度输出
        if save_scores and sim_records is not None:
            sims_path = os.path.join(out_dir, f"{dataset_name}_sim_earlyexit.jsonl")
            with open(sims_path, "w", encoding="utf-8") as f:
                for rec in sim_records:
                    f.write(json.dumps(rec, ensure_ascii=False) + "\n")
            print_master(f"[EE] Saved mid/last similarity records -> {sims_path}")

    elapsed = time.time() - start_time
    return score, elapsed

def make_layer_tag(keep_layers: int | None):
    return f"layer{keep_layers}" if keep_layers and keep_layers > 0 else "layerlast"

def dot_sim(a: np.ndarray, b: np.ndarray) -> np.ndarray:
    # a: [Nq, D], b: [Nc, D], both L2-normalized already if normalize=true
    return a @ b.T

def build_score_details(qid: int, cand_ids: list, score_vec: np.ndarray, ranked_indices: np.ndarray):
    return {
        "qid": int(qid),
        "cand_scores": [
            {"cand_id": str(cand_ids[i]), "score": float(score_vec[i])}
            for i in ranked_indices
        ]
    }

def top1_top2_margin(score_vec: np.ndarray) -> float:
    if len(score_vec) < 2:
        return float("inf")  # 只有一个候选时视作极大margin
    top2 = np.partition(score_vec, -2)[-2:]
    top2.sort()
    return float(top2[-1] - top2[-2])

def simulate_early_exit_by_margin(
    sims_mid: list[dict], sims_last: list[dict], labels: list[list[str]], metrics_to_report: list[str],
    taus: list[float], rank_global: bool
):
    """
    sims_mid / sims_last: 每个query一个dict: {cand_id: score}
    labels: 每个query的正样本cand_id列表
    返回：不同tau下的覆盖率、指标
    """
    assert len(sims_mid) == len(sims_last) == len(labels)
    N = len(labels)
    results = []

    from src.eval_utils.metrics import RankingMetrics
    metrics = RankingMetrics(metrics_to_report)

    # 预构造 用于metrics.evaluate 的pred_dict
    def to_pred_dicts(use_mid_mask: list[bool]) -> list[dict]:
        pred_dicts = []
        for qid in range(N):
            sims_use = sims_mid[qid] if use_mid_mask[qid] else sims_last[qid]
            # 排序
            ranked = sorted(sims_use.items(), key=lambda x: -x[1])
            pred_dicts.append({
                "prediction": [cid for cid, _ in ranked],
                "label": labels[qid],
                "rel_scores": None
            })
        return pred_dicts

    # 计算中间层margin
    margins = []
    for qid in range(N):
        # 取前两大分数的margin
        if len(sims_mid[qid]) == 0:
            margins.append(0.0)
            continue
        scores = np.array(list(sims_mid[qid].values()), dtype=np.float32)
        margins.append(top1_top2_margin(scores))

    margins = np.array(margins, dtype=np.float32)

    for tau in taus:
        use_mid_mask = (margins >= tau).tolist()
        pred_dicts = to_pred_dicts(use_mid_mask)
        score_dict = metrics.evaluate(pred_dicts)
        coverage = float(np.mean(use_mid_mask))  # 早停覆盖率
        results.append({
            "tau": tau,
            "coverage": coverage,
            **score_dict
        })
    return results

def top1_top2_margin_from_array(score_vec: np.ndarray) -> float:
    if score_vec is None or len(score_vec) == 0:
        return 0.0
    if len(score_vec) == 1:
        return float('inf')
    # 取前两大
    top2 = np.partition(score_vec, -2)[-2:]
    top2.sort()
    return float(top2[-1] - top2[-2])

logging.basicConfig(level=logging.INFO, format='[%(asctime)s] %(levelname)s [%(name)s:%(lineno)s] %(message)s')
logger = logging.getLogger(__name__)

# --- Global Dictionaries for Hooks (will be cleared before each encode_embeddings call) ---
timing_info = {}
token_info = {
    "vision_tokens": 0,
    "text_input_tokens": 0,  # Refers to the original text token count
    "text_output_tokens": 0, # Not directly applicable here as we are encoding, not generating. Will be 0.
    "total_llm_input_tokens": 0, # Refers to the total tokens LLM receives (visual + formatted text)
}

# --- Hook Functions Definition ---
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:
        # print(f"Warning: No pre-hook data for module {module.__class__.__name__} ({module_id})")
        return

    timing_info[module_id].append((time.time(), 'post', module.__class__.__name__))

    # Collect vision token count (only from Vision Transformer module's post hook)
    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] # For visual features, usually (batch_size, num_tokens, hidden_dim)
        elif hasattr(output, 'last_hidden_state'):
            token_info["vision_tokens"] = output.last_hidden_state.shape[1]


def register_model_hooks(model):
    registered_modules = []
    
    core_model = model.encoder if hasattr(model, "encoder") and model.encoder is not None else model
        
    # Vision module
    if hasattr(core_model, 'visual') and core_model.visual is not None:
        vision_module = core_model.visual
        vision_module.register_forward_pre_hook(timing_pre_hook)
        vision_module.register_forward_hook(timing_post_hook)
        registered_modules.append(vision_module)
        print_master(f"Registered hooks for vision module: {vision_module.__class__.__name__}")
    else:
        print_master(f"WARNING: No 'visual' attribute found on core_model ({type(core_model)}).")


    # Merger module (if inside visual) - it's part of the vision component
    if hasattr(core_model, 'visual') and hasattr(core_model.visual, 'merger') and core_model.visual.merger is not None:
        merger_module = core_model.visual.merger
        merger_module.register_forward_pre_hook(timing_pre_hook)
        merger_module.register_forward_hook(timing_post_hook)
        registered_modules.append(merger_module)
        print_master(f"Registered hooks for merger module: {merger_module.__class__.__name__}")
    else:
        print_master(f"WARNING: No 'merger' attribute found on core_model.visual ({type(getattr(core_model, 'visual', 'N/A'))}).")

    # Language model body
    if hasattr(core_model, 'model') and core_model.model is not None:
        llm_main_module = core_model.model
        llm_main_module.register_forward_pre_hook(timing_pre_hook)
        llm_main_module.register_forward_hook(timing_post_hook)
        registered_modules.append(llm_main_module)
        print_master(f"Registered hooks for LLM main module: {llm_main_module.__class__.__name__}")
    else:
        print_master(f"WARNING: No 'model' attribute found on core_model ({type(core_model)}).")


    # LM Head
    if hasattr(core_model, 'lm_head') and core_model.lm_head is not None:
        lm_head_module = core_model.lm_head
        lm_head_module.register_forward_pre_hook(timing_pre_hook)
        lm_head_module.register_forward_hook(timing_post_hook)
        registered_modules.append(lm_head_module)
        print_master(f"Registered hooks for LM head module: {lm_head_module.__class__.__name__}")
    else:
        print_master(f"WARNING: No 'lm_head' attribute found on core_model ({type(core_model)}).")


    if not registered_modules:
        print_master("Warning: No major modules found for hook registration. Check model architecture.")
    return registered_modules


def pad_dataset_to_divisible(dataset, world_size):
    num_samples = len(dataset)
    if num_samples % world_size == 0:
        return dataset, num_samples

    num_to_add = world_size - (num_samples % world_size)
    padded_size = num_samples + num_to_add

    padding_data = dataset.select([i % len(dataset) for i in range(num_to_add)])
    padded_dataset = concatenate_datasets([dataset, padding_data])
    return padded_dataset, padded_size

def encode_embeddings(
    model: MMEBModel,
    loader: DataLoader,
    training_args: TrainingArguments,
    model_args: ModelArguments,
    full_dataset: Dataset,
    encode_side: str,
    description: str = "Encoding"
) -> tuple[np.ndarray, list, list, list]:  # CHANGED: + list for img_token_masks
    """
    Encodes embeddings for a given dataset using the model, handling both standard and
    late-interaction models in a DDP-safe manner.
    Returns:
        - embeddings: np.ndarray
        - infos_or_ids: list
        - batch_stats_list: list
        - img_token_masks: list[None | list[bool]]  # NEW
    """
    local_rank = dist.get_rank() if dist.is_initialized() else 0
    world_size = dist.get_world_size() if dist.is_initialized() else 1

    # Check if the model is a late-interaction type
    is_late_interaction = (model_args.model_backbone == COLPALI)

    local_embeds = []
    local_gt_infos = []
    local_max_len = 0
    
    # --- New: List to store statistics for each batch ---
    batch_stats_list = []

    # --- NEW: Collect masks ---
    local_img_token_masks = []   # post image mask per sample
    local_txt_token_masks = []   # NEW: post text mask per sample
    local_post_attn_masks = []   # NEW: post attention_mask per sample (after prune, 1/0)

    # --- NEW: per-sample token reduction records ---
    local_token_records = []     # 每条样本一个 dict，含 pre/post/delta 数量

    model.eval()

    # Register hooks for the model once per encode_embeddings call
    registered_hooks = register_model_hooks(model)

    # --- NEW: helpers to取mask并序列化 ---
    def _search_key(obj, key: str):
        # 递归搜索 dict/list/tuple，找到指定 key
        if isinstance(obj, dict):
            if key in obj:
                return obj[key]
            for v in obj.values():
                r = _search_key(v, key)
                if r is not None:
                    return r
        elif isinstance(obj, (list, tuple)):
            for v in obj:
                r = _search_key(v, key)
                if r is not None:
                    return r
        return None

    def _to_serializable_mask_list(mask_list, batch_size: int):
        # 将模型返回的 mask（list/tensor/ndarray/None）转成 [None | list[bool]] * B
        if mask_list is None:
            return [None] * batch_size

        out = []
        if isinstance(mask_list, (list, tuple)):
            for m in mask_list:
                if m is None:
                    out.append(None)
                elif torch.is_tensor(m):
                    out.append(m.detach().cpu().tolist())
                elif isinstance(m, np.ndarray):
                    out.append(m.tolist())
                else:
                    # already python list/bool
                    out.append(m)
        elif torch.is_tensor(mask_list):
            # 若是 2D 张量（B, L），直接 tolist() -> list[list[bool/int]]
            out = mask_list.detach().cpu().tolist()
        elif isinstance(mask_list, np.ndarray):
            out = mask_list.tolist()
        else:
            # 未知类型，保守返回 None 占位
            out = [None] * batch_size

        # 长度对齐 batch_size
        if isinstance(out, list):
            if len(out) < batch_size:
                out = out + [None] * (batch_size - len(out))
            elif len(out) > batch_size:
                out = out[:batch_size]
        return out
    
    def _to_bool_lists(m, batch_size: int):
        lst = _to_serializable_mask_list(m, batch_size)
        # 归一化成 list[ list[bool] | None ]
        out = []
        for x in lst:
            if x is None:
                out.append(None)
            else:
                # x 可能是 list[int] 或 list[bool]
                out.append([bool(int(v)) for v in x])
        return out

    with torch.no_grad():
        for inputs, dataset_info in tqdm(loader, desc=f"{description} (rank {local_rank})", disable=local_rank > 0):
            # --- Reset statistics for each inference pass ---
            timing_info.clear()
            token_info["vision_tokens"] = 0
            token_info["text_input_tokens"] = 0
            token_info["text_output_tokens"] = 0
            token_info["total_llm_input_tokens"] = 0

            inputs = batch_to_device(inputs, training_args.device)
            current_batch_size = inputs['input_ids'].shape[0] if 'input_ids' in inputs and inputs['input_ids'] is not None else 1

            with torch.autocast(enabled=True, dtype=torch.bfloat16, device_type="cuda"):
                start_inference_time = time.time()
                # ---- NEW: 按侧开/关 AOP ----
                aop_cfg = getattr(model.encoder, "aop_prune_config", None)
                _orig_enabled = None
                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)
                    aop_cfg["enabled"] = bool(side_enable and _orig_enabled)
                    setattr(model.encoder, "aop_prune_config", aop_cfg)
                if encode_side == "qry":
                    output = model(qry=inputs)
                    reps = output["qry_reps"].detach()
                    local_gt_infos.extend(dataset_info)
                else:
                    output = model(tgt=inputs)
                    reps = output["tgt_reps"].detach()
                    local_gt_infos.extend([info["cand_name"] for info in dataset_info])
                # ---- NEW: 恢复 enabled（避免影响下个 encode_side）----
                if isinstance(aop_cfg, dict) and _orig_enabled is not None:
                    aop_cfg["enabled"] = _orig_enabled
                    setattr(model.encoder, "aop_prune_config", aop_cfg)
                end_inference_time = time.time()

            # --- NEW: 提取 post-prune 的 image/text 掩码 与 post attention_mask ---
            img_masks_raw = None
            txt_masks_raw = None
            post_attn_raw = None

            if isinstance(output, dict):
                img_masks_raw = _search_key(output, "image_token_bool_masks")
                txt_masks_raw = _search_key(output, "text_token_bool_masks")     # NEW
                post_attn_raw = _search_key(output, "post_attention_mask")       # NEW（我们的 MMEBModel.forward 里带了这个键）

            # 兼容：若挂在 model 上
            if img_masks_raw is None and hasattr(model, "image_token_bool_masks"):
                img_masks_raw = getattr(model, "image_token_bool_masks")
            if txt_masks_raw is None and hasattr(model, "text_token_bool_masks"):
                txt_masks_raw = getattr(model, "text_token_bool_masks")
            if post_attn_raw is None and hasattr(model, "post_attention_mask"):
                post_attn_raw = getattr(model, "post_attention_mask")

            img_masks_serializable = _to_serializable_mask_list(img_masks_raw, current_batch_size)
            txt_masks_serializable = _to_serializable_mask_list(txt_masks_raw, current_batch_size)     # NEW
            post_attn_serializable = _to_serializable_mask_list(post_attn_raw, current_batch_size)     # NEW

            local_img_token_masks.extend(img_masks_serializable)
            local_txt_token_masks.extend(txt_masks_serializable)        # NEW
            local_post_attn_masks.extend(post_attn_serializable)        # NEW

            # --- NEW: 计算本 batch 的 pre/post/delta 数量并累计 ---
            cfg = getattr(model.encoder, "config", None)
            # pre masks 来自 inputs（删前）
            input_ids = inputs.get("input_ids", None)
            attn2d_pre = inputs.get("attention_mask", None)
            if input_ids is None or attn2d_pre is None or cfg is None:
                # 无法统计，留空
                pre_vis_counts = [0] * current_batch_size
                pre_txt_counts = [0] * current_batch_size
                pre_tot_counts = [0] * current_batch_size
            else:
                iid = input_ids
                am = attn2d_pre.to(torch.bool)
                image_token_id = getattr(cfg, "image_token_id", None)
                video_token_id = getattr(cfg, "video_token_id", None)
                bos_id = getattr(cfg, "bos_token_id", None)
                eos_id = getattr(cfg, "eos_token_id", None)
                pad_id = getattr(cfg, "pad_token_id", None)

                is_image = (iid == image_token_id) if (image_token_id is not None and image_token_id >= 0) else torch.zeros_like(iid, dtype=torch.bool)
                is_video = (iid == video_token_id) if (video_token_id is not None and video_token_id >= 0) else torch.zeros_like(iid, dtype=torch.bool)
                is_vision = is_image | is_video

                is_special = torch.zeros_like(iid, dtype=torch.bool)
                for tid in [bos_id, eos_id, pad_id]:
                    if tid is not None and tid >= 0:
                        is_special |= (iid == tid)

                pre_txt_mask = am & (~is_vision) & (~is_special)
                pre_vis_mask = am & is_vision

                pre_vis_counts = pre_vis_mask.sum(dim=1).tolist()
                pre_txt_counts = pre_txt_mask.sum(dim=1).tolist()
                pre_tot_counts = am.sum(dim=1).tolist()

            # post masks（删后）来自模型输出；与 post_attn 做与运算
            post_text_masks = _to_bool_lists(txt_masks_raw, current_batch_size)      # list[ list[bool] | None ]
            post_image_masks = _to_bool_lists(img_masks_raw, current_batch_size)
            post_attn_masks  = _to_bool_lists(post_attn_raw, current_batch_size)

            sum_pre_text = 0; sum_post_text = 0
            sum_pre_vis  = 0; sum_post_vis  = 0
            sum_pre_tot  = 0; sum_post_tot  = 0

            for i in range(current_batch_size):
                pre_text = int(pre_txt_counts[i]) if i < len(pre_txt_counts) else 0
                pre_vis  = int(pre_vis_counts[i]) if i < len(pre_vis_counts) else 0
                pre_tot  = int(pre_tot_counts[i]) if i < len(pre_tot_counts) else 0

                # post 计数：mask 可能为 None
                m_text = post_text_masks[i] if post_text_masks is not None and i < len(post_text_masks) else None
                m_img  = post_image_masks[i] if post_image_masks is not None and i < len(post_image_masks) else None
                m_attn = post_attn_masks[i]  if post_attn_masks is not None  and i < len(post_attn_masks)  else None

                if m_attn is None:
                    post_text = 0; post_vis = 0; post_tot = 0
                else:
                    # 与 attention_mask 后统计 True 的数
                    if m_text is not None:
                        post_text = sum(1 for a, t in zip(m_attn, m_text) if a and t)
                    else:
                        post_text = 0
                    if m_img is not None:
                        post_vis = sum(1 for a, v in zip(m_attn, m_img) if a and v)
                    else:
                        post_vis = 0
                    post_tot = sum(1 for a in m_attn if a)

                # 累计 batch 级
                sum_pre_text += pre_text; sum_post_text += post_text
                sum_pre_vis  += pre_vis;  sum_post_vis  += post_vis
                sum_pre_tot  += pre_tot;  sum_post_tot  += post_tot

                # 保存 per-sample 记录（用于 JSONL）
                local_token_records.append({
                    "side": encode_side,
                    "pre":  {"text": pre_text,  "vision": pre_vis,  "total": pre_tot},
                    "post": {"text": post_text, "vision": post_vis, "total": post_tot},
                    "delta":{"text": pre_text - post_text, "vision": pre_vis - post_vis, "total": pre_tot - post_tot},
                })

            # --- Update total LLM input tokens after the model call ---
            if 'input_ids' in inputs and inputs['input_ids'] is not None:
                token_info["total_llm_input_tokens"] = inputs['input_ids'].shape[1]
                token_info["text_input_tokens"] = token_info["total_llm_input_tokens"] - token_info["vision_tokens"]
                token_info["text_input_tokens"] = max(0, token_info["text_input_tokens"])

            # --- Collect and Store Batch Statistics ---
            batch_inference_time = end_inference_time - start_inference_time
            
            current_batch_stats = {
                "batch_size": current_batch_size,
                "total_inference_time_seconds": batch_inference_time,
                "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"],
                }
            }
            current_batch_stats["token_reduction"] = {
                "sum_pre_text":  sum_pre_text,
                "sum_post_text": sum_post_text,
                "sum_pre_vision":  sum_pre_vis,
                "sum_post_vision": sum_post_vis,
                "sum_pre_total": sum_pre_tot,
                "sum_post_total": sum_post_tot,
            }

            # Calculate and store module timings for the current batch
            for module_obj in registered_hooks:
                module_id = id(module_obj)
                module_name = module_obj.__class__.__name__
                times = timing_info.get(module_id, [])
                durations = []
                pre_times = {} 
                for t, event_type, _ in times:
                    if event_type == 'pre':
                        pre_times[module_id] = t
                    elif event_type == 'post' and module_id in pre_times:
                        duration = t - pre_times.pop(module_id)
                        durations.append(duration)
                
                if durations:
                    current_batch_stats["module_inference_times"][module_name] = {
                        "total": sum(durations),
                        "count": len(durations),
                        "avg": sum(durations) / len(durations)
                    }
                else:
                    current_batch_stats["module_inference_times"][module_name] = {
                        "total": 0.0,
                        "count": 0,
                        "avg": 0.0
                    }
            
            batch_stats_list.append(current_batch_stats)

            # --- Debug prints (optional) ---
            print_rank(f"\n--- Inference Statistics for {encode_side} batch (Rank {local_rank}) ---")
            print_rank(f"Batch Inference took: {batch_inference_time:.4f} seconds")
            print_rank("--- Module Inference Timing Statistics ---")
            for module_name, stats in current_batch_stats["module_inference_times"].items():
                print_rank(f"**{module_name}**: Total: {stats['total']:.6f}s, Count: {stats['count']}, Avg: {stats['avg']:.6f}s")
            print_rank("--- Token Count Statistics ---")
            print_rank(f"**视觉 token 数量**: {current_batch_stats['token_counts']['visual_tokens']}")
            print_rank(f"**语言输入 token 数量 (仅原始文本)**: {current_batch_stats['token_counts']['language_input_tokens_raw']}")
            print_rank(f"**LLM总输入 token 数量 (包含视觉 + 格式化文本)**: {current_batch_stats['token_counts']['llm_total_input_tokens']}")
            print_rank(f"**语言输出 token 数量**: {current_batch_stats['token_counts']['language_output_tokens']}")

            if is_late_interaction and reps.dim() == 3:
                local_max_len = max(local_max_len, reps.shape[1])

            local_embeds.append(reps)

    if not local_embeds:
        # Handle cases where a rank gets no data
        return np.array([]), [], [], []  # CHANGED: 4个返回值

    # === DDP Synchronization and Padding for Late-Interaction Models ===
    if is_late_interaction:
        if dist.is_initialized():
            # 1: global max length
            local_max_len_tensor = torch.tensor(local_max_len, device=training_args.device)
            dist.all_reduce(local_max_len_tensor, op=dist.ReduceOp.MAX)
            global_max_len = local_max_len_tensor.item()
        else:
            global_max_len = local_max_len

        # 2: pad to global max length
        padded_embeds = []
        for reps_batch in local_embeds:
            if reps_batch.dim() == 3:
                B, L, H = reps_batch.shape
                padding_size = global_max_len - L
                padded_batch = F.pad(reps_batch, (0, 0, 0, padding_size), "constant", 0)
                padded_embeds.append(padded_batch)
            else:
                padded_embeds.append(reps_batch)

        embeds_tensor = torch.cat(padded_embeds, dim=0).contiguous()
    else:
        embeds_tensor = torch.cat(local_embeds, dim=0).contiguous()

    # === Gather embeddings and keys from all ranks ===
    if dist.is_initialized() and full_dataset.num_rows >= world_size:
        print_master(f"Gathering {encode_side} embeddings across all ranks...")

        # tensor gather
        output_shape = list(embeds_tensor.shape)
        output_shape[0] = full_dataset.num_rows
        embeds_tensor = embeds_tensor.to(training_args.device)
        gathered_embeds_tensor = torch.empty(output_shape, dtype=embeds_tensor.dtype, device=training_args.device)
        dist.all_gather_into_tensor(gathered_embeds_tensor, embeds_tensor)
        final_embeddings = gathered_embeds_tensor.cpu().float().numpy()

        # object gather for infos and stats
        gathered_gt_infos = [None for _ in range(world_size)]
        dist.all_gather_object(gathered_gt_infos, local_gt_infos)
        all_gt_infos = [key for rank_keys in gathered_gt_infos for key in rank_keys]

        gathered_batch_stats = [None for _ in range(world_size)]
        dist.all_gather_object(gathered_batch_stats, batch_stats_list)
        all_batch_stats = [stats for rank_stats in gathered_batch_stats for stats in rank_stats]

        # --- NEW: gather masks ---
        gathered_masks = [None for _ in range(world_size)]
        dist.all_gather_object(gathered_masks, local_img_token_masks)
        all_img_token_masks = [m for rank_list in gathered_masks for m in rank_list]
        # NEW: gather text masks
        gathered_txt_masks = [None for _ in range(world_size)]
        dist.all_gather_object(gathered_txt_masks, local_txt_token_masks)
        all_txt_token_masks = [m for rank_list in gathered_txt_masks for m in rank_list]

        # NEW: gather post attention masks（如需）
        gathered_post_attn = [None for _ in range(world_size)]
        dist.all_gather_object(gathered_post_attn, local_post_attn_masks)
        all_post_attn_masks = [m for rank_list in gathered_post_attn for m in rank_list]

        # NEW: gather token records
        gathered_token_recs = [None for _ in range(world_size)]
        dist.all_gather_object(gathered_token_recs, local_token_records)
        all_token_records = [r for rank_list in gathered_token_recs for r in rank_list]
    else:
        all_gt_infos = local_gt_infos
        final_embeddings = embeds_tensor.cpu().float().numpy()
        all_batch_stats = batch_stats_list
        all_img_token_masks = local_img_token_masks  # NEW
        all_txt_token_masks = local_txt_token_masks
        all_post_attn_masks = local_post_attn_masks
        all_token_records = local_token_records

    return final_embeddings, all_gt_infos, all_batch_stats, all_img_token_masks, all_txt_token_masks, all_token_records

# === NEW: 一次前向同时导出 cand 的中间层和最后一层向量 ===
def encode_candidates_both_layers(
    model: MMEBModel,
    loader: DataLoader,
    training_args: TrainingArguments,
    model_args: ModelArguments,
    full_dataset: Dataset,
    mid_layer: int,
) -> tuple[np.ndarray, np.ndarray, list]:
    """
    单次forward到最后一层，直接从 hidden_states 取：
      - mid_hidden = hidden_states[mid_layer]   # 表示经过 mid_layer 层后的状态（见Qwen2_5_VLModel的all_hidden_states定义）
      - last_hidden = hidden_states[-1]         # 最后一层norm后的状态
    然后用 _pooling(attention_mask) 取句向量，返回：
      - cand_mid_embeds: np.ndarray [Nc, D]
      - cand_last_embeds: np.ndarray [Nc, D]
      - cand_ids: list[str]
    说明：
      - cand 侧默认不做 AOP 剪枝（AOP_APPLY=qry 时天然关闭），因此 mid/last 的序列长度一致，可直接用原 attention_mask 做池化。
    """
    local_rank = dist.get_rank() if dist.is_initialized() else 0
    model.eval()

    all_mid = []
    all_last = []
    all_ids = []

    with torch.no_grad():
        for inputs, dataset_info in tqdm(loader, desc=f"Candidates[BOTH] (rank {local_rank})", disable=local_rank > 0):
            inputs = batch_to_device(inputs, training_args.device)
            # cand 侧确保不触发 AOP（如果你的 AOP_APPLY=qry/both，会在底模按侧门控；此处再做一次保险）
            aop_cfg = getattr(model.encoder, "aop_prune_config", None)
            _orig_enabled = None
            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 == "cand")
                aop_cfg["enabled"] = bool(side_enable and _orig_enabled)
                setattr(model.encoder, "aop_prune_config", aop_cfg)

            with torch.autocast(enabled=True, dtype=torch.bfloat16, device_type="cuda"):
                # 关键：一次forward拿全层的hidden_states
                out = model.encoder(
                    **inputs,
                    return_dict=True,
                    output_hidden_states=True,  # 必须
                    stop_at_layer=None,         # 走全层
                )

            # 取 hidden_states 并索引中间层/最后一层
            hs_list = out.hidden_states
            assert hs_list is not None and len(hs_list) > mid_layer, \
                f"hidden_states is None or too short. Need index {mid_layer}, got len={0 if hs_list is None else len(hs_list)}"
            mid_hs  = hs_list[mid_layer]      # [B, L, D]：等价“经过 mid_layer 层后的状态”（即 pre-layer(mid_layer+1)）
            last_hs = hs_list[-1]             # [B, L, D]：最终norm后的状态

            # 用原 attention_mask 池化（cand侧未剪枝）
            am = inputs.get("attention_mask", None)
            if am is not None and hasattr(am, "device"):
                if am.device != mid_hs.device:
                    am = am.to(mid_hs.device)

            reps_mid  = model._pooling(mid_hs, am)   # [B, D]
            reps_last = model._pooling(last_hs, am)  # [B, D]

            all_mid.append(reps_mid.detach().float().cpu())
            all_last.append(reps_last.detach().float().cpu())
            all_ids.extend([info["cand_name"] for info in dataset_info])

            # 恢复 AOP 开关（避免影响其它侧）
            if isinstance(aop_cfg, dict) and _orig_enabled is not None:
                aop_cfg["enabled"] = _orig_enabled
                setattr(model.encoder, "aop_prune_config", aop_cfg)

    if not all_mid:
        return np.array([]), np.array([]), []

    cand_mid_embeds  = torch.cat(all_mid,  dim=0).numpy()
    cand_last_embeds = torch.cat(all_last, dim=0).numpy()
    return cand_mid_embeds, cand_last_embeds, all_ids

def main():
    # ----------------------- Distributed init -----------------------
    if "RANK" in os.environ and dist.is_available() and not dist.is_initialized():
        dist.init_process_group(backend="nccl", timeout=datetime.timedelta(minutes=60))
    local_rank = dist.get_rank() if dist.is_initialized() else 0
    world_size = dist.get_world_size() if dist.is_initialized() else 1

    print_master("Distributed init debug info:")
    print_master(f"RANK: {os.environ.get('RANK')}")
    print_master(f"LOCAL_RANK: {os.environ.get('LOCAL_RANK')}")
    print_master(f"WORLD_SIZE: {os.environ.get('WORLD_SIZE')}")
    print_master(f"MASTER_ADDR: {os.environ.get('MASTER_ADDR')}")
    print_master(f"MASTER_PORT: {os.environ.get('MASTER_PORT')}")
    if dist.is_initialized():
        print_rank(f"dist.get_rank(): {dist.get_rank()}")
        print_rank(f"dist.get_world_size(): {dist.get_world_size()}")

    # 兼容 torchrun 参数
    for arg in sys.argv:
        if arg.startswith("--local-rank="):
            rank = arg.split("=")[1]
            sys.argv.remove(arg)
            sys.argv.append('--local_rank')
            sys.argv.append(rank)

    # ----------------------- Parse args -----------------------
    parser = HfArgumentParser((ModelArguments, DataArguments, TrainingArguments))
    model_args, data_args, training_args = parser.parse_args_into_dataclasses()
    model_args: ModelArguments
    data_args: DataArguments
    training_args: TrainingArguments
    os.makedirs(data_args.encode_output_path, exist_ok=True)

    # 支持多层评测（优先 LM_LAYERS，兼容 MID_LM_LAYER）
    layers_to_eval = get_env_eval_layers()
    print_master(f"Eval layers (qry/tgt): {layers_to_eval}")

    # ----------------------- Model loading -----------------------
    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'Model Backbone: {model_args.model_backbone}')

    # 仅 rank0 下载，其他rank等待缓存
    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] Loading the model from Huggingface: {model_args.model_name}...")
    if torch.distributed.is_initialized():
        torch.distributed.barrier()
    if local_rank != 0:
        print_rank(f"Loading the model from cache...")
        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)

    # ---- NEW: AOP 剪裁配置注入（驱动底模里已实现的 AOP 逻辑）----
    aop_cfg = get_env_aop_config()
    if aop_cfg["enabled"]:
        # 把配置塞到底模；底模 forward 中读取该 dict 并执行剪裁
        setattr(model.encoder, "aop_prune_config", aop_cfg)
        # 可选：为了便于在判定层取注意力或手算 qk，覆盖注意力实现
        attn_override = aop_cfg.get("attn_impl_override", "")
        if attn_override:
            try:
                if 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] override attn impl: {prev} -> {attn_override} (仅测试建议)")
            except Exception as e:
                print_master(f"[AOP] try override attn impl failed: {e}")

        print_master("[AOP] AOP-Prune enabled with config: " + json.dumps({
            "apply_to": aop_cfg["apply_to"],
            "layer_idx": aop_cfg["layer_idx"],
            "mode": aop_cfg["mode"],
            "delta": aop_cfg["delta"],
            "K_hat": aop_cfg["K_hat"],
            "keep_ratio": aop_cfg["keep_ratio"],
            "min_keep": aop_cfg["min_keep"],
            "use_bias": aop_cfg["use_bias"],
            "margin_mid?": (aop_cfg["margin_mid"] is not None),
            "prune_text": aop_cfg.get("prune_text", False),
            "keep_ratio_text": aop_cfg.get("keep_ratio_text", None),
            "keep_ratio_vision": aop_cfg.get("keep_ratio_vision", None),
            "selection": aop_cfg.get("selection", "aop"),
            "attn_agg": aop_cfg.get("attn_agg", "mean"),
        }))
    else:
        print_master("[AOP] disabled (set AOP_ENABLED=1 to enable)")
        
    # 确保“最后一层”时不裁层（避免类里默认20层的坑）
    model.set_inference_layers(qry_layers=None, tgt_layers=None)

    with open(data_args.dataset_config, 'r') as yaml_file:
        dataset_configs = yaml.safe_load(yaml_file)

    # ----------------------- Main evaluation loop -----------------------
    for dataset_idx, (dataset_name, task_config) in enumerate(dataset_configs.items()):
        if dist.is_initialized():
            dist.barrier()
        print_master(f"\n--- Evaluating {dataset_name} ---")

        # 根据 data_basedir 修正路径
        if data_args.data_basedir is not None:
            for key in ["image_root", "video_root", "frame_root", "clip_root", "data_path"]:
                if data_args.data_basedir and task_config.get(key):
                    task_config[key] = os.path.join(data_args.data_basedir, task_config[key])

        # 构建数据集
        full_eval_qry_dataset, corpus = AutoEvalPairDataset.instantiate(model_args=model_args, data_args=data_args, **task_config)
        full_eval_cand_dataset = generate_cand_dataset(full_eval_qry_dataset, corpus)
        eval_qry_dataset, eval_cand_dataset = full_eval_qry_dataset, full_eval_cand_dataset

        if dist.is_initialized():
            world_size = dist.get_world_size()
            padded_qry_dataset, _ = pad_dataset_to_divisible(full_eval_qry_dataset, world_size)
            padded_cand_dataset, _ = pad_dataset_to_divisible(full_eval_cand_dataset, world_size)
            eval_qry_dataset = split_dataset_by_node(padded_qry_dataset, rank=local_rank, world_size=world_size)
            eval_cand_dataset = split_dataset_by_node(padded_cand_dataset, rank=local_rank, world_size=world_size)
        else:
            padded_qry_dataset, padded_cand_dataset = full_eval_qry_dataset, full_eval_cand_dataset


        # === EE-only: 仅在线早停推理（先确保两份 candidate 向量）===
        ee_cfg = get_env_ee_config()
        assert ee_cfg["enabled"], "EE_ENABLED must be 1 for EE-only pipeline."

        # 依据 EE_LAYER 构造 tag
        mid_layer = int(ee_cfg["layer"])
        mid_tag = make_layer_tag(mid_layer)     # e.g., layer12
        last_tag = "layerlast"

        # 准备路径
        cand_mid_path  = os.path.join(data_args.encode_output_path, f"{dataset_name}_tgt_{mid_tag}")
        cand_last_path = os.path.join(data_args.encode_output_path, f"{dataset_name}_tgt_{last_tag}")

        # 构造 cand DataLoader（一次性，不切分）
        eval_cand_collator = MultimodalEvalDataCollator(processor, model_args, data_args, "cand")
        eval_cand_loader = DataLoader(
            full_eval_cand_dataset,
            batch_size=training_args.per_device_eval_batch_size,
            collate_fn=eval_cand_collator,
            num_workers=training_args.dataloader_num_workers
        )

        # === 替换为：一次前向，导出 cand 的 mid/last 两份向量 ===
        need_mid  = (not os.path.exists(cand_mid_path))
        need_last = (not os.path.exists(cand_last_path))

        if need_mid or need_last:
            print_master(f"[{dataset_name}] EE-only: encoding candidates BOTH layers in one pass (mid={mid_tag}, last={last_tag}) ...")
            # 走全层（不截层）
            model.set_inference_layers(qry_layers=None, tgt_layers=None)

            cand_embeds_mid, cand_embeds_last, all_cand_ids = encode_candidates_both_layers(
                model=model,
                loader=eval_cand_loader,
                training_args=training_args,
                model_args=model_args,
                full_dataset=full_eval_cand_dataset,
                mid_layer=mid_layer,
            )
            if local_rank == 0:
                if need_mid:
                    cand_embed_dict_mid = {cid: emb for cid, emb in zip(all_cand_ids, cand_embeds_mid)}
                    with open(cand_mid_path, "wb") as f:
                        pickle.dump(cand_embed_dict_mid, f)
                    print_master(f"[{dataset_name}] EE-only: saved {mid_tag} candidate embeddings -> {cand_mid_path}")
                if need_last:
                    cand_embed_dict_last = {cid: emb for cid, emb in zip(all_cand_ids, cand_embeds_last)}
                    with open(cand_last_path, "wb") as f:
                        pickle.dump(cand_embed_dict_last, f)
                    print_master(f"[{dataset_name}] EE-only: saved {last_tag} candidate embeddings -> {cand_last_path}")
        else:
            print_master(f"[{dataset_name}] EE-only: reuse existing candidates (mid={cand_mid_path}, last={cand_last_path})")

        if dist.is_initialized():
            dist.barrier()

        # 3) 在线早停门控 + 子集续跑（不做离线分层评分/曲线）
        if local_rank == 0:
            with open(cand_mid_path, "rb") as f:
                cand_mid_dict = pickle.load(f)
            with open(cand_last_path, "rb") as f:
                cand_last_dict = pickle.load(f)

            rank_global = task_config.get("eval_type", "global") == "global"
            print_master(f"[{dataset_name}] Run ONLINE early-exit at layer={ee_cfg['layer']}, method={ee_cfg['method']}, tau={ee_cfg['tau']}, topk={ee_cfg['topk']}, global={rank_global}")
            run_early_exit_queries(
                model=model,
                processor=processor,
                model_args=model_args,
                data_args=data_args,
                training_args=training_args,
                qry_dataset=full_eval_qry_dataset,   # 全量 query
                cand_mid_dict=cand_mid_dict,
                cand_last_dict=cand_last_dict,
                ee_cfg=ee_cfg,
                dataset_name=dataset_name,
                out_dir=data_args.encode_output_path,
                global_ranking=rank_global,
            )

        if dist.is_initialized():
            dist.barrier()
        # === EE-only 结束；直接进入下一个数据集 ===
        continue


if __name__ == '__main__':
    main()