src/model/model_gp.py

from typing import Dict
import torch
import torch.distributed as dist
from torch import nn, Tensor
from transformers import PreTrainedModel, AutoModelForCausalLM, AutoConfig
from src.model.vlm_backbone.qwen2_5_vl_gp.configuration import Qwen2_5_VL_GPConfig
from src.model.vlm_backbone.qwen2_vl_gp.configuration import Qwen2VL_GPConfig
from peft import LoraConfig, get_peft_model, PeftModel
from src.model.processor import QWEN2_5_VL_TOKENSELECTION
from src.arguments_gp import ModelArguments, TrainingArguments
from src.model.processor import LLAVA_NEXT, QWEN2_VL, QWEN2_VL_GP, PHI3V, get_backbone_name, print_master, QWEN2_5_VL, \
    backbone2model, QWEN2_VL_TOKENSELECTION, QWEN2_5_VL_TOKENSELECTION, E5_V

from src.arguments_gp import ModelArguments
from src.model.processor import LLAVA_NEXT, QWEN2_VL, QWEN2_VL_GP, QWEN2_5_VL_GP, PHI3V, get_backbone_name, print_master, QWEN2_5_VL, INTERNVIDEO2, \
    QWEN2_VL_TOKENSELECTION, backbone2model, GME, VLM_IMAGE_TOKENS, LamRA, LamRA_QWEN2_5, COLPALI
from src.model.baseline_backbone.colpali import ColPali
from src.model.baseline_backbone.gme.gme_inference import GmeQwen2VL
from src.model.baseline_backbone.lamra.lamra_inference import LamRAQwen2VL
from src.model.baseline_backbone.lamra.lamra_qwen25_inference import LamRAQwen25VL
from src.model.baseline_backbone.phi3_v.modeling_phi3_v import Phi3VForCausalLM
from src.model.baseline_backbone.llava_next import LlavaNextForConditionalGeneration

from transformers import modeling_utils
if not hasattr(modeling_utils, "ALL_PARALLEL_STYLES") or modeling_utils.ALL_PARALLEL_STYLES is None:
    modeling_utils.ALL_PARALLEL_STYLES = ["tp", "none", "colwise", 'rowwise']

# ADD mask output
class MMEBModel(nn.Module):
    TRANSFORMER_CLS = AutoModelForCausalLM

    def __init__(self,
                 encoder: PreTrainedModel,
                 pooling: str = 'last',
                 normalize: bool = False,
                 temperature: float = 0.02,
                 ):
        super().__init__()
        self.config = encoder.config
        self.encoder = encoder
        self.pooling = pooling
        self.normalize = normalize
        self.temperature = temperature
        self.cross_entropy = nn.CrossEntropyLoss(reduction='mean')
        self.is_ddp = dist.is_initialized()
        if self.is_ddp:
            self.process_rank = dist.get_rank()
            self.world_size = dist.get_world_size()

        self._last_image_token_bool_masks = None
        self._last_image_token_mask_logits = None
        self._last_le_loss = None

    def encode_input(self, input):
        self._last_image_token_bool_masks = None
        if getattr(self, "model_backbone", None) == INTERNVIDEO2:
            if "input_ids" in input.keys():
                # text side
                text_output = self.encoder.get_text_encoder()(
                    input["input_ids"],
                    attention_mask=input["attention_mask"],
                    return_dict=True,
                    mode="text",
                )
                text_embeds = text_output.last_hidden_state
                pooled_text_embeds = text_embeds[:, 0]
                pooled_output = self.encoder.text_proj(pooled_text_embeds)
                pooled_output /= pooled_output.norm(dim=-1, keepdim=True)
                return pooled_output
            else:
                _, vfeat = self.encoder.encode_vision(input["pixel_values"], test=True)
                vfeat = self.encoder.vision_proj(vfeat)
                vfeat /= vfeat.norm(dim=-1, keepdim=True)
                return vfeat
        elif getattr(self, "model_backbone", None) in [GME, LamRA, LamRA_QWEN2_5]:
            # pooled_output = self.encoder(**input, return_dict=True, output_hidden_states=True)
            texts = [text.replace(VLM_IMAGE_TOKENS[QWEN2_VL] + '\n', '') for text in input["texts"]] # we are actually passing video queries so this should not happen
            images = []
            for imgs in input['images']:
                # if multi images are given, select the middle frame only
                if isinstance(imgs, list):
                    imgs = imgs[len(imgs) // 2]
                    assert not isinstance(imgs, list) # make sure we have extracted the middle frame and it is no longer a list
                    images.append(imgs)
                else:
                    images.append(imgs)
            pooled_output = self.encoder.get_fused_embeddings(texts=texts, images=images)
            return pooled_output
        elif getattr(self, "model_backbone", None) == COLPALI:
            pooled_output = self.encoder(**input, return_dict=True, output_hidden_states=True)
            return pooled_output
        elif getattr(self, "model_backbone", None) == LLAVA_NEXT:
            input['pixel_values'] = input['pixel_values'].squeeze(dim=1)
            input['image_sizes'] = input['image_sizes'].squeeze(dim=1)
            hidden_states = self.encoder(**input, return_dict=True, output_hidden_states=True)
            hidden_states = hidden_states.hidden_states[-1]
            pooled_output = self._pooling(hidden_states, input['attention_mask'])
            return pooled_output
        elif getattr(self, "model_backbone", None) in {QWEN2_VL_GP, QWEN2_5_VL_GP}:
            # do_sel = self.gp_do_selection
            # print('do_sel:', do_sel)
            # exit()
            outputs = self.encoder(
                **input,
                return_dict=True,
                output_hidden_states=True,         # 如需可设 False
                do_selection=True,
                delay_selection=False, # if do_sel else True,  # 开裁剪时立刻裁剪；训练默认不裁剪
                use_cache=False,
                use_ref_masks=False,               # 推理也用预测掩码
            )
            self._last_image_token_bool_masks = getattr(outputs, "image_token_bool_masks", None)
            self._last_image_token_mask_logits = getattr(outputs, "image_token_mask_logits", None)
            self._last_le_loss = getattr(outputs, "le_loss", None)

            hs = outputs.hidden_states
            last_hidden = hs[-1] if isinstance(hs, (list, tuple)) else hs
            attn_mask = getattr(outputs, "attention_mask", None)
            if attn_mask is None:
                attn_mask = input["attention_mask"]
            attn = attn_mask
            pooled_output = self._pooling(last_hidden, attn)
            return pooled_output


        else:
            # 默认分支同样要鲁棒处理（有可能用到带 glimpse 的模型）
            outputs = self.encoder(**input, return_dict=True, output_hidden_states=True)

            img_masks = getattr(outputs, "image_token_bool_masks", None)
            if img_masks is None:
                try:
                    img_masks = outputs.get("image_token_bool_masks", None)
                except Exception:
                    pass
            self._last_image_token_bool_masks = img_masks

            hs = outputs.hidden_states
            if isinstance(hs, (list, tuple)):
                last_hidden = hs[-1]            # [B, L, D]
            else:
                last_hidden = hs                 # [B, L, D]（如果底模走了 glimpse 路径）
            attn = getattr(outputs, "attention_mask", None)
            if attn is None:
                attn = input["attention_mask"]

            pooled_output = self._pooling(last_hidden, attn)
            return pooled_output

    def _pooling(self, last_hidden_state, attention_mask):
        if self.pooling == 'last' or self.pooling == 'eos':
            left_padding = (attention_mask[:, -1].sum() == attention_mask.shape[0])
            batch_size = last_hidden_state.shape[0]
            if left_padding:
                # Get the vectors at the last position
                reps = last_hidden_state[torch.arange(batch_size), -1, :]
            else:
                # Calculate last 1 position in the original tensor
                eos_indices = attention_mask.sum(dim=1) - 1
                # Get the vectors at the last 1 position of each attention mask
                reps = last_hidden_state[
                    torch.arange(batch_size, device=last_hidden_state.device), eos_indices]
        else:
            raise NotImplementedError
        if self.normalize:
            reps = torch.nn.functional.normalize(reps, p=2, dim=-1)
        return reps

    @classmethod
    def build(cls, model_args: ModelArguments, **kwargs):
        config = AutoConfig.from_pretrained(model_args.model_name, trust_remote_code=True)
        model_backbone = get_backbone_name(hf_config=config)
        print_master(f'Loading backbone [{model_backbone}] from {model_args.model_name}')
        # Loading the base model
        if model_backbone == PHI3V:
            config._attn_implementation = "eager"
            config.padding_side = "right"
            config.use_cache = False
            base_model = Phi3VForCausalLM.from_pretrained(
                model_args.model_name,
                config=config,
                torch_dtype=torch.bfloat16,
                low_cpu_mem_usage=True,
            )
        elif model_backbone == LLAVA_NEXT:
            config.use_cache = False
            config.padding_side = "left"
            base_model = LlavaNextForConditionalGeneration.from_pretrained(
                model_args.model_name,
                config=config,
                torch_dtype=torch.bfloat16,
                low_cpu_mem_usage=True,
            )
        elif model_backbone in [QWEN2_VL, QWEN2_5_VL]:
            config._attn_implementation = "flash_attention_2"
            config.padding_side = "left"
            config.use_cache = False
            base_model = backbone2model[model_backbone].from_pretrained(
                model_args.model_name,
                config=config,
                torch_dtype=torch.bfloat16,
                low_cpu_mem_usage=True,
            )
        elif model_backbone in [QWEN2_VL_GP, QWEN2_5_VL_GP]:
            config._attn_implementation = "flash_attention_2"
            config.padding_side = "left"
            config.use_cache = False
            base_model = backbone2model[model_backbone].from_pretrained(
                model_args.model_name,
                config=config,
                torch_dtype=torch.bfloat16,
                low_cpu_mem_usage=True,
            )
        elif model_backbone in [QWEN2_VL_TOKENSELECTION, QWEN2_5_VL_TOKENSELECTION]:
            config._attn_implementation = "flash_attention_2"
            config.padding_side = "left"
            config.use_cache = False

            from .utils import parse_layer_type
            lm_qwen_layer = 28
            vis_qwen_layer = 32
            lm_skip_layer = parse_layer_type(model_args.lm_skip_layer, lm_qwen_layer)
            vis_skip_layer = parse_layer_type(model_args.vis_skip_layer, vis_qwen_layer)

            base_model = backbone2model[model_backbone].from_pretrained(
                model_args.model_name,
                config=config,
                torch_dtype=torch.bfloat16,
                low_cpu_mem_usage=True,
                lm_skip_layer=lm_skip_layer,
                vis_skip_layer=vis_skip_layer,
            )
        else:
            config.use_cache = False
            base_model = cls.TRANSFORMER_CLS.from_pretrained(
                model_args.model_name, **kwargs, config=config,
                attn_implementation="flash_attention_2",
                torch_dtype=torch.bfloat16,
                trust_remote_code=True)

        if model_args.lora:
            print_master(f'Loading lora adapter from {base_model}')
            lora_config = LoraConfig(
                r=model_args.lora_r,
                lora_alpha=model_args.lora_alpha,
                target_modules=model_args.lora_target_modules.split(','),
                lora_dropout=model_args.lora_dropout,
                init_lora_weights="gaussian",
                use_dora=True,
                inference_mode=False
            )
            lora_model = get_peft_model(base_model, lora_config)
            model = cls(
                encoder=lora_model,
                pooling=model_args.pooling,
                normalize=model_args.normalize,
                temperature=model_args.temperature
            )
        else:
            model = cls(
                encoder=base_model,
                pooling=model_args.pooling,
                normalize=model_args.normalize,
                temperature=model_args.temperature
            )
        # model.gp_do_selection = getattr(model_args, "gp_do_selection", False)
        return model


    @classmethod
    def load(cls, model_args: ModelArguments, is_trainable=True, **kwargs):
        # Loading the base model
        model_name_or_path = model_args.checkpoint_path if model_args.checkpoint_path else model_args.model_name
        config = AutoConfig.from_pretrained(model_name_or_path, trust_remote_code=True)
        if not hasattr(model_args, "model_backbone") or not model_args.model_backbone:
            model_backbone = get_backbone_name(hf_config=config, model_type=model_args.model_type)
            setattr(model_args, 'model_backbone', model_backbone)
        # print('model_args.model_backbone', model_args.model_backbone)
        # exit()
        print_master(f'Loading backbone [{model_args.model_backbone}] from {model_name_or_path}')
        if model_args.model_backbone in {LLAVA_NEXT, QWEN2_VL, QWEN2_5_VL, QWEN2_VL_TOKENSELECTION, QWEN2_5_VL_TOKENSELECTION, E5_V}:
            config = AutoConfig.from_pretrained(model_args.model_name, trust_remote_code=True)
            config._attn_implementation = "flash_attention_2"
            config.vision_config._attn_implementation = "flash_attention_2"
            base_model = backbone2model[model_args.model_backbone].from_pretrained(
                model_args.model_name,
                torch_dtype=torch.bfloat16,
                low_cpu_mem_usage=True,
                config=config
            )
        elif model_args.model_backbone in {QWEN2_5_VL_GP}:
            new_modules_dir = getattr(model_args, "new_modules_dir", None) or kwargs.get("new_modules_dir", None)

            if new_modules_dir:
                import os
                from huggingface_hub import hf_hub_download
                def _load_gp_config_from_dir_or_repo(repo_or_dir: str) -> Qwen2_5_VL_GPConfig:
                    # 支持本地目录或 HF 仓库 ID
                    if os.path.isdir(repo_or_dir):
                        cfg_path = os.path.join(repo_or_dir, "config.json")
                        if not os.path.exists(cfg_path):
                            raise FileNotFoundError(f"config.json not found in {repo_or_dir}")
                    else:
                        # 从 HF 仓库下载 config.json
                        cfg_path = hf_hub_download(repo_id=repo_or_dir, filename="config.json")

                    return Qwen2_5_VL_GPConfig.from_json_file(cfg_path)
                gp_cfg = _load_gp_config_from_dir_or_repo(new_modules_dir)
                gp_cfg._attn_implementation = "flash_attention_2"
                gp_cfg.padding_side = "left"
                gp_cfg.use_cache = False
                try:
                    gp_cfg.vision_config._attn_implementation = "flash_attention_2"
                except Exception:
                    if isinstance(gp_cfg.vision_config, dict):
                        gp_cfg.vision_config["_attn_implementation"] = "flash_attention_2"

                base_model = backbone2model[model_args.model_backbone].from_pretrained(
                    model_args.model_name,
                    torch_dtype=torch.bfloat16,
                    low_cpu_mem_usage=True,
                    config=gp_cfg,
                )
                # print('base_model:', base_model)
                # exit()
                base_model.load_new_modules(new_modules_dir)

                try:
                    print_master(f"[GP] attn_fuser = {type(base_model.attn_fuser).__name__}")
                except Exception:
                    print_master("[GP] attn_fuser not found (unexpected).")

            else:
                # 没有 new_modules_dir：保留你原来的 fallback（用 base config 构建 GPConfig，不加载新模块权重）
                base_cfg = AutoConfig.from_pretrained(model_args.model_name, trust_remote_code=True)
                gp_cfg = Qwen2_5_VL_GPConfig(**base_cfg.to_dict())
                gp_cfg._attn_implementation = "flash_attention_2"
                try:
                    gp_cfg.vision_config._attn_implementation = "flash_attention_2"
                except Exception:
                    if isinstance(gp_cfg.vision_config, dict):
                        gp_cfg.vision_config["_attn_implementation"] = "flash_attention_2"
                gp_cfg.padding_side = "left"
                gp_cfg.use_cache = False

                L = getattr(base_cfg, "num_hidden_layers", None) or 28
                gp_cfg.reduce_layer = min(max(L - 2, 1), L - 1)
                gp_cfg.selected_layers = [max(gp_cfg.reduce_layer - 2, 0), gp_cfg.reduce_layer]
                # gp_cfg.attn_fuse_type = "AttnFuserV1"
                # gp_cfg.selected_visual_layers = (8,)
                # gp_cfg.use_attention_logits = False
                # gp_cfg.use_zero_masks = False
                # gp_cfg.use_ref_masks = False
                # gp_cfg.le_layers = []
                # gp_cfg.le_length = 0

                base_model = backbone2model[model_args.model_backbone].from_pretrained(
                    model_args.model_name,
                    torch_dtype=torch.bfloat16,
                    low_cpu_mem_usage=True,
                    config=gp_cfg,
                )
        elif model_args.model_backbone == QWEN2_VL_GP:
            import os
            from huggingface_hub import snapshot_download
            # from transformers import AutoConfig
            from transformers.utils import is_flash_attn_2_available

            new_modules_dir = getattr(model_args, "new_modules_dir", None) or kwargs.get("new_modules_dir", None)
            if not new_modules_dir:
                raise ValueError("new_modules_dir is required for QWEN2_VL_GP")

            # 统一成本地目录
            if os.path.isdir(new_modules_dir):
                local_new_modules_dir = new_modules_dir
            else:
                local_new_modules_dir = snapshot_download(repo_id=new_modules_dir)

            # 读取 GP 模块自带的配置（只用于校验，不拿它来构建 base）
            gp_cfg_path = os.path.join(local_new_modules_dir, "config.json")
            if not os.path.exists(gp_cfg_path):
                raise FileNotFoundError(f"config.json not found in {local_new_modules_dir}")
            gp_cfg = Qwen2VL_GPConfig.from_json_file(gp_cfg_path)

            # 用 base 模型自己的 config 来构建，避免与 checkpoint 维度不一致
            base_cfg = AutoConfig.from_pretrained(model_args.model_name, trust_remote_code=True)

            # 早检测：关键结构参数必须一致（尤其 hidden_size）
            if hasattr(gp_cfg, "hidden_size") and hasattr(base_cfg, "hidden_size"):
                if gp_cfg.hidden_size != base_cfg.hidden_size:
                    raise ValueError(
                        f"[Mismatch] base_model hidden_size={base_cfg.hidden_size} "
                        f"!= gp_modules hidden_size={gp_cfg.hidden_size}. "
                        f"请更换为匹配的 base 模型或对应的 new_modules_dir。"
                    )

            # 只做运行时补丁，不改动结构
            base_cfg._attn_implementation = "flash_attention_2" if is_flash_attn_2_available() else "sdpa"
            base_cfg.padding_side = "left"
            base_cfg.use_cache = False
            if hasattr(base_cfg, "vision_config"):
                vc = base_cfg.vision_config
                if isinstance(vc, dict):
                    vc["_attn_implementation"] = base_cfg._attn_implementation
                else:
                    setattr(vc, "_attn_implementation", base_cfg._attn_implementation)

            # 构建并加载 base 权重（与 checkpoint 结构一致）
            ModelClass = backbone2model[model_args.model_backbone]
            base_model = ModelClass.from_pretrained(
                model_args.model_name,
                torch_dtype=torch.bfloat16,
                low_cpu_mem_usage=True,
                config=base_cfg,
                trust_remote_code=True,
            )

            # 注入 GP 模块（目录必须与 base 匹配）
            base_model.load_new_modules(local_new_modules_dir)

            try:
                print_master(f"[GP] attn_fuser = {type(base_model.attn_fuser).__name__}")
            except Exception:
                print_master("[GP] attn_fuser not found (unexpected).")

        elif model_args.model_backbone == PHI3V:
            config = AutoConfig.from_pretrained(model_args.model_name, trust_remote_code=True)
            config.use_cache = False
            config.padding_side = "right"
            base_model = Phi3VForCausalLM.from_pretrained(model_args.model_name, **kwargs, config=config,
                                                          torch_dtype=torch.bfloat16, trust_remote_code=True)
            base_model.padding_side = "right"
        elif model_args.model_backbone == INTERNVIDEO2:
            print_master(f'Loading backbone [{model_args.model_backbone}] from {"src/model/vlm_backbone/internvideo2/"}')
            config = AutoConfig.from_pretrained("src/model/vlm_backbone/internvideo2/",
                                                trust_remote_code=True)
            base_model = backbone2model[model_args.model_backbone].from_pretrained("src/model/vlm_backbone/internvideo2/", config=config,
                                                                                   trust_remote_code=True)
        elif model_args.model_backbone == GME:
            base_model = GmeQwen2VL(model_args.model_name, processor=kwargs['processor'])
            setattr(base_model, 'config', config)
        elif model_args.model_backbone == LamRA:
            base_model = LamRAQwen2VL(model_args.model_name)
            setattr(base_model, 'config', config)
        elif model_args.model_backbone == LamRA_QWEN2_5:
            base_model = LamRAQwen25VL(model_args.model_name)
            setattr(base_model, 'config', config)
        elif model_args.model_backbone == COLPALI:
            base_model = ColPali.from_pretrained(model_args.model_name)
            setattr(base_model, 'config', config)
        else:
            # Loading external base model from HF
            config = AutoConfig.from_pretrained(model_args.model_name, trust_remote_code=True)
            config.use_cache = False
            base_model = cls.TRANSFORMER_CLS.from_pretrained(
                model_name_or_path, **kwargs, config=config,
                torch_dtype=torch.bfloat16,
                trust_remote_code=True)

        # Building the model on top of the base
        if model_args.lora:
            print_master(f'Loading LoRA from {model_name_or_path}')
            lora_config = LoraConfig.from_pretrained(model_name_or_path)
            lora_model = PeftModel.from_pretrained(base_model, model_name_or_path, config=lora_config, is_trainable=is_trainable)
            lora_model.load_adapter(model_name_or_path, lora_model.active_adapter, is_trainable=is_trainable)
            if not is_trainable:
                lora_model = lora_model.merge_and_unload()
            model = cls(
                encoder=lora_model,
                pooling=model_args.pooling,
                normalize=model_args.normalize,
                temperature=model_args.temperature
            )
        else:
            model = cls(
                encoder=base_model,
                pooling=model_args.pooling,
                normalize=model_args.normalize,
                temperature=model_args.temperature
            )

        model.model_backbone = model_args.model_backbone
        return model

    def save(self, output_dir: str):
        self.encoder.save_pretrained(output_dir)

    def forward(self, qry: Dict[str, Tensor] = None, tgt: Dict[str, Tensor] = None, *args, **kwargs):
        qry_reps = self.encode_input(qry) if qry else None  # (bsz_per_device, dim)
        tgt_reps = self.encode_input(tgt) if tgt else None # (bsz_per_device, dim)
        # qry_reps, tgt_reps = None, None
        # if qry is not None:
        #     qry_reps = self.encode_input(qry, compression_rate=0.2)  # QRY 压缩率
        # if tgt is not None:
        #     tgt_reps = self.encode_input(tgt, compression_rate=0.4)  # CAND 压缩率

        if qry_reps is None or tgt_reps is None:
            # return {"qry_reps": qry_reps, "tgt_reps": tgt_reps}
            out = {"qry_reps": qry_reps, "tgt_reps": tgt_reps}
            # NEW: 透传底模的图像 token mask（如果有）
            img_masks = getattr(self, "_last_image_token_bool_masks", None)
            if img_masks is not None:
                out["image_token_bool_masks"] = img_masks
            return out

        if self.is_ddp:
            all_qry_reps = self._dist_gather_tensor(qry_reps)
            all_tgt_reps = self._dist_gather_tensor(tgt_reps)
        else:
            all_qry_reps = qry_reps
            all_tgt_reps = tgt_reps

        scores = self.compute_similarity(all_qry_reps, all_tgt_reps)
        scores = scores.view(all_qry_reps.size(0), -1)
        target = torch.arange(scores.size(0), device=scores.device, dtype=torch.long)
        target = target * (all_qry_reps.size(0) // all_tgt_reps.size(0))
        loss = self.cross_entropy(scores / self.temperature, target)
        if self.is_ddp:
            loss = loss * self.world_size

        return loss

    def _dist_gather_tensor(self, t: Tensor):
        t = t.contiguous()
        all_tensors = [torch.empty_like(t) for _ in range(self.world_size)]
        dist.all_gather(all_tensors, t)
        all_tensors[self.process_rank] = t
        all_tensors = torch.cat(all_tensors, dim=0)
        return all_tensors

    def compute_similarity(self, q_reps, p_reps):
        return torch.matmul(q_reps, p_reps.transpose(0, 1))

class AsymMMEBModel(nn.Module):
    def __init__(self, encoder_qry: PreTrainedModel, encoder_tgt: PreTrainedModel,
                 pooling: str = 'last', normalize: bool = False, temperature: float = 0.02):
        super().__init__()
        self.encoder_qry = encoder_qry  # Qwen2.5-VL-GP
        self.encoder_tgt = encoder_tgt  # Qwen2.5-VL
        self.pooling = pooling
        self.normalize = normalize
        self.temperature = temperature
        self.cross_entropy = nn.CrossEntropyLoss(reduction='mean')
        self.is_ddp = dist.is_initialized()
        self.world_size = dist.get_world_size() if self.is_ddp else 1

        # 仅 query 端控制裁剪
        self.gp_do_selection_qry = False  # 训练默认 False，推理可 True

        # debug/观测用
        self._last_image_token_bool_masks = None
        self._last_image_token_mask_logits = None
        self._last_le_loss = None

    def _pool(self, last_hidden_state, attention_mask):
        left_padding = (attention_mask[:, -1].sum() == attention_mask.shape[0])
        B = last_hidden_state.shape[0]
        if left_padding:
            reps = last_hidden_state[torch.arange(B), -1, :]
        else:
            idx = attention_mask.sum(dim=1) - 1
            reps = last_hidden_state[torch.arange(B, device=last_hidden_state.device), idx]
        if self.normalize:
            reps = torch.nn.functional.normalize(reps, p=2, dim=-1)
        return reps

    def encode_qry(self, inp: Dict[str, torch.Tensor]):
        # 训练检索对比建议不开裁剪；推理打开
        do_sel = self.gp_do_selection_qry
        out = self.encoder_qry(
            **inp,
            return_dict=True,
            output_hidden_states=True,
            do_selection=do_sel,
            delay_selection=False if do_sel else True,
            use_cache=False,
            use_ref_masks=False,
        )
        self._last_image_token_bool_masks = getattr(out, "image_token_bool_masks", None)
        self._last_image_token_mask_logits = getattr(out, "image_token_mask_logits", None)
        self._last_le_loss = getattr(out, "le_loss", None)

        hs = out.hidden_states
        last = hs[-1] if isinstance(hs, (list, tuple)) else hs
        attn = getattr(out, "attention_mask", None) or inp["attention_mask"]
        rep = self._pool(last, attn)
        return rep

    def encode_tgt(self, inp: Dict[str, torch.Tensor]):
        out = self.encoder_tgt(
            **inp,
            return_dict=True,
            output_hidden_states=True,
            use_cache=False,
        )
        hs = out.hidden_states
        last = hs[-1] if isinstance(hs, (list, tuple)) else hs
        attn = getattr(out, "attention_mask", None) or inp["attention_mask"]
        rep = self._pool(last, attn)
        return rep

    def forward(self, qry: Dict[str, torch.Tensor] = None, tgt: Dict[str, torch.Tensor] = None):
        qry_reps = self.encode_qry(qry) if qry else None
        tgt_reps = self.encode_tgt(tgt) if tgt else None
        if qry_reps is None or tgt_reps is None:
            return {"qry_reps": qry_reps, "tgt_reps": tgt_reps}

        if self.is_ddp:
            qry_all = self._gather(qry_reps)
            tgt_all = self._gather(tgt_reps)
        else:
            qry_all, tgt_all = qry_reps, tgt_reps

        scores = torch.matmul(qry_all, tgt_all.transpose(0, 1)).view(qry_all.size(0), -1)
        target = torch.arange(scores.size(0), device=scores.device, dtype=torch.long)
        target = target * (qry_all.size(0) // tgt_all.size(0))
        loss = self.cross_entropy(scores / self.temperature, target)
        if self.is_ddp:
            loss = loss * self.world_size
        return loss

    def _gather(self, t: torch.Tensor):
        t = t.contiguous()
        outs = [torch.empty_like(t) for _ in range(self.world_size)]
        dist.all_gather(outs, t)
        outs[dist.get_rank()] = t
        return torch.cat(outs, dim=0)

    @classmethod
    def load_asym(cls, model_args, processor=None, processor_tgt=None, is_trainable=True):
        # query encoder: GP
        from src.model.vlm_backbone.qwen2_5_vl_gp.configuration import Qwen2_5_VL_GPConfig
        from src.model.vlm_backbone.qwen2_5_vl_gp.modeling_qwen2_5_vl import Qwen2_5_VLForConditionalGeneration_GP as GPModel
        from transformers import AutoConfig
        new_modules_dir = getattr(model_args, "new_modules_dir", None)
        assert new_modules_dir is not None, "new_modules_dir 不能为空（query 端 GP 模块）"

        gp_cfg = Qwen2_5_VL_GPConfig.from_pretrained(new_modules_dir)
        gp_cfg._attn_implementation = "flash_attention_2"
        gp_cfg.padding_side = "left"
        gp_cfg.use_cache = False
        encoder_qry = GPModel.from_pretrained(
            model_args.model_name,
            torch_dtype=torch.bfloat16,
            low_cpu_mem_usage=True,
            config=gp_cfg,
        )
        encoder_qry.load_new_modules(new_modules_dir)

        # target encoder: vanilla
        from src.model.vlm_backbone.qwen2_5_vl import Qwen2_5_VLForConditionalGeneration as VanillaModel
        cfg_tgt = AutoConfig.from_pretrained(model_args.model_name, trust_remote_code=True)
        cfg_tgt._attn_implementation = "flash_attention_2"
        cfg_tgt.use_cache = False
        # 兼容视觉
        if hasattr(cfg_tgt, "vision_config"):
            try:
                cfg_tgt.vision_config._attn_implementation = "flash_attention_2"
            except Exception:
                pass
        encoder_tgt = VanillaModel.from_pretrained(
            model_args.model_name,
            torch_dtype=torch.bfloat16,
            low_cpu_mem_usage=True,
            config=cfg_tgt,
        )

        # 冻结（如果只训检索头/LoRA）
        if getattr(model_args, "image_encoder_freeze", True):
            for p in encoder_qry.parameters(): p.requires_grad = False
            for p in encoder_tgt.parameters(): p.requires_grad = False
            # 只放开 query 端 GP 新模块（如果你要训练选择头）
            if hasattr(encoder_qry, "new_modules_to_be_saved"):
                for mod in encoder_qry.new_modules_to_be_saved().values():
                    if isinstance(mod, torch.nn.Parameter):
                        mod.requires_grad = True
                    else:
                        for p in mod.parameters(): p.requires_grad = True

        return cls(encoder_qry=encoder_qry, encoder_tgt=encoder_tgt,
                   pooling=model_args.pooling, normalize=model_args.normalize,
                   temperature=model_args.temperature)