dump_ee_phaseA_features.py

import os
import json
import math
import time
import random
import datetime
import numpy as np
import torch
import torch.distributed as dist
from tqdm import tqdm
from torch.utils.data import DataLoader
from transformers import HfArgumentParser, AutoConfig
import yaml

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


# ---------------- Utils ----------------
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_int(v: str, default=None):
    try:
        return int(v) if v is not None else default
    except Exception:
        return default


def _parse_float(v: str, default=None):
    try:
        return float(v) if v is not None else default
    except Exception:
        return default


def get_env_aop_config():
    enabled = _parse_bool(os.environ.get("AOP_ENABLED"), False)
    apply_to = os.environ.get("AOP_APPLY", "qry").strip().lower()
    layer_idx = _parse_int(os.environ.get("AOP_LAYER"), None)
    mode = os.environ.get("AOP_MODE", "ratio").strip().lower()
    prune_vision = _parse_bool(os.environ.get("AOP_PRUNE_VISION"), True)
    prune_text = _parse_bool(os.environ.get("AOP_PRUNE_TEXT"), False)
    keep_ratio_v = _parse_float(os.environ.get("AOP_KEEP_RATIO_VISION"), None)
    keep_ratio_t = _parse_float(os.environ.get("AOP_KEEP_RATIO_TEXT"), None)
    attn_agg = os.environ.get("AOP_ATTENTION_AGG", "mean").strip().lower()
    ee_layer = _parse_int(os.environ.get("EE_LAYER"), None)

    return {
        "enabled": enabled,
        "apply_to": apply_to,
        "layer_idx": layer_idx,
        "mode": mode,
        "prune_vision": prune_vision,
        "prune_text": prune_text,
        "keep_ratio_vision": keep_ratio_v,
        "keep_ratio_text": keep_ratio_t,
        "attn_agg": attn_agg,
        "ee_layer": ee_layer,
    }


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)
    padding_data = dataset.select([i % len(dataset) for i in range(num_to_add)])
    padded_dataset = concatenate_datasets([dataset, padding_data])
    return padded_dataset, num_samples + num_to_add


# ---------- Candidate encode in one pass (mid+last) ----------
@torch.no_grad()
def encode_candidates_both_layers(model: MMEBModel, loader: DataLoader, training_args: TrainingArguments, mid_layer: int):
    model.eval()
    all_mid, all_last, all_ids = [], [], []
    for inputs, infos in tqdm(loader, desc="[DUMP] Cands[BOTH]", disable=False):
        inputs = batch_to_device(inputs, training_args.device)
        # cand 侧不启用 AOP
        aop_cfg = getattr(model.encoder, "aop_prune_config", None)
        if isinstance(aop_cfg, dict) and aop_cfg:
            aop_off = dict(aop_cfg)
            aop_off["enabled"] = False
            setattr(model.encoder, "aop_prune_config", aop_off)

        with torch.autocast(device_type="cuda", dtype=torch.bfloat16, enabled=True):
            out = model.encoder(
                **inputs,
                return_dict=True,
                output_hidden_states=True,
                stop_at_layer=None,
                compute_lm_head=False,
            )
        hs_list = out.hidden_states
        assert hs_list is not None and len(hs_list) > mid_layer, "hidden_states too short for mid_layer"
        mid_hs, last_hs = hs_list[mid_layer], hs_list[-1]
        am = inputs.get("attention_mask", None)
        if am is not None and hasattr(am, "device") and am.device != mid_hs.device:
            am = am.to(mid_hs.device)
        reps_mid = model._pooling(mid_hs, am).detach().float().cpu()
        reps_last = model._pooling(last_hs, am).detach().float().cpu()
        all_mid.append(reps_mid)
        all_last.append(reps_last)
        all_ids.extend([info["cand_name"] for info in infos])
    cand_mid = torch.cat(all_mid, dim=0).numpy()
    cand_last = torch.cat(all_last, dim=0).numpy()
    return cand_mid, cand_last, all_ids


# ---------- Build Phase A features ----------
@torch.no_grad()
def build_phaseA_features_global(
    reps_mid_t: torch.Tensor,  # [B,D] GPU
    cand_mid_t: torch.Tensor,  # [Nc,D] GPU
    am_mid: torch.Tensor,  # [B,L]
    input_ids: torch.Tensor,  # [B,L]
    cfg,  # model.encoder.config
    topk: int = 200,
    temp: float = 0.05,
):
    device = reps_mid_t.device
    B = reps_mid_t.size(0)
    # 相似度类
    scores_t = reps_mid_t @ cand_mid_t.T
    k = min(topk, scores_t.size(1))
    vals_t, _ = torch.topk(scores_t, k=k, dim=1)
    s1 = vals_t[:, 0]
    s2 = vals_t[:, 1] if k >= 2 else torch.zeros_like(s1)
    margin = s1 - s2
    p_t = torch.softmax(vals_t / max(temp, 1e-6), dim=1)
    H = -(p_t * (torch.log(p_t + 1e-12))).sum(dim=1) / math.log(max(k, 1))
    sum_p2 = (p_t**2).sum(dim=1)

    # 长度比例
    am = am_mid.to(torch.bool)
    iid = input_ids
    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) & am

    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)
    is_text = am & (~is_vision) & (~is_special)

    L_vis = is_vision.sum(dim=1).float()
    L_txt = is_text.sum(dim=1).float()
    L_tot = am.sum(dim=1).float().clamp(min=1.0)
    r_vis = L_vis / L_tot
    r_txt = L_txt / L_tot

    # 类型 one-hot
    is_I = ((L_vis > 0) & (L_txt == 0)).float()
    is_T = ((L_txt > 0) & (L_vis == 0)).float()
    is_IT = ((L_txt > 0) & (L_vis > 0)).float()

    feats = torch.stack([s1, s2, margin, H, sum_p2, L_txt, L_vis, L_tot, r_txt, r_vis, is_I, is_T, is_IT], dim=1)
    return feats  # [B,13]


@torch.no_grad()
def build_phaseA_features_local(
    reps_mid_t: torch.Tensor,  # [B,D]
    cand_mid_t: torch.Tensor,  # [Nc,D]
    am_mid: torch.Tensor,  # [B,L]
    input_ids: torch.Tensor,  # [B,L]
    cfg,
    per_sample_rows: list,  # list[list[int]]
    topk: int = 200,
    temp: float = 0.05,
):
    device = reps_mid_t.device
    B = reps_mid_t.size(0)
    s1_list, s2_list, H_list, sum_p2_list = [], [], [], []
    for b in range(B):
        rows = per_sample_rows[b]
        if len(rows) == 0:
            s1_list.append(torch.tensor(0.0, device=device))
            s2_list.append(torch.tensor(0.0, device=device))
            H_list.append(torch.tensor(1.0, device=device))
            sum_p2_list.append(torch.tensor(0.0, device=device))
            continue
        cmat = cand_mid_t[rows]
        sv = (reps_mid_t[b:b+1] @ cmat.T)[0]
        k = min(topk, sv.size(0))
        vals, _ = torch.topk(sv, k=k, dim=0)
        s1_list.append(vals[0])
        s2_list.append(vals[1] if k >= 2 else torch.tensor(0.0, device=device, dtype=vals.dtype))
        p = torch.softmax(vals / max(temp, 1e-6), dim=0)
        H_list.append((-(p * (torch.log(p + 1e-12))).sum() / math.log(max(k, 1))))
        sum_p2_list.append((p**2).sum())
    s1 = torch.stack(s1_list)
    s2 = torch.stack(s2_list)
    H = torch.stack(H_list)
    sum_p2 = torch.stack(sum_p2_list)
    margin = s1 - s2

    am = am_mid.to(torch.bool)
    iid = input_ids
    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) & am

    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)
    is_text = am & (~is_vision) & (~is_special)

    L_vis = is_vision.sum(dim=1).float()
    L_txt = is_text.sum(dim=1).float()
    L_tot = am.sum(dim=1).float().clamp(min=1.0)
    r_vis = L_vis / L_tot
    r_txt = L_txt / L_tot

    is_I = ((L_vis > 0) & (L_txt == 0)).float()
    is_T = ((L_txt > 0) & (L_vis == 0)).float()
    is_IT = ((L_txt > 0) & (L_vis > 0)).float()

    feats = torch.stack([s1, s2, margin, H, sum_p2, L_txt, L_vis, L_tot, r_txt, r_vis, is_I, is_T, is_IT], dim=1)
    return feats  # [B,13]


# ---------- Label(y_exit) ----------
def compute_label_top1_equal_global(scores_mid: np.ndarray, scores_last: np.ndarray) -> np.ndarray:
    top1_mid = scores_mid.argmax(axis=1)
    top1_last = scores_last.argmax(axis=1)
    return (top1_mid == top1_last).astype(np.int32)


def compute_label_top1_equal_local(scores_mid_list, scores_last_list):
    y = []
    for sv_mid, sv_last in zip(scores_mid_list, scores_last_list):
        if sv_mid.size == 0 or sv_last.size == 0:
            y.append(0)
        else:
            y.append(int(int(sv_mid.argmax()) == int(sv_last.argmax())))
    return np.array(y, dtype=np.int32)


# ---------------- Main dump ----------------
def main():
    # DDP 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

    parser = HfArgumentParser((ModelArguments, DataArguments, TrainingArguments))
    model_args, data_args, training_args = parser.parse_args_into_dataclasses()
    os.makedirs(data_args.encode_output_path, exist_ok=True)

    # Load model
    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)

    if local_rank == 0:
        processor = load_processor(model_args, data_args)
        model = MMEBModel.load(model_args, is_trainable=False, processor=processor)
    if dist.is_initialized():
        dist.barrier()
    if local_rank != 0:
        processor = load_processor(model_args, data_args)
        time.sleep(random.randint(2 * local_rank, 3 * local_rank))
        model = MMEBModel.load(model_args, is_trainable=False, processor=processor)

    model.eval()
    model = model.to(training_args.device, dtype=torch.bfloat16)

    # AOP env
    aop_cfg_env = get_env_aop_config()

    # EE layer / feat topk / label rule
    ee_layer = int(os.environ.get("EE_LAYER", os.environ.get("AOP_LAYER", "12")))
    feat_topk = int(os.environ.get("EE_FEAT_TOPK", "200"))
    force_no_aop = os.environ.get("DUMP_EXIT_NO_AOP", "1").strip().lower() in {"1", "true", "yes", "on"}
    print_master(f"[DUMP] EE_LAYER={ee_layer}, FEAT_TOPK={feat_topk}, NO_AOP={force_no_aop}")

    # load datasets yaml
    with open(data_args.dataset_config, 'r', encoding='utf-8') as yf:
        dataset_configs = yaml.safe_load(yf)

    for dataset_name, task_cfg in dataset_configs.items():
        if dist.is_initialized():
            dist.barrier()
        print_master(f"\n[DUMP] {dataset_name}")

        # fix paths by data_basedir
        if data_args.data_basedir:
            for key in ["image_root", "video_root", "frame_root", "clip_root", "data_path"]:
                if task_cfg.get(key):
                    task_cfg[key] = os.path.join(data_args.data_basedir, task_cfg[key])

        full_qry, corpus = AutoEvalPairDataset.instantiate(model_args=model_args, data_args=data_args, **task_cfg)
        full_cand = generate_cand_dataset(full_qry, corpus)

        qry_collator = MultimodalEvalDataCollator(processor, model_args, data_args, "qry")
        cand_collator = MultimodalEvalDataCollator(processor, model_args, data_args, "cand")
        qry_loader = DataLoader(
            full_qry,
            batch_size=training_args.per_device_eval_batch_size,
            collate_fn=qry_collator,
            num_workers=training_args.dataloader_num_workers,
        )
        cand_loader = DataLoader(
            full_cand,
            batch_size=training_args.per_device_eval_batch_size,
            collate_fn=cand_collator,
            num_workers=training_args.dataloader_num_workers,
        )

        # cands mid/last (one pass)
        cand_mid_np, cand_last_np, cand_ids = encode_candidates_both_layers(model, cand_loader, training_args, mid_layer=ee_layer)
        cand_id2row = {str(cid): i for i, cid in enumerate(cand_ids)}
        device = training_args.device
        cand_mid_t = torch.from_numpy(cand_mid_np).to(device=device, dtype=torch.bfloat16)
        cand_last_t = None  # 延迟到用到时再搬 GPU

        # output paths
        feat_path = os.path.join(data_args.encode_output_path, f"{dataset_name}_phaseA_features.jsonl")
        scaler_path = os.path.join(data_args.encode_output_path, f"{dataset_name}_phaseA_scaler.json")

        # stats for scaler
        sum_feat = None
        sum2_feat = None
        n_feat = 0

        # ranking mode
        rank_global = task_cfg.get("eval_type", "global") == "global"

        qid_global = 0
        with open(feat_path, "w", encoding="utf-8") as fout:
            for inputs, infos in tqdm(qry_loader, desc=f"[DUMP] Qrys", disable=False):
                inputs = batch_to_device(inputs, device)
                B = inputs["input_ids"].size(0)

                # --- mid 前向: 强制禁用 AOP (Phase A建议无AOP) ---
                aop_cfg_cur = getattr(model.encoder, "aop_prune_config", None)
                orig_aop = None
                if force_no_aop and isinstance(aop_cfg_cur, dict):
                    orig_aop = dict(aop_cfg_cur)
                    aop_off = dict(aop_cfg_cur)
                    aop_off["enabled"] = False
                    setattr(model.encoder, "aop_prune_config", aop_off)

                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_layer),
                        compute_lm_head=False,
                        return_intermediate_state=True,
                    )

                if orig_aop is not None:
                    setattr(model.encoder, "aop_prune_config", orig_aop)

                hs_mid = getattr(out_mid, "last_hidden_state", None)
                if hs_mid is None:
                    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)

                # --- Phase A 轻量特征 ---
                if rank_global:
                    feats_t = build_phaseA_features_global(reps_mid_t, cand_mid_t, am_mid, inputs["input_ids"], model.encoder.config, topk=feat_topk)
                else:
                    rows_list = []
                    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]
                        rows_list.append(rows)
                    feats_t = build_phaseA_features_local(reps_mid_t, cand_mid_t, am_mid, inputs["input_ids"], model.encoder.config, rows_list, topk=feat_topk)
                feats_np = feats_t.detach().float().cpu().numpy()  # [B, 13]

                # --- last 续跑 (用于 y_exit 标签) ---
                # 强制禁用 AOP
                aop_cfg_cur = getattr(model.encoder, "aop_prune_config", None)
                orig_aop2 = None
                if force_no_aop and isinstance(aop_cfg_cur, dict):
                    orig_aop2 = dict(aop_cfg_cur)
                    aop_off2 = dict(aop_cfg_cur)
                    aop_off2["enabled"] = False
                    setattr(model.encoder, "aop_prune_config", aop_off2)

                interm = getattr(out_mid, "intermediate_state", None)
                assert interm is not None, "intermediate_state missing; ensure return_intermediate_state=True"
                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"])
                resume_state = {
                    "hidden_states": hs,
                    "attention_mask": am,
                    "position_ids": pos,
                    "vision_mask": vm,
                    "text_mask": tm,
                    "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,
                    )

                if orig_aop2 is not None:
                    setattr(model.encoder, "aop_prune_config", orig_aop2)

                hs_last = getattr(out_last, "last_hidden_state", None)
                if hs_last is None:
                    hs_last = out_last.hidden_states[-1]
                am_last = getattr(out_last, "attention_mask", None)
                if am_last is None:
                    am_last = am
                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)

                # --- 计算标签 y_exit: top1_equal ---
                if rank_global:
                    if cand_last_t is None:
                        cand_last_t = torch.from_numpy(cand_last_np).to(device=device, dtype=torch.bfloat16)
                    sim_mid = (reps_mid_t @ cand_mid_t.T).detach().float().cpu().numpy()
                    sim_last = (reps_last_t @ cand_last_t.T).detach().float().cpu().numpy()
                    y = compute_label_top1_equal_global(sim_mid, sim_last)
                else:
                    y_list = []
                    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:
                            y_list.append(0)
                            continue
                        c_mid = cand_mid_t[rows]
                        if cand_last_t is None:
                            cand_last_t = torch.from_numpy(cand_last_np).to(device=device, dtype=torch.bfloat16)
                        c_last = cand_last_t[rows]
                        sv_mid = (reps_mid_t[b:b+1] @ c_mid.T)[0].detach().float().cpu().numpy()
                        sv_last = (reps_last_t[b:b+1] @ c_last.T)[0].detach().float().cpu().numpy()
                        y_list.append(int(int(sv_mid.argmax()) == int(sv_last.argmax())))
                    y = np.array(y_list, dtype=np.int32)

                # --- 写样本 & 累计 scaler ---
                # 类型：从特征恢复
                L_txt = feats_np[:, 5]
                L_vis = feats_np[:, 6]
                types = np.where((L_vis > 0) & (L_txt == 0), "I", np.where((L_txt > 0) & (L_vis == 0), "T", "IT"))

                for b in range(B):
                    row = {
                        "dataset": dataset_name,
                        "qid": int(qid_global + b),
                        "type": str(types[b]),
                        "feats": feats_np[b].tolist(),
                        "y_exit": int(y[b]),
                    }
                    fout.write(json.dumps(row, ensure_ascii=False) + "\n")

                # scaler
                if sum_feat is None:
                    sum_feat = feats_np.sum(axis=0)
                    sum2_feat = (feats_np**2).sum(axis=0)
                else:
                    sum_feat += feats_np.sum(axis=0)
                    sum2_feat += (feats_np**2).sum(axis=0)
                n_feat += feats_np.shape[0]
                qid_global += B

        # 保存 scaler
        if n_feat > 0 and (not dist.is_initialized() or local_rank == 0):
            mean = (sum_feat / n_feat).tolist()
            var = (sum2_feat / n_feat - (sum_feat / n_feat) ** 2).tolist()
            std = [float(max(1e-6, math.sqrt(max(0.0, v)))) for v in var]
            with open(scaler_path, "w", encoding="utf-8") as f:
                json.dump({"mean": mean, "std": std, "in_dim": len(mean), "n": n_feat}, f, indent=2, ensure_ascii=False)
            print_master(f"[DUMP] {dataset_name}: features -> {feat_path}, scaler -> {scaler_path}, n={n_feat}")

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


# datasets.concatenate_datasets 需要导入（上方 pad_dataset_to_divisible 用，若不使用可忽略）
from datasets import concatenate_datasets

if __name__ == "__main__":
    main()