# 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 # from torch.utils.data import DataLoader # from tqdm import tqdm # from transformers import HfArgumentParser, AutoConfig # from datasets import Dataset, concatenate_datasets # from datasets.distributed import split_dataset_by_node # from src.arguments import ModelArguments, DataArguments, TrainingArguments # from src.data.collator.eval_collator import MultimodalEvalDataCollator # from src.data.eval_dataset.base_eval_dataset import AutoEvalPairDataset, generate_cand_dataset # from src.eval_utils.metrics import RankingMetrics # from src.model.model import MMEBModel # from src.model.processor import get_backbone_name, load_processor, COLPALI # from src.utils import batch_to_device, print_rank, print_master # import multiprocessing # from multiprocessing import Pool, cpu_count # logging.basicConfig(level=logging.INFO, format='[%(asctime)s] %(levelname)s [%(name)s:%(lineno)s] %(message)s') # logger = logging.getLogger(__name__) # ############################################### # # 计时开始 # def start_timer(name, timing_dict): # timing_dict[name] = time.time() # # 计时结束 # def end_timer(name, timing_dict): # end_time = time.time() # if name in timing_dict: # timing_dict[name] = end_time - timing_dict[name] # # # 放在 main 函数之前，或者单独放在 utils.py 中 # # def register_hooks(model, timing_dict): # # # --- vision_encoder hook --- # # def vision_forward_hook(module, input, output): # # print_master(f"[vision_encoder] output shape: {output.shape}") # # print_master(f"[vision_encoder] num_image_tokens: {output.shape[1]}") # # start_timer("vision_encoder", timing_dict) # # def vision_forward_post_hook(module, input, output): # # end_timer("vision_encoder", timing_dict) # # model.encoder.visual.register_forward_hook(vision_forward_hook) # # model.encoder.visual.register_forward_hook(vision_forward_post_hook) # # # --- merger hook --- # # def merger_forward_hook(module, input, output): # # print_master(f"[merger] before merger - input shape: {input[0].shape}") # # print_master(f"[merger] before merger - num_tokens: {input[0].shape[1]}") # # start_timer("merger", timing_dict) # # def merger_forward_post_hook(module, input, output): # # print_master(f"[merger] after merger - output shape: {output.shape}") # # print_master(f"[merger] after merger - num_tokens: {output.shape[1]}") # # end_timer("merger", timing_dict) # # if hasattr(model.encoder.visual, 'merger'): # # model.encoder.visual.merger.register_forward_hook(merger_forward_hook) # # model.encoder.visual.merger.register_forward_hook(merger_forward_post_hook) # # # --- decoder hook --- # # def decoder_forward_hook(module, input, output): # # # 这里更新为接收 input 和 output 参数 # # if isinstance(input, tuple) and len(input) > 0: # # print_master(f"[llm_decoder] input shape: {input[0].shape}") # # print_master(f"[llm_decoder] total_tokens (image+text): {input[0].shape[1]}") # # start_timer("llm_decoder", timing_dict) # # def decoder_forward_post_hook(module, input, output): # # end_timer("llm_decoder", timing_dict) # # model.encoder.model.register_forward_hook(decoder_forward_hook) # # model.encoder.model.register_forward_hook(decoder_forward_post_hook) # # # --- lm_head hook --- # # def lm_head_forward_hook(module, input, output): # # start_timer("lm_head", timing_dict) # # def lm_head_forward_post_hook(module, input, output): # # end_timer("lm_head", timing_dict) # # model.encoder.lm_head.register_forward_hook(lm_head_forward_hook) # # model.encoder.lm_head.register_forward_hook(lm_head_forward_post_hook) # def register_timing_hooks(model, timing_dict): # def make_hooks(name): # def pre_hook(module, input): # timing_dict[f"{name}_start"] = time.time() # def forward_hook(module, input, output): # elapsed = time.time() - timing_dict[f"{name}_start"] # timing_dict[name] = elapsed # 记录时间 # print_master(f"[{name}] took {elapsed * 1000:.2f} ms") # return pre_hook, forward_hook # # vision encoder # pre, post = make_hooks("vision_encoder") # model.encoder.visual.register_forward_pre_hook(pre) # model.encoder.visual.register_forward_hook(post) # # merger # if hasattr(model.encoder.visual, 'merger'): # pre, post = make_hooks("merger") # model.encoder.visual.merger.register_forward_pre_hook(pre) # model.encoder.visual.merger.register_forward_hook(post) # # decoder # pre, post = make_hooks("llm_decoder") # model.encoder.model.register_forward_pre_hook(pre) # model.encoder.model.register_forward_hook(post) # # lm_head # pre, post = make_hooks("lm_head") # model.encoder.lm_head.register_forward_pre_hook(pre) # model.encoder.lm_head.register_forward_hook(post) # ##################################################### # 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", # timing_dict: dict | None = None # ) -> tuple[np.ndarray, list]: # """ # Encodes embeddings for a given dataset using the model, handling both standard and # late-interaction models in a DDP-safe manner. # """ # 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 # model.eval() # with torch.no_grad(): # for inputs, dataset_info in tqdm(loader, desc=f"{description} (rank {local_rank})", disable=local_rank > 0): # inputs = batch_to_device(inputs, training_args.device) # with torch.autocast(enabled=True, dtype=torch.bfloat16, device_type="cuda"): # # Determine if encoding query or target based on available keys # if encode_side == "qry": # output = model(qry=inputs) # reps = output["qry_reps"].detach() # local_gt_infos.extend(dataset_info) # to retain all information per query # else: # output = model(tgt=inputs) # reps = output["tgt_reps"].detach() # local_gt_infos.extend([info["cand_name"] for info in dataset_info]) # to retain ground-truth labels # 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([]), [] # # === DDP Synchronization and Padding for Late-Interaction Models === # if is_late_interaction: # if dist.is_initialized(): # # 1. Find the global maximum sequence length across all ranks # 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 all local embeddings to the 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: # Should not happen if model is consistently late-interaction # padded_embeds.append(reps_batch) # embeds_tensor = torch.cat(padded_embeds, dim=0).contiguous() # else: # Standard dense models # 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...") # # Use the more efficient all_gather_into_tensor for tensors # 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() # # Gather metadata, for which all_gather_object is appropriate # 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] # else: # all_gt_infos = local_gt_infos # final_embeddings = embeds_tensor.cpu().float().numpy() # print_master(f"Timing results for {description}:") # for k, v in timing_dict.items(): # if not k.startswith('_'): # print_master(f" {k}: {v:.4f} sec") # return final_embeddings, all_gt_infos # def main(): # 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 # # DEBUG PRINTS for Distributed Setup # 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()}") # 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) # 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) # # --- 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}') # # --- DDP-Safe Model Loading --- # # Step 1: Only the master process (rank 0) downloads the model. # 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}...") # # Step 2: All processes wait here. The non-master processes will pause # # until the master process (rank 0) finishes downloading and exits this barrier. # if torch.distributed.is_initialized(): # torch.distributed.barrier() # # Step 3: Now that the model is cached, the non-master processes load it from the local cache. # 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) # with open(data_args.dataset_config, 'r') as yaml_file: # dataset_configs = yaml.safe_load(yaml_file) # ############################################################################# # import time # timing_dict = {} # register_hooks(model, timing_dict) # 注册 hooks，开始计时 # ############################################################################## # # --- Main Evaluation Loop --- # for dataset_idx, (dataset_name, task_config) in enumerate(dataset_configs.items()): # # 0. load dataset # if dist.is_initialized(): # dist.barrier() # print_master(f"--- Evaluating {dataset_name} ---") # query_embed_path = os.path.join(data_args.encode_output_path, f"{dataset_name}_qry") # cand_embed_path = os.path.join(data_args.encode_output_path, f"{dataset_name}_tgt") # dataset_info_path = os.path.join(data_args.encode_output_path, f"{dataset_name}_info.jsonl") # do_query = not os.path.exists(query_embed_path) or not os.path.exists(dataset_info_path) # do_cand = not os.path.exists(cand_embed_path) # if do_query or do_cand: # if data_args.data_basedir is not None: # # Construct full paths for data files if --data_basedir is provided # 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 # # Pad datasets to be divisible by world_size before splitting # if dist.is_initialized(): # 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 # # --- 1. Compute Query Embeddings --- # if do_query: # print_master("Encoding queries...") # eval_qry_collator = MultimodalEvalDataCollator(processor, model_args, data_args, "qry") # eval_qry_loader = DataLoader(eval_qry_dataset, batch_size=training_args.per_device_eval_batch_size, collate_fn=eval_qry_collator, num_workers=training_args.dataloader_num_workers) # query_embeds, gt_infos = encode_embeddings(model, eval_qry_loader, training_args, model_args, padded_qry_dataset, encode_side="qry", description=f"Queries for {dataset_name}", timing_dict=timing_dict) # query_embeds = query_embeds[:len(full_eval_qry_dataset)] # world_size>1, trim the padded data points # gt_infos = gt_infos[:len(full_eval_qry_dataset)] # if local_rank == 0: # with open(query_embed_path, 'wb') as f: # pickle.dump(query_embeds, f) # with open(dataset_info_path, 'w') as f: # for info in gt_infos: # f.write(json.dumps(info) + '\n') # print_master(f"Saved query embeddings to {query_embed_path}") # if dist.is_initialized(): # dist.barrier() # # --- 2. Compute Candidate Embeddings --- # if do_cand: # print_master("Encoding candidates...") # eval_cand_collator = MultimodalEvalDataCollator(processor, model_args, data_args, "cand") # eval_cand_loader = DataLoader(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_embeds, all_cand_ids = encode_embeddings(model, eval_cand_loader, training_args, model_args, padded_cand_dataset, encode_side="cand", description=f"Candidates for {dataset_name}", timing_dict=timing_dict) # cand_embeds = cand_embeds[:len(full_eval_cand_dataset)] # world_size>1, trim the padded data points # all_cand_ids = all_cand_ids[:len(full_eval_cand_dataset)] # if local_rank == 0: # cand_embed_dict = {cand_id: embed for cand_id, embed in zip(all_cand_ids, cand_embeds)} # with open(cand_embed_path, 'wb') as f: pickle.dump(cand_embed_dict, f) # print_master(f"Saved candidate embeddings to {cand_embed_path}") # if dist.is_initialized(): # dist.barrier() # # --- 3. Compute Scores (on master rank only) --- # if local_rank == 0: # score_path = os.path.join(data_args.encode_output_path, f"{dataset_name}_score.json") # if os.path.exists(score_path): # try: # with open(score_path, "r") as f: # score_dict = json.load(f) # print_master(f"Score of {dataset_name} (loaded from previous run): {score_path}") # formatted = {k: f"{v:.4f}" for k, v in score_dict.items()} # print_master(formatted) # continue # except Exception as e: # print_master(f"Failed to load score for {dataset_name}, skipping {dataset_name}") # with open(query_embed_path, 'rb') as f: qry_embeds = pickle.load(f) # with open(cand_embed_path, 'rb') as f: cand_embed_dict = pickle.load(f) # gt_infos = [json.loads(l) for l in open(dataset_info_path)] # pred_dicts = [] # rank_against_all_candidates = task_config.get("eval_type", "global") == "global" # if rank_against_all_candidates: # cand_keys = list(cand_embed_dict.keys()) # cand_embeds = np.stack([cand_embed_dict[key] for key in cand_keys]) # # Handle late-interaction scoring # if qry_embeds.ndim == 3: # Query: [N_q, L_q, H] | Candidate: [N_c, L_c, H] # qry_embed = torch.from_numpy(qry_embeds) # cand_embeds = [torch.from_numpy(np.array(t)) for t in cand_embeds] # scores = processor.score(qry_embed, cand_embeds, batch_size=64) # use ColPali score function # ranked_candids = torch.argsort(-scores, dim=1).cpu().numpy().tolist() # else: # Dense # cosine_scores = np.dot(qry_embeds, cand_embeds.T) # ranked_candids = np.argsort(-cosine_scores, axis=1) # for qid, (ranked_candid, gt_info) in tqdm(enumerate(zip(ranked_candids, gt_infos)), desc=f"Calculating scores for {dataset_name}"): # rel_docids = gt_info["label_name"] if isinstance(gt_info["label_name"], list) else [gt_info["label_name"]] # rel_scores = gt_info["rel_scores"] if "rel_scores" in gt_info else None # assert rel_scores is None or len(rel_docids) == len(rel_scores) # pred_dicts.append({ # "prediction": [cand_keys[i] for i in ranked_candid], # "label": rel_docids, # "rel_scores": rel_scores, # }) # else: # for qid, (qry_embed, gt_info) in tqdm(enumerate(zip(qry_embeds, gt_infos)), desc=f"Calculating scores for {dataset_name}"): # cand_embeds = np.stack([cand_embed_dict[key] for key in gt_info["cand_names"]]) # if qry_embeds.ndim == 3: # Query: [N_q, L_q, H] | Candidate: [N_c, L_c, H] # qry_embed = torch.from_numpy(np.array(qry_embed)).unsqueeze(0) # cand_embeds = [torch.from_numpy(np.array(t)) for t in cand_embeds] # scores = processor.score(qry_embed, cand_embeds, batch_size=1024) # use ColPali score function # ranked_candids = torch.argsort(-scores, dim=1).cpu().numpy().tolist()[0] # else: # cosine_score = np.dot(qry_embed, cand_embeds.T) # ranked_candids = np.argsort(-cosine_score) # rel_docids = gt_info["label_name"] if isinstance(gt_info["label_name"], list) else [gt_info["label_name"]] # rel_scores = gt_info["rel_scores"] if "rel_scores" in gt_info else None # assert rel_scores is None or len(rel_docids) == len(rel_scores) # pred_dicts.append({ # "prediction": [gt_info["cand_names"][i] for i in ranked_candids], # "label": rel_docids, # "rel_scores": rel_scores, # }) # score_path = os.path.join(data_args.encode_output_path, f"{dataset_name}_score.json") # pred_path = os.path.join(data_args.encode_output_path, f"{dataset_name}_pred.jsonl") # metrics_to_report = task_config["metrics"] if task_config.get("metrics", None) is not None else ["hit", "ndcg", "precision", "recall", "f1", "map", "mrr"] # metrics = RankingMetrics(metrics_to_report) # score_dict = metrics.evaluate(pred_dicts) # formatted = {k: f"{v:.4f}" for k, v in score_dict.items()} # score_dict["num_pred"] = len(pred_dicts) # score_dict["num_data"] = len(gt_infos) # print_master(f"Score of {dataset_name}:") # print_master(formatted) # print_master(f"Outputting final score to: {score_path}") # with open(score_path, "w") as f: # json.dump(score_dict, f, indent=4) # with open(pred_path, "w") as f: # for pred in pred_dicts: # f.write(json.dumps(pred) + '\n') # if __name__ == "__main__": # main() ################################################################################################### #直接打印输出对应模块时间和token数量 # 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 # from torch.utils.data import DataLoader # from tqdm import tqdm # from transformers import HfArgumentParser, AutoConfig, AutoTokenizer#, Qwen2VLForConditionalGeneration # from datasets import Dataset, concatenate_datasets # from datasets.distributed import split_dataset_by_node # from src.model.vlm_backbone.qwen2_vl.modeling_qwen2_vl 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 import MMEBModel # from src.model.processor import get_backbone_name, load_processor, COLPALI # from src.utils import batch_to_device, print_rank, print_master # import multiprocessing # from multiprocessing import Pool, cpu_count # 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] # elif hasattr(output, 'last_hidden_state'): # token_info["vision_tokens"] = output.last_hidden_state.shape[1] # # --- Hook Functions Definition --- # # (timing_pre_hook and timing_post_hook remain as previously corrected with debug prints) # def register_model_hooks(model): # registered_modules = [] # core_model = model # print_master(f"DEBUG: Initial model type in register_model_hooks: {type(model)}") # if hasattr(model, 'encoder') and model.encoder is not None: # print_master(f"DEBUG: model has 'encoder' attribute. Type of model.encoder: {type(model.encoder)}") # # 使用从 'src' 路径导入的 Qwen2VLForConditionalGeneration 进行检查 # if isinstance(model.encoder, _Qwen2VLForConditionalGeneration_src): # print_master("Detected MMEBModel structure, registering hooks on model.encoder's sub-modules.") # core_model = model.encoder # else: # print_master(f"WARNING: model.encoder is not an instance of _Qwen2VLForConditionalGeneration_src. Its type is {type(model.encoder)}. Hooks will be registered on top-level model if applicable.") # else: # print_master("WARNING: Model structure does not have an 'encoder' attribute. Registering hooks directly on top-level modules.") # # 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]: # """ # Encodes embeddings for a given dataset using the model, handling both standard and # late-interaction models in a DDP-safe manner. # """ # 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 # model.eval() # # Register hooks for the model once per encode_embeddings call # # This assumes `model` is the MMEBModel instance that wraps the actual HuggingFace model # # You might need to adjust this if MMEBModel internally manages multiple sub-models # registered_hooks = register_model_hooks(model) # <--- FIX: Assign the return value here # # Initialize a tokenizer for text token counting (needs to be from the same model path) # temp_tokenizer = AutoTokenizer.from_pretrained(model_args.model_name) # 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 # Encoding doesn't generate text output tokens # token_info["total_llm_input_tokens"] = 0 # inputs = batch_to_device(inputs, training_args.device) # with torch.autocast(enabled=True, dtype=torch.bfloat16, device_type="cuda"): # # Determine if encoding query or target based on available keys # # This is where the forward pass happens, triggering hooks # start_inference_time = time.time() # if encode_side == "qry": # output = model(qry=inputs) # reps = output["qry_reps"].detach() # local_gt_infos.extend(dataset_info) # to retain all information per query # else: # output = model(tgt=inputs) # reps = output["tgt_reps"].detach() # local_gt_infos.extend([info["cand_name"] for info in dataset_info]) # to retain ground-truth labels # end_inference_time = time.time() # # --- Update total LLM input tokens after the model call --- # # This requires knowing which part of `inputs` corresponds to the LLM's full input. # # Assuming `inputs.input_ids` directly goes into the LLM part of Qwen2-VL. # if 'input_ids' in inputs and inputs['input_ids'] is not None: # token_info["total_llm_input_tokens"] = inputs['input_ids'].shape[1] # # Approximation for text_input_tokens (if not explicitly available from collator) # # This assumes visual tokens are a prefix and the rest are text/special tokens. # token_info["text_input_tokens"] = token_info["total_llm_input_tokens"] - token_info["vision_tokens"] # # Ensure it's not negative # token_info["text_input_tokens"] = max(0, token_info["text_input_tokens"]) # # --- Print Inference Timing and Token Statistics per Batch --- # print_rank(f"\n--- Inference Statistics for {encode_side} batch (Rank {local_rank}) ---") # print_rank(f"Batch Inference took: {end_inference_time - start_inference_time:.4f} seconds") # # Calculate and print module timings # print_rank("--- Module Inference Timing Statistics ---") # 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 = {} # Store start times for each pre-hook # 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: # print_rank(f"**{module_name}**: Total: {sum(durations):.6f}s, Count: {len(durations)}, Avg: {sum(durations)/len(durations):.6f}s") # else: # print_rank(f"**{module_name}**: No complete timing data found for this batch.") # print_rank("--- Token Count Statistics ---") # print_rank(f"**视觉 token 数量**: {token_info['vision_tokens']}") # print_rank(f"**语言输入 token 数量 (仅原始文本)**: {token_info['text_input_tokens']}") # print_rank(f"**LLM总输入 token 数量 (包含视觉 + 格式化文本)**: {token_info['total_llm_input_tokens']}") # print_rank(f"**语言输出 token 数量**: {token_info['text_output_tokens']}") # Will be 0 for encoding # 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([]), [] # # === DDP Synchronization and Padding for Late-Interaction Models === # if is_late_interaction: # if dist.is_initialized(): # # 1. Find the global maximum sequence length across all ranks # 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 all local embeddings to the 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: # Should not happen if model is consistently late-interaction # padded_embeds.append(reps_batch) # embeds_tensor = torch.cat(padded_embeds, dim=0).contiguous() # else: # Standard dense models # 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...") # # Use the more efficient all_gather_into_tensor for tensors # 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() # # Gather metadata, for which all_gather_object is appropriate # 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] # else: # all_gt_infos = local_gt_infos # final_embeddings = embeds_tensor.cpu().float().numpy() # return final_embeddings, all_gt_infos # def main(): # 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 # # DEBUG PRINTS for Distributed Setup # 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()}") # 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) # 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) # # --- 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}') # # --- DDP-Safe Model Loading --- # # Step 1: Only the master process (rank 0) downloads the model. # 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}...") # # Step 2: All processes wait here. The non-master processes will pause # # until the master process (rank 0) finishes downloading and exits this barrier. # if torch.distributed.is_initialized(): # torch.distributed.barrier() # # Step 3: Now that the model is cached, the non-master processes load it from the local cache. # 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) # 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()): # # 0. load dataset # if dist.is_initialized(): # dist.barrier() # print_master(f"--- Evaluating {dataset_name} ---") # query_embed_path = os.path.join(data_args.encode_output_path, f"{dataset_name}_qry") # cand_embed_path = os.path.join(data_args.encode_output_path, f"{dataset_name}_tgt") # dataset_info_path = os.path.join(data_args.encode_output_path, f"{dataset_name}_info.jsonl") # do_query = not os.path.exists(query_embed_path) or not os.path.exists(dataset_info_path) # do_cand = not os.path.exists(cand_embed_path) # if do_query or do_cand: # if data_args.data_basedir is not None: # # Construct full paths for data files if --data_basedir is provided # 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 # # Pad datasets to be divisible by world_size before splitting # if dist.is_initialized(): # 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 # # --- 1. Compute Query Embeddings --- # if do_query: # print_master("Encoding queries...") # eval_qry_collator = MultimodalEvalDataCollator(processor, model_args, data_args, "qry") # eval_qry_loader = DataLoader(eval_qry_dataset, batch_size=training_args.per_device_eval_batch_size, collate_fn=eval_qry_collator, num_workers=training_args.dataloader_num_workers) # query_embeds, gt_infos = encode_embeddings(model, eval_qry_loader, training_args, model_args, padded_qry_dataset, encode_side="qry", description=f"Queries for {dataset_name}") # query_embeds = query_embeds[:len(full_eval_qry_dataset)] # world_size>1, trim the padded data points # gt_infos = gt_infos[:len(full_eval_qry_dataset)] # if local_rank == 0: # with open(query_embed_path, 'wb') as f: # pickle.dump(query_embeds, f) # with open(dataset_info_path, 'w') as f: # for info in gt_infos: # f.write(json.dumps(info) + '\n') # print_master(f"Saved query embeddings to {query_embed_path}") # if dist.is_initialized(): # dist.barrier() # # --- 2. Compute Candidate Embeddings --- # if do_cand: # print_master("Encoding candidates...") # eval_cand_collator = MultimodalEvalDataCollator(processor, model_args, data_args, "cand") # eval_cand_loader = DataLoader(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_embeds, all_cand_ids = encode_embeddings(model, eval_cand_loader, training_args, model_args, padded_cand_dataset, encode_side="cand", description=f"Candidates for {dataset_name}") # cand_embeds = cand_embeds[:len(full_eval_cand_dataset)] # world_size>1, trim the padded data points # all_cand_ids = all_cand_ids[:len(full_eval_cand_dataset)] # if local_rank == 0: # cand_embed_dict = {cand_id: embed for cand_id, embed in zip(all_cand_ids, cand_embeds)} # with open(cand_embed_path, 'wb') as f: pickle.dump(cand_embed_dict, f) # print_master(f"Saved candidate embeddings to {cand_embed_path}") # if dist.is_initialized(): # dist.barrier() # # --- 3. Compute Scores (on master rank only) --- # if local_rank == 0: # score_path = os.path.join(data_args.encode_output_path, f"{dataset_name}_score.json") # if os.path.exists(score_path): # try: # with open(score_path, "r") as f: # score_dict = json.load(f) # print_master(f"Score of {dataset_name} (loaded from previous run): {score_path}") # formatted = {k: f"{v:.4f}" for k, v in score_dict.items()} # print_master(formatted) # continue # except Exception as e: # print_master(f"Failed to load score for {dataset_name}, skipping {dataset_name}") # with open(query_embed_path, 'rb') as f: qry_embeds = pickle.load(f) # with open(cand_embed_path, 'rb') as f: cand_embed_dict = pickle.load(f) # gt_infos = [json.loads(l) for l in open(dataset_info_path)] # pred_dicts = [] # rank_against_all_candidates = task_config.get("eval_type", "global") == "global" # if rank_against_all_candidates: # cand_keys = list(cand_embed_dict.keys()) # cand_embeds = np.stack([cand_embed_dict[key] for key in cand_keys]) # # Handle late-interaction scoring # if qry_embeds.ndim == 3: # Query: [N_q, L_q, H] | Candidate: [N_c, L_c, H] # qry_embed = torch.from_numpy(qry_embeds) # cand_embeds = [torch.from_numpy(np.array(t)) for t in cand_embeds] # scores = processor.score(qry_embed, cand_embeds, batch_size=64) # use ColPali score function # ranked_candids = torch.argsort(-scores, dim=1).cpu().numpy().tolist() # else: # Dense # cosine_scores = np.dot(qry_embeds, cand_embeds.T) # ranked_candids = np.argsort(-cosine_scores, axis=1) # for qid, (ranked_candid, gt_info) in tqdm(enumerate(zip(ranked_candids, gt_infos)), desc=f"Calculating scores for {dataset_name}"): # rel_docids = gt_info["label_name"] if isinstance(gt_info["label_name"], list) else [gt_info["label_name"]] # rel_scores = gt_info["rel_scores"] if "rel_scores" in gt_info else None # assert rel_scores is None or len(rel_docids) == len(rel_scores) # pred_dicts.append({ # "prediction": [cand_keys[i] for i in ranked_candid], # "label": rel_docids, # "rel_scores": rel_scores, # }) # else: # for qid, (qry_embed, gt_info) in tqdm(enumerate(zip(qry_embeds, gt_infos)), desc=f"Calculating scores for {dataset_name}"): # cand_embeds = np.stack([cand_embed_dict[key] for key in gt_info["cand_names"]]) # if qry_embeds.ndim == 3: # Query: [N_q, L_q, H] | Candidate: [N_c, L_c, H] # qry_embed = torch.from_numpy(np.array(qry_embed)).unsqueeze(0) # cand_embeds = [torch.from_numpy(np.array(t)) for t in cand_embeds] # scores = processor.score(qry_embed, cand_embeds, batch_size=1024) # use ColPali score function # ranked_candids = torch.argsort(-scores, dim=1).cpu().numpy().tolist()[0] # else: # cosine_score = np.dot(qry_embed, cand_embeds.T) # ranked_candids = np.argsort(-cosine_score) # rel_docids = gt_info["label_name"] if isinstance(gt_info["label_name"], list) else [gt_info["label_name"]] # rel_scores = gt_info["rel_scores"] if "rel_scores" in gt_info else None # assert rel_scores is None or len(rel_docids) == len(rel_scores) # pred_dicts.append({ # "prediction": [gt_info["cand_names"][i] for i in ranked_candids], # "label": rel_docids, # "rel_scores": rel_scores, # }) # score_path = os.path.join(data_args.encode_output_path, f"{dataset_name}_score.json") # pred_path = os.path.join(data_args.encode_output_path, f"{dataset_name}_pred.jsonl") # metrics_to_report = task_config["metrics"] if task_config.get("metrics", None) is not None else ["hit", "ndcg", "precision", "recall", "f1", "map", "mrr"] # metrics = RankingMetrics(metrics_to_report) # score_dict = metrics.evaluate(pred_dicts) # formatted = {k: f"{v:.4f}" for k, v in score_dict.items()} # score_dict["num_pred"] = len(pred_dicts) # score_dict["num_data"] = len(gt_infos) # print_master(f"Score of {dataset_name}:") # print_master(formatted) # print_master(f"Outputting final score to: {score_path}") # with open(score_path, "w") as f: # json.dump(score_dict, f, indent=4) # with open(pred_path, "w") as f: # for pred in pred_dicts: # f.write(json.dumps(pred) + '\n') # 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 # 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 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 import MMEBModel # from src.model.processor import get_backbone_name, load_processor, COLPALI # from src.utils import batch_to_device, print_rank, print_master # logging.basicConfig(level=logging.INFO, format='[%(asctime)s] %(levelname)s [%(name)s:%(lineno)s] %(message)s') # logger = logging.getLogger(__name__) # # --- 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 # # print_master(f"DEBUG: Initial model type in register_model_hooks: {type(model)}") # if hasattr(model, 'encoder') and model.encoder is not None: # # print_master(f"DEBUG: model has 'encoder' attribute. Type of model.encoder: {type(model.encoder)}") # # 使用从 'src' 路径导入的 Qwen2VLForConditionalGeneration 进行检查 # if isinstance(model.encoder, _Qwen2VLForConditionalGeneration_src): # # print_master("Detected MMEBModel structure, registering hooks on model.encoder's sub-modules.") # core_model = model.encoder # else: # print_master(f"WARNING: model.encoder is not an instance of _Qwen2VLForConditionalGeneration_src. Its type is {type(model.encoder)}. Hooks will be registered on top-level model if applicable.") # else: # print_master("WARNING: Model structure does not have an 'encoder' attribute. Registering hooks directly on top-level modules.") # # 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]: # Added list to return type for batch_stats # """ # Encodes embeddings for a given dataset using the model, handling both standard and # late-interaction models in a DDP-safe manner. # """ # 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 = [] # model.eval() # # Register hooks for the model once per encode_embeddings call # registered_hooks = register_model_hooks(model) # with torch.no_grad(): # for inputs, dataset_info in tqdm(loader, desc=f"{description} (rank {local_rank})", disable=local_rank > 0): # # --- Reset 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 else 1 # Determine actual batch size # with torch.autocast(enabled=True, dtype=torch.bfloat16, device_type="cuda"): # start_inference_time = time.time() # 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]) # end_inference_time = time.time() # # --- Update total LLM input tokens after the model call --- # if 'input_ids' in inputs and inputs['input_ids'] is not None: # # `inputs['input_ids'].shape[1]` gives the sequence length, # # which is the number of tokens per item in the batch. # # To get total tokens for the batch, multiply by batch size. # token_info["total_llm_input_tokens"] = inputs['input_ids'].shape[1] # # Approximation for text_input_tokens # 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"]) # Ensure not negative # # --- 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"], # } # } # # 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) # Append the stats for this batch # # --- Print Inference Timing and Token Statistics per Batch (Optional, for debugging) --- # 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([]), [], [] # Return empty list for batch_stats_list as well # # === DDP Synchronization and Padding for Late-Interaction Models === # if is_late_interaction: # if dist.is_initialized(): # # 1. Find the global maximum sequence length across all ranks # 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 all local embeddings to the 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: # Should not happen if model is consistently late-interaction # padded_embeds.append(reps_batch) # embeds_tensor = torch.cat(padded_embeds, dim=0).contiguous() # else: # Standard dense models # 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...") # # Use the more efficient all_gather_into_tensor for tensors # 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() # # Gather metadata, for which all_gather_object is appropriate # 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] # # --- New: Gather batch_stats_list from all ranks --- # 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] # else: # all_gt_infos = local_gt_infos # final_embeddings = embeds_tensor.cpu().float().numpy() # all_batch_stats = batch_stats_list # If not DDP, just use local list # return final_embeddings, all_gt_infos, all_batch_stats # def main(): # 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 # # DEBUG PRINTS for Distributed Setup # 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()}") # 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) # 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) # # --- 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}') # # --- DDP-Safe Model Loading --- # # Step 1: Only the master process (rank 0) downloads the model. # 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}...") # # Step 2: All processes wait here. The non-master processes will pause # # until the master process (rank 0) finishes downloading and exits this barrier. # if torch.distributed.is_initialized(): # torch.distributed.barrier() # # Step 3: Now that the model is cached, the non-master processes load it from the local cache. # 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) # 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()): # # Initialize task-level statistics accumulators # task_total_stats = { # "total_inference_time_seconds": 0.0, # "module_inference_times": { # "Qwen2VisionTransformerPretrainedModel": {"total": 0.0, "count": 0}, # "PatchMerger": {"total": 0.0, "count": 0}, # "Qwen2VLModel": {"total": 0.0, "count": 0}, # "Linear": {"total": 0.0, "count": 0}, # }, # "token_counts": { # "visual_tokens": 0, # "language_input_tokens_raw": 0, # "llm_total_input_tokens": 0, # "language_output_tokens": 0, # }, # "data_point_count": 0 # Number of image-text pairs processed # } # if dist.is_initialized(): # dist.barrier() # print_master(f"\n--- Evaluating {dataset_name} ---") # query_embed_path = os.path.join(data_args.encode_output_path, f"{dataset_name}_qry") # cand_embed_path = os.path.join(data_args.encode_output_path, f"{dataset_name}_tgt") # dataset_info_path = os.path.join(data_args.encode_output_path, f"{dataset_name}_info.jsonl") # # New: Define path for inference statistics output # inference_stats_path = os.path.join(data_args.encode_output_path, f"{dataset_name}_inference_stats.json") # do_query = not os.path.exists(query_embed_path) or not os.path.exists(dataset_info_path) # do_cand = not os.path.exists(cand_embed_path) # if do_query or do_cand: # if data_args.data_basedir is not None: # # Construct full paths for data files if --data_basedir is provided # 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 # # Pad datasets to be divisible by world_size before splitting # if dist.is_initialized(): # 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 # # --- 1. Compute Query Embeddings --- # if do_query: # print_master("Encoding queries...") # eval_qry_collator = MultimodalEvalDataCollator(processor, model_args, data_args, "qry") # eval_qry_loader = DataLoader(eval_qry_dataset, batch_size=training_args.per_device_eval_batch_size, collate_fn=eval_qry_collator, num_workers=training_args.dataloader_num_workers) # # Modified: capture batch_stats_list # query_embeds, gt_infos, qry_batch_stats = encode_embeddings(model, eval_qry_loader, training_args, model_args, padded_qry_dataset, encode_side="qry", description=f"Queries for {dataset_name}") # # Accumulate query statistics # for batch_stat in qry_batch_stats: # batch_size = batch_stat["batch_size"] # task_total_stats["total_inference_time_seconds"] += batch_stat["total_inference_time_seconds"] # for module_name, module_stats in batch_stat["module_inference_times"].items(): # if module_name in task_total_stats["module_inference_times"]: # task_total_stats["module_inference_times"][module_name]["total"] += module_stats["total"] # task_total_stats["module_inference_times"][module_name]["count"] += module_stats["count"] # count here is per-module-call, not per-item # # Token counts are per-item for 'visual_tokens', 'llm_total_input_tokens'. # # For 'text_input_tokens', it's calculated based on sequence length, so it's also total tokens in the batch. # # We need to average it later by total data_point_count. # # However, your current hook logic collects the token count for a *single* item if batch_size=1, # # or for the full batch if it processes sequentially. # # Let's assume the `token_info` collected by hooks reflects the *current batch*. # # To get per-item average later, we sum up totals and divide by total data points. # task_total_stats["token_counts"]["visual_tokens"] += batch_stat["token_counts"]["visual_tokens"] * batch_size # Corrected assumption: visual_tokens are per-item, multiplied by batch_size to get total for batch # task_total_stats["token_counts"]["language_input_tokens_raw"] += batch_stat["token_counts"]["language_input_tokens_raw"] * batch_size # Corrected # task_total_stats["token_counts"]["llm_total_input_tokens"] += batch_stat["token_counts"]["llm_total_input_tokens"] * batch_size # Corrected # task_total_stats["token_counts"]["language_output_tokens"] += batch_stat["token_counts"]["language_output_tokens"] * batch_size # Corrected # task_total_stats["data_point_count"] += batch_size # Accumulate the number of processed items # query_embeds = query_embeds[:len(full_eval_qry_dataset)] # gt_infos = gt_infos[:len(full_eval_qry_dataset)] # if local_rank == 0: # with open(query_embed_path, 'wb') as f: # pickle.dump(query_embeds, f) # with open(dataset_info_path, 'w') as f: # for info in gt_infos: # f.write(json.dumps(info) + '\n') # print_master(f"Saved query embeddings to {query_embed_path}") # if dist.is_initialized(): # dist.barrier() # # --- 2. Compute Candidate Embeddings --- # if do_cand: # print_master("Encoding candidates...") # eval_cand_collator = MultimodalEvalDataCollator(processor, model_args, data_args, "cand") # eval_cand_loader = DataLoader(eval_cand_dataset, batch_size=training_args.per_device_eval_batch_size, collate_fn=eval_cand_collator, num_workers=training_args.dataloader_num_workers) # # Modified: capture batch_stats_list # cand_embeds, all_cand_ids, cand_batch_stats = encode_embeddings(model, eval_cand_loader, training_args, model_args, padded_cand_dataset, encode_side="cand", description=f"Candidates for {dataset_name}") # # Accumulate candidate statistics (similar logic as query) # for batch_stat in cand_batch_stats: # batch_size = batch_stat["batch_size"] # task_total_stats["total_inference_time_seconds"] += batch_stat["total_inference_time_seconds"] # for module_name, module_stats in batch_stat["module_inference_times"].items(): # if module_name in task_total_stats["module_inference_times"]: # task_total_stats["module_inference_times"][module_name]["total"] += module_stats["total"] # task_total_stats["module_inference_times"][module_name]["count"] += module_stats["count"] # task_total_stats["token_counts"]["visual_tokens"] += batch_stat["token_counts"]["visual_tokens"] * batch_size # task_total_stats["token_counts"]["language_input_tokens_raw"] += batch_stat["token_counts"]["language_input_tokens_raw"] * batch_size # task_total_stats["token_counts"]["llm_total_input_tokens"] += batch_stat["token_counts"]["llm_total_input_tokens"] * batch_size # task_total_stats["token_counts"]["language_output_tokens"] += batch_stat["token_counts"]["language_output_tokens"] * batch_size # task_total_stats["data_point_count"] += batch_size # Accumulate the number of processed items # cand_embeds = cand_embeds[:len(full_eval_cand_dataset)] # all_cand_ids = all_cand_ids[:len(full_eval_cand_dataset)] # if local_rank == 0: # cand_embed_dict = {cand_id: embed for cand_id, embed in zip(all_cand_ids, cand_embeds)} # with open(cand_embed_path, 'wb') as f: pickle.dump(cand_embed_dict, f) # print_master(f"Saved candidate embeddings to {cand_embed_path}") # if dist.is_initialized(): # dist.barrier() # # --- New: Calculate and Save Task-level Inference Statistics (on master rank only) --- # if local_rank == 0: # if task_total_stats["data_point_count"] > 0: # final_task_stats = { # "task_name": dataset_name, # "data_point_count": task_total_stats["data_point_count"], # "inference_times": { # "total_inference_time_seconds": task_total_stats["total_inference_time_seconds"], # "avg_inference_time_per_item_seconds": task_total_stats["total_inference_time_seconds"] / task_total_stats["data_point_count"], # "module_average_times_per_call": {}, # Average per call to the module # "module_total_times_seconds": {}, # Total time spent in the module # "module_calls_count": {}, # Number of times the module was called # }, # "token_counts": { # "total_visual_tokens": task_total_stats["token_counts"]["visual_tokens"], # "avg_visual_tokens_per_item": task_total_stats["token_counts"]["visual_tokens"] / task_total_stats["data_point_count"], # "total_language_input_tokens_raw": task_total_stats["token_counts"]["language_input_tokens_raw"], # "avg_language_input_tokens_raw_per_item": task_total_stats["token_counts"]["language_input_tokens_raw"] / task_total_stats["data_point_count"], # "total_llm_total_input_tokens": task_total_stats["token_counts"]["llm_total_input_tokens"], # "avg_llm_total_input_tokens_per_item": task_total_stats["token_counts"]["llm_total_input_tokens"] / task_total_stats["data_point_count"], # "total_language_output_tokens": task_total_stats["token_counts"]["language_output_tokens"], # "avg_language_output_tokens_per_item": task_total_stats["token_counts"]["language_output_tokens"] / task_total_stats["data_point_count"], # } # } # for module_name, stats in task_total_stats["module_inference_times"].items(): # final_task_stats["inference_times"]["module_total_times_seconds"][module_name] = stats["total"] # final_task_stats["inference_times"]["module_calls_count"][module_name] = stats["count"] # if stats["count"] > 0: # final_task_stats["inference_times"]["module_average_times_per_call"][module_name] = stats["total"] / stats["count"] # else: # final_task_stats["inference_times"]["module_average_times_per_call"][module_name] = 0.0 # with open(inference_stats_path, 'w', encoding='utf-8') as f: # json.dump(final_task_stats, f, ensure_ascii=False, indent=4) # print_master(f"Inference statistics for {dataset_name} saved to: {inference_stats_path}") # else: # print_master(f"No data processed for {dataset_name}, skipping inference statistics output.") # # --- 3. Compute Scores (on master rank only) --- # score_path = os.path.join(data_args.encode_output_path, f"{dataset_name}_score.json") # if os.path.exists(score_path): # try: # with open(score_path, "r") as f: # score_dict = json.load(f) # print_master(f"Score of {dataset_name} (loaded from previous run): {score_path}") # formatted = {k: f"{v:.4f}" for k, v in score_dict.items()} # print_master(formatted) # # No `continue` here, as we want to ensure other files are processed/generated # except Exception as e: # print_master(f"Failed to load score for {dataset_name}, proceeding to recompute. Error: {e}") # # Proceed with score computation if not loaded or failed to load # with open(query_embed_path, 'rb') as f: qry_embeds = pickle.load(f) # with open(cand_embed_path, 'rb') as f: cand_embed_dict = pickle.load(f) # gt_infos = [json.loads(l) for l in open(dataset_info_path)] # pred_dicts = [] # rank_against_all_candidates = task_config.get("eval_type", "global") == "global" # if rank_against_all_candidates: # cand_keys = list(cand_embed_dict.keys()) # cand_embeds = np.stack([cand_embed_dict[key] for key in cand_keys]) # # Handle late-interaction scoring # if qry_embeds.ndim == 3: # Query: [N_q, L_q, H] | Candidate: [N_c, L_c, H] # qry_embed = torch.from_numpy(qry_embeds) # cand_embeds = [torch.from_numpy(np.array(t)) for t in cand_embeds] # scores = processor.score(qry_embed, cand_embeds, batch_size=64) # use ColPali score function # ranked_candids = torch.argsort(-scores, dim=1).cpu().numpy().tolist() # else: # Dense # cosine_scores = np.dot(qry_embeds, cand_embeds.T) # ranked_candids = np.argsort(-cosine_scores, axis=1) # for qid, (ranked_candid, gt_info) in tqdm(enumerate(zip(ranked_candids, gt_infos)), desc=f"Calculating scores for {dataset_name}"): # rel_docids = gt_info["label_name"] if isinstance(gt_info["label_name"], list) else [gt_info["label_name"]] # rel_scores = gt_info["rel_scores"] if "rel_scores" in gt_info else None # assert rel_scores is None or len(rel_docids) == len(rel_scores) # pred_dicts.append({ # "prediction": [cand_keys[i] for i in ranked_candid], # "label": rel_docids, # "rel_scores": rel_scores, # }) # else: # for qid, (qry_embed, gt_info) in tqdm(enumerate(zip(qry_embeds, gt_infos)), desc=f"Calculating scores for {dataset_name}"): # cand_embeds = np.stack([cand_embed_dict[key] for key in gt_info["cand_names"]]) # if qry_embeds.ndim == 3: # Query: [N_q, L_q, H] | Candidate: [N_c, L_c, H] # qry_embed = torch.from_numpy(np.array(qry_embed)).unsqueeze(0) # cand_embeds = [torch.from_numpy(np.array(t)) for t in cand_embeds] # scores = processor.score(qry_embed, cand_embeds, batch_size=1024) # use ColPali score function # ranked_candids = torch.argsort(-scores, dim=1).cpu().numpy().tolist()[0] # else: # cosine_score = np.dot(qry_embed, cand_embeds.T) # ranked_candids = np.argsort(-cosine_score) # rel_docids = gt_info["label_name"] if isinstance(gt_info["label_name"], list) else [gt_info["label_name"]] # rel_scores = gt_info["rel_scores"] if "rel_scores" in gt_info else None # assert rel_scores is None or len(rel_docids) == len(rel_scores) # pred_dicts.append({ # "prediction": [gt_info["cand_names"][i] for i in ranked_candids], # "label": rel_docids, # "rel_scores": rel_scores, # }) # score_path = os.path.join(data_args.encode_output_path, f"{dataset_name}_score.json") # pred_path = os.path.join(data_args.encode_output_path, f"{dataset_name}_pred.jsonl") # metrics_to_report = task_config["metrics"] if task_config.get("metrics", None) is not None else ["hit", "ndcg", "precision", "recall", "f1", "map", "mrr"] # metrics = RankingMetrics(metrics_to_report) # score_dict = metrics.evaluate(pred_dicts) # formatted = {k: f"{v:.4f}" for k, v in score_dict.items()} # score_dict["num_pred"] = len(pred_dicts) # score_dict["num_data"] = len(gt_infos) # print_master(f"Score of {dataset_name}:") # print_master(formatted) # print_master(f"Outputting final score to: {score_path}") # with open(score_path, "w") as f: # json.dump(score_dict, f, indent=4) # with open(pred_path, "w") as f: # for pred in pred_dicts: # f.write(json.dumps(pred) + '\n') # if __name__ == "__main__": # main() ########################################################################################## # ######################################################################################## # #分query和cand进行统计 # 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 # 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 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 import MMEBModel # from src.model.processor import get_backbone_name, load_processor, COLPALI # from src.utils import batch_to_device, print_rank, print_master # logging.basicConfig(level=logging.INFO, format='[%(asctime)s] %(levelname)s [%(name)s:%(lineno)s] %(message)s') # logger = logging.getLogger(__name__) # # --- 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 # # print_master(f"DEBUG: Initial model type in register_model_hooks: {type(model)}") # if hasattr(model, 'encoder') and model.encoder is not None: # # print_master(f"DEBUG: model has 'encoder' attribute. Type of model.encoder: {type(model.encoder)}") # # 使用从 'src' 路径导入的 Qwen2VLForConditionalGeneration 进行检查 # if isinstance(model.encoder, _Qwen2VLForConditionalGeneration_src): # # print_master("Detected MMEBModel structure, registering hooks on model.encoder's sub-modules.") # core_model = model.encoder # else: # print_master(f"WARNING: model.encoder is not an instance of _Qwen2VLForConditionalGeneration_src. Its type is {type(model.encoder)}. Hooks will be registered on top-level model if applicable.") # else: # print_master("WARNING: Model structure does not have an 'encoder' attribute. Registering hooks directly on top-level modules.") # # 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]: # Added list to return type for batch_stats # """ # Encodes embeddings for a given dataset using the model, handling both standard and # late-interaction models in a DDP-safe manner. # """ # 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 = [] # model.eval() # # Register hooks for the model once per encode_embeddings call # registered_hooks = register_model_hooks(model) # with torch.no_grad(): # for inputs, dataset_info in tqdm(loader, desc=f"{description} (rank {local_rank})", disable=local_rank > 0): # # --- Reset 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 else 1 # Determine actual batch size # with torch.autocast(enabled=True, dtype=torch.bfloat16, device_type="cuda"): # start_inference_time = time.time() # 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]) # end_inference_time = time.time() # # --- Update total LLM input tokens after the model call --- # if 'input_ids' in inputs and inputs['input_ids'] is not None: # # `inputs['input_ids'].shape[1]` gives the sequence length, # # which is the number of tokens per item in the batch. # # To get total tokens for the batch, multiply by batch size. # token_info["total_llm_input_tokens"] = inputs['input_ids'].shape[1] # # Approximation for text_input_tokens # 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"]) # Ensure not negative # # --- 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"], # } # } # # 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) # Append the stats for this batch # # --- Print Inference Timing and Token Statistics per Batch (Optional, for debugging) --- # 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([]), [], [] # Return empty list for batch_stats_list as well # # === DDP Synchronization and Padding for Late-Interaction Models === # if is_late_interaction: # if dist.is_initialized(): # # 1. Find the global maximum sequence length across all ranks # 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 all local embeddings to the 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: # Should not happen if model is consistently late-interaction # padded_embeds.append(reps_batch) # embeds_tensor = torch.cat(padded_embeds, dim=0).contiguous() # else: # Standard dense models # 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...") # # Use the more efficient all_gather_into_tensor for tensors # 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() # # Gather metadata, for which all_gather_object is appropriate # 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] # # --- New: Gather batch_stats_list from all ranks --- # 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] # else: # all_gt_infos = local_gt_infos # final_embeddings = embeds_tensor.cpu().float().numpy() # all_batch_stats = batch_stats_list # If not DDP, just use local list # return final_embeddings, all_gt_infos, all_batch_stats # def main(): # 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 # # DEBUG PRINTS for Distributed Setup # 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()}") # 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) # 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) # # --- 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}') # # --- DDP-Safe Model Loading --- # # Step 1: Only the master process (rank 0) downloads the model. # 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}...") # # Step 2: All processes wait here. The non-master processes will pause # # until the master process (rank 0) finishes downloading and exits this barrier. # if torch.distributed.is_initialized(): # torch.distributed.barrier() # # Step 3: Now that the model is cached, the non-master processes load it from the local cache. # 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) # 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()): # # Initialize task-level statistics accumulators for QUERY # query_total_stats = { # "total_inference_time_seconds": 0.0, # "module_inference_times": { # "Qwen2VisionTransformerPretrainedModel": {"total": 0.0, "count": 0}, # "PatchMerger": {"total": 0.0, "count": 0}, # "Qwen2VLModel": {"total": 0.0, "count": 0}, # "Linear": {"total": 0.0, "count": 0}, # }, # "token_counts": { # "visual_tokens": 0, # "language_input_tokens_raw": 0, # "llm_total_input_tokens": 0, # "language_output_tokens": 0, # }, # "data_point_count": 0 # Number of image-text pairs processed # } # # Initialize task-level statistics accumulators for CANDIDATE # cand_total_stats = { # "total_inference_time_seconds": 0.0, # "module_inference_times": { # "Qwen2VisionTransformerPretrainedModel": {"total": 0.0, "count": 0}, # "PatchMerger": {"total": 0.0, "count": 0}, # "Qwen2VLModel": {"total": 0.0, "count": 0}, # "Linear": {"total": 0.0, "count": 0}, # }, # "token_counts": { # "visual_tokens": 0, # "language_input_tokens_raw": 0, # "llm_total_input_tokens": 0, # "language_output_tokens": 0, # }, # "data_point_count": 0 # Number of image-text pairs processed # } # if dist.is_initialized(): # dist.barrier() # print_master(f"\n--- Evaluating {dataset_name} ---") # query_embed_path = os.path.join(data_args.encode_output_path, f"{dataset_name}_qry") # cand_embed_path = os.path.join(data_args.encode_output_path, f"{dataset_name}_tgt") # dataset_info_path = os.path.join(data_args.encode_output_path, f"{dataset_name}_info.jsonl") # # New: Define distinct paths for query and candidate inference statistics output # query_inference_stats_path = os.path.join(data_args.encode_output_path, f"{dataset_name}_qry_inference_stats.json") # cand_inference_stats_path = os.path.join(data_args.encode_output_path, f"{dataset_name}_cand_inference_stats.json") # do_query = not os.path.exists(query_embed_path) or not os.path.exists(dataset_info_path) # do_cand = not os.path.exists(cand_embed_path) # if do_query or do_cand: # if data_args.data_basedir is not None: # # Construct full paths for data files if --data_basedir is provided # 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 # # Pad datasets to be divisible by world_size before splitting # if dist.is_initialized(): # 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 # # --- 1. Compute Query Embeddings --- # if do_query: # print_master("Encoding queries...") # eval_qry_collator = MultimodalEvalDataCollator(processor, model_args, data_args, "qry") # eval_qry_loader = DataLoader(eval_qry_dataset, batch_size=training_args.per_device_eval_batch_size, collate_fn=eval_qry_collator, num_workers=training_args.dataloader_num_workers) # # Modified: capture batch_stats_list # query_embeds, gt_infos, qry_batch_stats = encode_embeddings(model, eval_qry_loader, training_args, model_args, padded_qry_dataset, encode_side="qry", description=f"Queries for {dataset_name}") # # Accumulate query statistics # for batch_stat in qry_batch_stats: # batch_size = batch_stat["batch_size"] # query_total_stats["total_inference_time_seconds"] += batch_stat["total_inference_time_seconds"] # for module_name, module_stats in batch_stat["module_inference_times"].items(): # if module_name in query_total_stats["module_inference_times"]: # query_total_stats["module_inference_times"][module_name]["total"] += module_stats["total"] # query_total_stats["module_inference_times"][module_name]["count"] += module_stats["count"] # query_total_stats["token_counts"]["visual_tokens"] += batch_stat["token_counts"]["visual_tokens"] * batch_size # query_total_stats["token_counts"]["language_input_tokens_raw"] += batch_stat["token_counts"]["language_input_tokens_raw"] * batch_size # query_total_stats["token_counts"]["llm_total_input_tokens"] += batch_stat["token_counts"]["llm_total_input_tokens"] * batch_size # query_total_stats["token_counts"]["language_output_tokens"] += batch_stat["token_counts"]["language_output_tokens"] * batch_size # query_total_stats["data_point_count"] += batch_size # Accumulate the number of processed items # query_embeds = query_embeds[:len(full_eval_qry_dataset)] # gt_infos = gt_infos[:len(full_eval_qry_dataset)] # if local_rank == 0: # with open(query_embed_path, 'wb') as f: # pickle.dump(query_embeds, f) # with open(dataset_info_path, 'w') as f: # for info in gt_infos: # f.write(json.dumps(info) + '\n') # print_master(f"Saved query embeddings to {query_embed_path}") # # Save query-specific inference statistics # if query_total_stats["data_point_count"] > 0: # final_query_stats = { # "task_name": dataset_name, # "encode_side": "query", # "data_point_count": query_total_stats["data_point_count"], # "inference_times": { # "total_inference_time_seconds": query_total_stats["total_inference_time_seconds"], # "avg_inference_time_per_item_seconds": query_total_stats["total_inference_time_seconds"] / query_total_stats["data_point_count"], # "module_average_times_per_call": {}, # "module_total_times_seconds": {}, # "module_calls_count": {}, # }, # "token_counts": { # "total_visual_tokens": query_total_stats["token_counts"]["visual_tokens"], # "avg_visual_tokens_per_item": query_total_stats["token_counts"]["visual_tokens"] / query_total_stats["data_point_count"], # "total_language_input_tokens_raw": query_total_stats["token_counts"]["language_input_tokens_raw"], # "avg_language_input_tokens_raw_per_item": query_total_stats["token_counts"]["language_input_tokens_raw"] / query_total_stats["data_point_count"], # "total_llm_total_input_tokens": query_total_stats["token_counts"]["llm_total_input_tokens"], # "avg_llm_total_input_tokens_per_item": query_total_stats["token_counts"]["llm_total_input_tokens"] / query_total_stats["data_point_count"], # "total_language_output_tokens": query_total_stats["token_counts"]["language_output_tokens"], # "avg_language_output_tokens_per_item": query_total_stats["token_counts"]["language_output_tokens"] / query_total_stats["data_point_count"], # } # } # for module_name, stats in query_total_stats["module_inference_times"].items(): # final_query_stats["inference_times"]["module_total_times_seconds"][module_name] = stats["total"] # final_query_stats["inference_times"]["module_calls_count"][module_name] = stats["count"] # if stats["count"] > 0: # final_query_stats["inference_times"]["module_average_times_per_call"][module_name] = stats["total"] / stats["count"] # else: # final_query_stats["inference_times"]["module_average_times_per_call"][module_name] = 0.0 # with open(query_inference_stats_path, 'w', encoding='utf-8') as f: # json.dump(final_query_stats, f, ensure_ascii=False, indent=4) # print_master(f"Query inference statistics for {dataset_name} saved to: {query_inference_stats_path}") # else: # print_master(f"No query data processed for {dataset_name}, skipping query inference statistics output.") # if dist.is_initialized(): # dist.barrier() # # --- 2. Compute Candidate Embeddings --- # if do_cand: # print_master("Encoding candidates...") # eval_cand_collator = MultimodalEvalDataCollator(processor, model_args, data_args, "cand") # eval_cand_loader = DataLoader(eval_cand_dataset, batch_size=training_args.per_device_eval_batch_size, collate_fn=eval_cand_collator, num_workers=training_args.dataloader_num_workers) # # Modified: capture batch_stats_list # cand_embeds, all_cand_ids, cand_batch_stats = encode_embeddings(model, eval_cand_loader, training_args, model_args, padded_cand_dataset, encode_side="cand", description=f"Candidates for {dataset_name}") # # Accumulate candidate statistics (similar logic as query) # for batch_stat in cand_batch_stats: # batch_size = batch_stat["batch_size"] # cand_total_stats["total_inference_time_seconds"] += batch_stat["total_inference_time_seconds"] # for module_name, module_stats in batch_stat["module_inference_times"].items(): # if module_name in cand_total_stats["module_inference_times"]: # cand_total_stats["module_inference_times"][module_name]["total"] += module_stats["total"] # cand_total_stats["module_inference_times"][module_name]["count"] += module_stats["count"] # cand_total_stats["token_counts"]["visual_tokens"] += batch_stat["token_counts"]["visual_tokens"] * batch_size # cand_total_stats["token_counts"]["language_input_tokens_raw"] += batch_stat["token_counts"]["language_input_tokens_raw"] * batch_size # cand_total_stats["token_counts"]["llm_total_input_tokens"] += batch_stat["token_counts"]["llm_total_input_tokens"] * batch_size # cand_total_stats["token_counts"]["language_output_tokens"] += batch_stat["token_counts"]["language_output_tokens"] * batch_size # cand_total_stats["data_point_count"] += batch_size # Accumulate the number of processed items # cand_embeds = cand_embeds[:len(full_eval_cand_dataset)] # all_cand_ids = all_cand_ids[:len(full_eval_cand_dataset)] # if local_rank == 0: # cand_embed_dict = {cand_id: embed for cand_id, embed in zip(all_cand_ids, cand_embeds)} # with open(cand_embed_path, 'wb') as f: pickle.dump(cand_embed_dict, f) # print_master(f"Saved candidate embeddings to {cand_embed_path}") # # Save candidate-specific inference statistics # if cand_total_stats["data_point_count"] > 0: # final_cand_stats = { # "task_name": dataset_name, # "encode_side": "candidate", # "data_point_count": cand_total_stats["data_point_count"], # "inference_times": { # "total_inference_time_seconds": cand_total_stats["total_inference_time_seconds"], # "avg_inference_time_per_item_seconds": cand_total_stats["total_inference_time_seconds"] / cand_total_stats["data_point_count"], # "module_average_times_per_call": {}, # "module_total_times_seconds": {}, # "module_calls_count": {}, # }, # "token_counts": { # "total_visual_tokens": cand_total_stats["token_counts"]["visual_tokens"], # "avg_visual_tokens_per_item": cand_total_stats["token_counts"]["visual_tokens"] / cand_total_stats["data_point_count"], # "total_language_input_tokens_raw": cand_total_stats["token_counts"]["language_input_tokens_raw"], # "avg_language_input_tokens_raw_per_item": cand_total_stats["token_counts"]["language_input_tokens_raw"] / cand_total_stats["data_point_count"], # "total_llm_total_input_tokens": cand_total_stats["token_counts"]["llm_total_input_tokens"], # "avg_llm_total_input_tokens_per_item": cand_total_stats["token_counts"]["llm_total_input_tokens"] / cand_total_stats["data_point_count"], # "total_language_output_tokens": cand_total_stats["token_counts"]["language_output_tokens"], # "avg_language_output_tokens_per_item": cand_total_stats["token_counts"]["language_output_tokens"] / cand_total_stats["data_point_count"], # } # } # for module_name, stats in cand_total_stats["module_inference_times"].items(): # final_cand_stats["inference_times"]["module_total_times_seconds"][module_name] = stats["total"] # final_cand_stats["inference_times"]["module_calls_count"][module_name] = stats["count"] # if stats["count"] > 0: # final_cand_stats["inference_times"]["module_average_times_per_call"][module_name] = stats["total"] / stats["count"] # else: # final_cand_stats["inference_times"]["module_average_times_per_call"][module_name] = 0.0 # with open(cand_inference_stats_path, 'w', encoding='utf-8') as f: # json.dump(final_cand_stats, f, ensure_ascii=False, indent=4) # print_master(f"Candidate inference statistics for {dataset_name} saved to: {cand_inference_stats_path}") # else: # print_master(f"No candidate data processed for {dataset_name}, skipping candidate inference statistics output.") # if dist.is_initialized(): # dist.barrier() # # --- 3. Compute Scores (on master rank only) --- # score_path = os.path.join(data_args.encode_output_path, f"{dataset_name}_score.json") # #################################################################################### # pred_path = os.path.join(data_args.encode_output_path, f"{dataset_name}_pred.jsonl") # score_detail_path = os.path.join(data_args.encode_output_path, f"{dataset_name}_score_details.jsonl") # 新文件，存相似度分数 # def append_score_detail(score_detail_list, qid, ranked_indices, score_vector, cand_ids, labels): # """追加一个 query 的候选分数详情""" # score_detail_list.append({ # "qid": int(qid), # "cand_scores": [ # {"cand_id": str(cand_ids[i]), "score": float(score_vector[i])} # for i in ranked_indices # ], # "label": labels # }) # #################################################################################### # if local_rank == 0: # if os.path.exists(score_path): # try: # with open(score_path, "r") as f: # score_dict = json.load(f) # print_master(f"Score of {dataset_name} (loaded from previous run): {score_path}") # formatted = {k: f"{v:.4f}" for k, v in score_dict.items()} # print_master(formatted) # # No `continue` here, as we want to ensure other files are processed/generated # except Exception as e: # print_master(f"Failed to load score for {dataset_name}, proceeding to recompute. Error: {e}") # # Proceed with score computation if not loaded or failed to load # with open(query_embed_path, 'rb') as f: qry_embeds = pickle.load(f) # with open(cand_embed_path, 'rb') as f: cand_embed_dict = pickle.load(f) # gt_infos = [json.loads(l) for l in open(dataset_info_path)] # pred_dicts = [] # score_detail_dicts = []################################### # rank_against_all_candidates = task_config.get("eval_type", "global") == "global" # # if rank_against_all_candidates: # # cand_keys = list(cand_embed_dict.keys()) # # cand_embeds = np.stack([cand_embed_dict[key] for key in cand_keys]) # # # Handle late-interaction scoring # # if qry_embeds.ndim == 3: # Query: [N_q, L_q, H] | Candidate: [N_c, L_c, H] # # qry_embed = torch.from_numpy(qry_embeds) # # cand_embeds = [torch.from_numpy(np.array(t)) for t in cand_embeds] # # scores = processor.score(qry_embed, cand_embeds, batch_size=64) # use ColPali score function # # ranked_candids = torch.argsort(-scores, dim=1).cpu().numpy().tolist() # # scores = scores.cpu().numpy() # # else: # Dense # # cosine_scores = np.dot(qry_embeds, cand_embeds.T) # # ranked_candids = np.argsort(-cosine_scores, axis=1) # ##################################################### # if rank_against_all_candidates: # cand_keys = list(cand_embed_dict.keys()) # cand_embeds = np.stack([cand_embed_dict[key] for key in cand_keys]) # if qry_embeds.ndim == 3: # Late-interaction # qry_embed_t = torch.from_numpy(qry_embeds) # cand_embeds_t = [torch.from_numpy(np.array(t)) for t in cand_embeds] # sim_matrix = processor.score(qry_embed_t, cand_embeds_t, batch_size=64).cpu().numpy() # [N_q, N_c] # else: # Dense # sim_matrix = np.dot(qry_embeds, cand_embeds.T) # [N_q, N_c] # ranked_all = np.argsort(-sim_matrix, axis=1) # ######################################################### # for qid, (ranked_candid, gt_info) in tqdm(enumerate(zip(ranked_candids, gt_infos)), desc=f"Calculating scores for {dataset_name}"): # rel_docids = gt_info["label_name"] if isinstance(gt_info["label_name"], list) else [gt_info["label_name"]] # rel_scores = gt_info["rel_scores"] if "rel_scores" in gt_info else None # assert rel_scores is None or len(rel_docids) == len(rel_scores) # pred_dicts.append({ # "prediction": [cand_keys[i] for i in ranked_candid], # "label": rel_docids, # "rel_scores": rel_scores, # }) # ################################# 新增：详细相似度字典 # append_score_detail(score_detail_dicts, qid, ranked_indices, sim_matrix[qid], cand_keys, rel_docids) # ######################################## # # else: # # for qid, (qry_embed, gt_info) in tqdm(enumerate(zip(qry_embeds, gt_infos)), desc=f"Calculating scores for {dataset_name}"): # # cand_embeds = np.stack([cand_embed_dict[key] for key in gt_info["cand_names"]]) # # if qry_embeds.ndim == 3: # Query: [N_q, L_q, H] | Candidate: [N_c, L_c, H] # # qry_embed = torch.from_numpy(np.array(qry_embed)).unsqueeze(0) # # cand_embeds = [torch.from_numpy(np.array(t)) for t in cand_embeds] # # scores = processor.score(qry_embed, cand_embeds, batch_size=1024) # use ColPali score function # # ranked_candids = torch.argsort(-scores, dim=1).cpu().numpy().tolist()[0] # # else: # # cosine_score = np.dot(qry_embed, cand_embeds.T) # # ranked_candids = np.argsort(-cosine_score) # # rel_docids = gt_info["label_name"] if isinstance(gt_info["label_name"], list) else [gt_info["label_name"]] # # rel_scores = gt_info["rel_scores"] if "rel_scores" in gt_info else None # # assert rel_scores is None or len(rel_docids) == len(rel_scores) # # pred_dicts.append({ # # "prediction": [gt_info["cand_names"][i] for i in ranked_candids], # # "label": rel_docids, # # "rel_scores": rel_scores, # # }) # ####################################################################### # else: # 非全局 # for qid, (qry_embed, gt_info) in tqdm(enumerate(zip(qry_embeds, gt_infos)), desc=f"Calculating scores for {dataset_name}"): # cand_ids_local = gt_info["cand_names"] # cand_embeds = np.stack([cand_embed_dict[key] for key in cand_ids_local]) # if qry_embeds.ndim == 3: # Late-interaction # qry_embed_t = torch.from_numpy(np.array(qry_embed)).unsqueeze(0) # [1, Lq, H] # cand_embeds_t = [torch.from_numpy(np.array(t)) for t in cand_embeds] # sim_vec = processor.score(qry_embed_t, cand_embeds_t, batch_size=1024).cpu().numpy()[0] # [N_c] # else: # Dense # sim_vec = np.dot(qry_embed, cand_embeds.T) # [N_c] # ranked_indices = np.argsort(-sim_vec) # rel_docids = gt_info["label_name"] if isinstance(gt_info["label_name"], list) else [gt_info["label_name"]] # rel_scores = gt_info["rel_scores"] if "rel_scores" in gt_info else None # assert rel_scores is None or len(rel_docids) == len(rel_scores) # pred_dicts.append({ # "prediction": [cand_ids_local[i] for i in ranked_indices], # "label": rel_docids, # "rel_scores": rel_scores, # }) # # 新增：分数详情 # append_score_detail(score_detail_dicts, qid, ranked_indices, sim_vec, cand_ids_local, rel_docids) # ########################################## 保存预测和分数 # with open(score_detail_path, "w") as f: # 新增 # for detail in score_detail_dicts: # f.write(json.dumps(detail) + '\n') # print_master(f"Detailed score file saved to: {score_detail_path}") # metrics_to_report = task_config["metrics"] if task_config.get("metrics", None) is not None else ["hit", "ndcg", "precision", "recall", "f1", "map", "mrr"] # metrics = RankingMetrics(metrics_to_report) # score_dict = metrics.evaluate(pred_dicts) # formatted = {k: f"{v:.4f}" for k, v in score_dict.items()} # score_dict["num_pred"] = len(pred_dicts) # score_dict["num_data"] = len(gt_infos) # print_master(f"Score of {dataset_name}:") # print_master(formatted) # print_master(f"Outputting final score to: {score_path}") # with open(score_path, "w") as f: # json.dump(score_dict, f, indent=4) # with open(pred_path, "w") as f: # for pred in pred_dicts: # f.write(json.dumps(pred) + '\n') # #################################################################### # score_path = os.path.join(data_args.encode_output_path, f"{dataset_name}_score.json") # pred_path = os.path.join(data_args.encode_output_path, f"{dataset_name}_pred.jsonl") # metrics_to_report = task_config["metrics"] if task_config.get("metrics", None) is not None else ["hit", "ndcg", "precision", "recall", "f1", "map", "mrr"] # metrics = RankingMetrics(metrics_to_report) # score_dict = metrics.evaluate(pred_dicts) # formatted = {k: f"{v:.4f}" for k, v in score_dict.items()} # score_dict["num_pred"] = len(pred_dicts) # score_dict["num_data"] = len(gt_infos) # print_master(f"Score of {dataset_name}:") # print_master(formatted) # print_master(f"Outputting final score to: {score_path}") # with open(score_path, "w") as f: # json.dump(score_dict, f, indent=4) # with open(pred_path, "w") as f: # for pred in pred_dicts: # f.write(json.dumps(pred) + '\n') # if __name__ == '__main__': # main() ###################################################################################################### ####################################################################################################### #为了可视化把mask值也输出 ######################################################################################## #分query和cand进行统计 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 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.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 import MMEBModel from src.model.processor import get_backbone_name, load_processor, COLPALI from src.utils import batch_to_device, print_rank, print_master logging.basicConfig(level=logging.INFO, format='[%(asctime)s] %(levelname)s [%(name)s:%(lineno)s] %(message)s') logger = logging.getLogger(__name__) # --- 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 # print_master(f"DEBUG: Initial model type in register_model_hooks: {type(model)}") if hasattr(model, 'encoder') and model.encoder is not None: print_master(f"DEBUG: model has 'encoder' attribute. Type of model.encoder: {type(model.encoder)}") else: print_master("WARNING: Model structure does not have an 'encoder' attribute. Registering hooks directly on top-level modules.") # 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 image token masks locally --- local_img_token_masks = [] # 每个样本一个元素：None 或 [bool, ...] 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 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() if encode_side == "qry": output = model(qry=inputs) # torch.set_printoptions(threshold=10000) # print('output:', output) # exit() 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]) end_inference_time = time.time() # --- NEW: 提取并保存本 batch 的 image_token_bool_masks --- # 期望 MMEBModel 的 output 中直接或间接包含 'image_token_bool_masks' img_masks_raw = None if isinstance(output, dict): img_masks_raw = _search_key(output, "image_token_bool_masks") # 可选：若你在 MMEBModel 上挂了属性，也可以尝试读取 if img_masks_raw is None and hasattr(model, "image_token_bool_masks"): img_masks_raw = getattr(model, "image_token_bool_masks") img_masks_serializable = _to_serializable_mask_list(img_masks_raw, current_batch_size) local_img_token_masks.extend(img_masks_serializable) # --- 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"], } } # 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] 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 return final_embeddings, all_gt_infos, all_batch_stats, all_img_token_masks # CHANGED def main(): 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 # DEBUG PRINTS for Distributed Setup 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()}") 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) 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) # --- 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}') # --- DDP-Safe Model Loading --- # Step 1: Only the master process (rank 0) downloads the model. 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}...") # Step 2: All processes wait here. The non-master processes will pause # until the master process (rank 0) finishes downloading and exits this barrier. if torch.distributed.is_initialized(): torch.distributed.barrier() # Step 3: Now that the model is cached, the non-master processes load it from the local cache. 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) 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()): # Initialize task-level statistics accumulators for QUERY query_total_stats = { "total_inference_time_seconds": 0.0, "module_inference_times": { "Qwen2VisionTransformerPretrainedModel": {"total": 0.0, "count": 0}, "PatchMerger": {"total": 0.0, "count": 0}, "Qwen2VLModel": {"total": 0.0, "count": 0}, "Linear": {"total": 0.0, "count": 0}, }, "token_counts": { "visual_tokens": 0, "language_input_tokens_raw": 0, "llm_total_input_tokens": 0, "language_output_tokens": 0, }, "data_point_count": 0 # Number of image-text pairs processed } # Initialize task-level statistics accumulators for CANDIDATE cand_total_stats = { "total_inference_time_seconds": 0.0, "module_inference_times": { "Qwen2VisionTransformerPretrainedModel": {"total": 0.0, "count": 0}, "PatchMerger": {"total": 0.0, "count": 0}, "Qwen2VLModel": {"total": 0.0, "count": 0}, "Linear": {"total": 0.0, "count": 0}, }, "token_counts": { "visual_tokens": 0, "language_input_tokens_raw": 0, "llm_total_input_tokens": 0, "language_output_tokens": 0, }, "data_point_count": 0 # Number of image-text pairs processed } if dist.is_initialized(): dist.barrier() print_master(f"\n--- Evaluating {dataset_name} ---") query_embed_path = os.path.join(data_args.encode_output_path, f"{dataset_name}_qry") cand_embed_path = os.path.join(data_args.encode_output_path, f"{dataset_name}_tgt") dataset_info_path = os.path.join(data_args.encode_output_path, f"{dataset_name}_info.jsonl") # New: Define distinct paths for query and candidate inference statistics output query_inference_stats_path = os.path.join(data_args.encode_output_path, f"{dataset_name}_qry_inference_stats.json") cand_inference_stats_path = os.path.join(data_args.encode_output_path, f"{dataset_name}_cand_inference_stats.json") do_query = not os.path.exists(query_embed_path) or not os.path.exists(dataset_info_path) do_cand = not os.path.exists(cand_embed_path) if do_query or do_cand: if data_args.data_basedir is not None: # Construct full paths for data files if --data_basedir is provided 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 # Pad datasets to be divisible by world_size before splitting if dist.is_initialized(): 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 # --- 1. Compute Query Embeddings --- if do_query: print_master("Encoding queries...") eval_qry_collator = MultimodalEvalDataCollator(processor, model_args, data_args, "qry") eval_qry_loader = DataLoader(eval_qry_dataset, batch_size=training_args.per_device_eval_batch_size, collate_fn=eval_qry_collator, num_workers=training_args.dataloader_num_workers) # Modified: capture batch_stats_list query_embeds, gt_infos, qry_batch_stats, qry_img_masks = encode_embeddings(model, eval_qry_loader, training_args, model_args, padded_qry_dataset, encode_side="qry", description=f"Queries for {dataset_name}") # Accumulate query statistics for batch_stat in qry_batch_stats: batch_size = batch_stat["batch_size"] query_total_stats["total_inference_time_seconds"] += batch_stat["total_inference_time_seconds"] for module_name, module_stats in batch_stat["module_inference_times"].items(): if module_name in query_total_stats["module_inference_times"]: query_total_stats["module_inference_times"][module_name]["total"] += module_stats["total"] query_total_stats["module_inference_times"][module_name]["count"] += module_stats["count"] query_total_stats["token_counts"]["visual_tokens"] += batch_stat["token_counts"]["visual_tokens"] * batch_size query_total_stats["token_counts"]["language_input_tokens_raw"] += batch_stat["token_counts"]["language_input_tokens_raw"] * batch_size query_total_stats["token_counts"]["llm_total_input_tokens"] += batch_stat["token_counts"]["llm_total_input_tokens"] * batch_size query_total_stats["token_counts"]["language_output_tokens"] += batch_stat["token_counts"]["language_output_tokens"] * batch_size query_total_stats["data_point_count"] += batch_size # Accumulate the number of processed items query_embeds = query_embeds[:len(full_eval_qry_dataset)] gt_infos = gt_infos[:len(full_eval_qry_dataset)] if local_rank == 0: with open(query_embed_path, 'wb') as f: pickle.dump(query_embeds, f) with open(dataset_info_path, 'w') as f: for info in gt_infos: f.write(json.dumps(info) + '\n') print_master(f"Saved query embeddings to {query_embed_path}") qry_img_masks_path = os.path.join(data_args.encode_output_path, f"{dataset_name}_qry_img_token_masks.jsonl") with open(qry_img_masks_path, 'w', encoding='utf-8') as f: for i, m in enumerate(qry_img_masks[:len(full_eval_qry_dataset)]): f.write(json.dumps({"index": i, "mask": m}, ensure_ascii=False) + "\n") print_master(f"Saved query image token masks to {qry_img_masks_path}") # Save query-specific inference statistics if query_total_stats["data_point_count"] > 0: final_query_stats = { "task_name": dataset_name, "encode_side": "query", "data_point_count": query_total_stats["data_point_count"], "inference_times": { "total_inference_time_seconds": query_total_stats["total_inference_time_seconds"], "avg_inference_time_per_item_seconds": query_total_stats["total_inference_time_seconds"] / query_total_stats["data_point_count"], "module_average_times_per_call": {}, "module_total_times_seconds": {}, "module_calls_count": {}, }, "token_counts": { "total_visual_tokens": query_total_stats["token_counts"]["visual_tokens"], "avg_visual_tokens_per_item": query_total_stats["token_counts"]["visual_tokens"] / query_total_stats["data_point_count"], "total_language_input_tokens_raw": query_total_stats["token_counts"]["language_input_tokens_raw"], "avg_language_input_tokens_raw_per_item": query_total_stats["token_counts"]["language_input_tokens_raw"] / query_total_stats["data_point_count"], "total_llm_total_input_tokens": query_total_stats["token_counts"]["llm_total_input_tokens"], "avg_llm_total_input_tokens_per_item": query_total_stats["token_counts"]["llm_total_input_tokens"] / query_total_stats["data_point_count"], "total_language_output_tokens": query_total_stats["token_counts"]["language_output_tokens"], "avg_language_output_tokens_per_item": query_total_stats["token_counts"]["language_output_tokens"] / query_total_stats["data_point_count"], } } for module_name, stats in query_total_stats["module_inference_times"].items(): final_query_stats["inference_times"]["module_total_times_seconds"][module_name] = stats["total"] final_query_stats["inference_times"]["module_calls_count"][module_name] = stats["count"] if stats["count"] > 0: final_query_stats["inference_times"]["module_average_times_per_call"][module_name] = stats["total"] / stats["count"] else: final_query_stats["inference_times"]["module_average_times_per_call"][module_name] = 0.0 with open(query_inference_stats_path, 'w', encoding='utf-8') as f: json.dump(final_query_stats, f, ensure_ascii=False, indent=4) print_master(f"Query inference statistics for {dataset_name} saved to: {query_inference_stats_path}") else: print_master(f"No query data processed for {dataset_name}, skipping query inference statistics output.") if dist.is_initialized(): dist.barrier() # --- 2. Compute Candidate Embeddings --- if do_cand: print_master("Encoding candidates...") eval_cand_collator = MultimodalEvalDataCollator(processor, model_args, data_args, "cand") eval_cand_loader = DataLoader(eval_cand_dataset, batch_size=training_args.per_device_eval_batch_size, collate_fn=eval_cand_collator, num_workers=training_args.dataloader_num_workers) # Modified: capture batch_stats_list cand_embeds, all_cand_ids, cand_batch_stats, cand_img_masks = encode_embeddings(model, eval_cand_loader, training_args, model_args, padded_cand_dataset, encode_side="cand", description=f"Candidates for {dataset_name}") # Accumulate candidate statistics (similar logic as query) for batch_stat in cand_batch_stats: batch_size = batch_stat["batch_size"] cand_total_stats["total_inference_time_seconds"] += batch_stat["total_inference_time_seconds"] for module_name, module_stats in batch_stat["module_inference_times"].items(): if module_name in cand_total_stats["module_inference_times"]: cand_total_stats["module_inference_times"][module_name]["total"] += module_stats["total"] cand_total_stats["module_inference_times"][module_name]["count"] += module_stats["count"] cand_total_stats["token_counts"]["visual_tokens"] += batch_stat["token_counts"]["visual_tokens"] * batch_size cand_total_stats["token_counts"]["language_input_tokens_raw"] += batch_stat["token_counts"]["language_input_tokens_raw"] * batch_size cand_total_stats["token_counts"]["llm_total_input_tokens"] += batch_stat["token_counts"]["llm_total_input_tokens"] * batch_size cand_total_stats["token_counts"]["language_output_tokens"] += batch_stat["token_counts"]["language_output_tokens"] * batch_size cand_total_stats["data_point_count"] += batch_size # Accumulate the number of processed items cand_embeds = cand_embeds[:len(full_eval_cand_dataset)] all_cand_ids = all_cand_ids[:len(full_eval_cand_dataset)] if local_rank == 0: cand_embed_dict = {cand_id: embed for cand_id, embed in zip(all_cand_ids, cand_embeds)} with open(cand_embed_path, 'wb') as f: pickle.dump(cand_embed_dict, f) print_master(f"Saved candidate embeddings to {cand_embed_path}") cand_img_masks_path = os.path.join(data_args.encode_output_path, f"{dataset_name}_cand_img_token_masks.jsonl") with open(cand_img_masks_path, 'w', encoding='utf-8') as f: for cid, m in zip(all_cand_ids[:len(full_eval_cand_dataset)], cand_img_masks[:len(full_eval_cand_dataset)]): f.write(json.dumps({"cand_id": str(cid), "mask": m}, ensure_ascii=False) + "\n") print_master(f"Saved candidate image token masks to {cand_img_masks_path}") # Save candidate-specific inference statistics if cand_total_stats["data_point_count"] > 0: final_cand_stats = { "task_name": dataset_name, "encode_side": "candidate", "data_point_count": cand_total_stats["data_point_count"], "inference_times": { "total_inference_time_seconds": cand_total_stats["total_inference_time_seconds"], "avg_inference_time_per_item_seconds": cand_total_stats["total_inference_time_seconds"] / cand_total_stats["data_point_count"], "module_average_times_per_call": {}, "module_total_times_seconds": {}, "module_calls_count": {}, }, "token_counts": { "total_visual_tokens": cand_total_stats["token_counts"]["visual_tokens"], "avg_visual_tokens_per_item": cand_total_stats["token_counts"]["visual_tokens"] / cand_total_stats["data_point_count"], "total_language_input_tokens_raw": cand_total_stats["token_counts"]["language_input_tokens_raw"], "avg_language_input_tokens_raw_per_item": cand_total_stats["token_counts"]["language_input_tokens_raw"] / cand_total_stats["data_point_count"], "total_llm_total_input_tokens": cand_total_stats["token_counts"]["llm_total_input_tokens"], "avg_llm_total_input_tokens_per_item": cand_total_stats["token_counts"]["llm_total_input_tokens"] / cand_total_stats["data_point_count"], "total_language_output_tokens": cand_total_stats["token_counts"]["language_output_tokens"], "avg_language_output_tokens_per_item": cand_total_stats["token_counts"]["language_output_tokens"] / cand_total_stats["data_point_count"], } } for module_name, stats in cand_total_stats["module_inference_times"].items(): final_cand_stats["inference_times"]["module_total_times_seconds"][module_name] = stats["total"] final_cand_stats["inference_times"]["module_calls_count"][module_name] = stats["count"] if stats["count"] > 0: final_cand_stats["inference_times"]["module_average_times_per_call"][module_name] = stats["total"] / stats["count"] else: final_cand_stats["inference_times"]["module_average_times_per_call"][module_name] = 0.0 with open(cand_inference_stats_path, 'w', encoding='utf-8') as f: json.dump(final_cand_stats, f, ensure_ascii=False, indent=4) print_master(f"Candidate inference statistics for {dataset_name} saved to: {cand_inference_stats_path}") else: print_master(f"No candidate data processed for {dataset_name}, skipping candidate inference statistics output.") if dist.is_initialized(): dist.barrier() # --- 3. Compute Scores (on master rank only) --- score_path = os.path.join(data_args.encode_output_path, f"{dataset_name}_score.json") #################################################################################### pred_path = os.path.join(data_args.encode_output_path, f"{dataset_name}_pred.jsonl") score_detail_path = os.path.join(data_args.encode_output_path, f"{dataset_name}_score_details.jsonl") # 新文件，存相似度分数 def append_score_detail(score_detail_list, qid, ranked_indices, score_vector, cand_ids, labels): """追加一个 query 的候选分数详情""" score_detail_list.append({ "qid": int(qid), "cand_scores": [ {"cand_id": str(cand_ids[i]), "score": float(score_vector[i])} for i in ranked_indices ], "label": labels }) #################################################################################### if local_rank == 0: if os.path.exists(score_path): try: with open(score_path, "r") as f: score_dict = json.load(f) print_master(f"Score of {dataset_name} (loaded from previous run): {score_path}") formatted = {k: f"{v:.4f}" for k, v in score_dict.items()} print_master(formatted) # No `continue` here, as we want to ensure other files are processed/generated except Exception as e: print_master(f"Failed to load score for {dataset_name}, proceeding to recompute. Error: {e}") # Proceed with score computation if not loaded or failed to load with open(query_embed_path, 'rb') as f: qry_embeds = pickle.load(f) with open(cand_embed_path, 'rb') as f: cand_embed_dict = pickle.load(f) gt_infos = [json.loads(l) for l in open(dataset_info_path)] pred_dicts = [] score_detail_dicts = []################################### rank_against_all_candidates = task_config.get("eval_type", "global") == "global" # if rank_against_all_candidates: # cand_keys = list(cand_embed_dict.keys()) # cand_embeds = np.stack([cand_embed_dict[key] for key in cand_keys]) # # Handle late-interaction scoring # if qry_embeds.ndim == 3: # Query: [N_q, L_q, H] | Candidate: [N_c, L_c, H] # qry_embed = torch.from_numpy(qry_embeds) # cand_embeds = [torch.from_numpy(np.array(t)) for t in cand_embeds] # scores = processor.score(qry_embed, cand_embeds, batch_size=64) # use ColPali score function # ranked_candids = torch.argsort(-scores, dim=1).cpu().numpy().tolist() # scores = scores.cpu().numpy() # else: # Dense # cosine_scores = np.dot(qry_embeds, cand_embeds.T) # ranked_candids = np.argsort(-cosine_scores, axis=1) ##################################################### if rank_against_all_candidates: cand_keys = list(cand_embed_dict.keys()) cand_embeds = np.stack([cand_embed_dict[key] for key in cand_keys]) if qry_embeds.ndim == 3: # Late-interaction qry_embed_t = torch.from_numpy(qry_embeds) cand_embeds_t = [torch.from_numpy(np.array(t)) for t in cand_embeds] sim_matrix = processor.score(qry_embed_t, cand_embeds_t, batch_size=64).cpu().numpy() # [N_q, N_c] else: # Dense sim_matrix = np.dot(qry_embeds, cand_embeds.T) # [N_q, N_c] ranked_candids = np.argsort(-sim_matrix, axis=1) ######################################################### for qid, (ranked_candid, gt_info) in tqdm(enumerate(zip(ranked_candids, gt_infos)), desc=f"Calculating scores for {dataset_name}"): rel_docids = gt_info["label_name"] if isinstance(gt_info["label_name"], list) else [gt_info["label_name"]] rel_scores = gt_info["rel_scores"] if "rel_scores" in gt_info else None assert rel_scores is None or len(rel_docids) == len(rel_scores) pred_dicts.append({ "prediction": [cand_keys[i] for i in ranked_candid], "label": rel_docids, "rel_scores": rel_scores, }) ################################# 新增：详细相似度字典 append_score_detail(score_detail_dicts, qid, ranked_candid, sim_matrix[qid], cand_keys, rel_docids) ######################################## # else: # for qid, (qry_embed, gt_info) in tqdm(enumerate(zip(qry_embeds, gt_infos)), desc=f"Calculating scores for {dataset_name}"): # cand_embeds = np.stack([cand_embed_dict[key] for key in gt_info["cand_names"]]) # if qry_embeds.ndim == 3: # Query: [N_q, L_q, H] | Candidate: [N_c, L_c, H] # qry_embed = torch.from_numpy(np.array(qry_embed)).unsqueeze(0) # cand_embeds = [torch.from_numpy(np.array(t)) for t in cand_embeds] # scores = processor.score(qry_embed, cand_embeds, batch_size=1024) # use ColPali score function # ranked_candids = torch.argsort(-scores, dim=1).cpu().numpy().tolist()[0] # else: # cosine_score = np.dot(qry_embed, cand_embeds.T) # ranked_candids = np.argsort(-cosine_score) # rel_docids = gt_info["label_name"] if isinstance(gt_info["label_name"], list) else [gt_info["label_name"]] # rel_scores = gt_info["rel_scores"] if "rel_scores" in gt_info else None # assert rel_scores is None or len(rel_docids) == len(rel_scores) # pred_dicts.append({ # "prediction": [gt_info["cand_names"][i] for i in ranked_candids], # "label": rel_docids, # "rel_scores": rel_scores, # }) ####################################################################### else: # 非全局 for qid, (qry_embed, gt_info) in tqdm(enumerate(zip(qry_embeds, gt_infos)), desc=f"Calculating scores for {dataset_name}"): cand_ids_local = gt_info["cand_names"] cand_embeds = np.stack([cand_embed_dict[key] for key in cand_ids_local]) if qry_embeds.ndim == 3: # Late-interaction qry_embed_t = torch.from_numpy(np.array(qry_embed)).unsqueeze(0) # [1, Lq, H] cand_embeds_t = [torch.from_numpy(np.array(t)) for t in cand_embeds] sim_vec = processor.score(qry_embed_t, cand_embeds_t, batch_size=1024).cpu().numpy()[0] # [N_c] else: # Dense sim_vec = np.dot(qry_embed, cand_embeds.T) # [N_c] ranked_indices = np.argsort(-sim_vec) rel_docids = gt_info["label_name"] if isinstance(gt_info["label_name"], list) else [gt_info["label_name"]] rel_scores = gt_info["rel_scores"] if "rel_scores" in gt_info else None assert rel_scores is None or len(rel_docids) == len(rel_scores) pred_dicts.append({ "prediction": [cand_ids_local[i] for i in ranked_indices], "label": rel_docids, "rel_scores": rel_scores, }) # 新增：分数详情 append_score_detail(score_detail_dicts, qid, ranked_indices, sim_vec, cand_ids_local, rel_docids) ########################################## 保存预测和分数 with open(score_detail_path, "w") as f: # 新增 for detail in score_detail_dicts: f.write(json.dumps(detail) + '\n') print_master(f"Detailed score file saved to: {score_detail_path}") metrics_to_report = task_config["metrics"] if task_config.get("metrics", None) is not None else ["hit", "ndcg", "precision", "recall", "f1", "map", "mrr"] metrics = RankingMetrics(metrics_to_report) score_dict = metrics.evaluate(pred_dicts) formatted = {k: f"{v:.4f}" for k, v in score_dict.items()} score_dict["num_pred"] = len(pred_dicts) score_dict["num_data"] = len(gt_infos) print_master(f"Score of {dataset_name}:") print_master(formatted) print_master(f"Outputting final score to: {score_path}") with open(score_path, "w") as f: json.dump(score_dict, f, indent=4) with open(pred_path, "w") as f: for pred in pred_dicts: f.write(json.dumps(pred) + '\n') #################################################################### score_path = os.path.join(data_args.encode_output_path, f"{dataset_name}_score.json") pred_path = os.path.join(data_args.encode_output_path, f"{dataset_name}_pred.jsonl") metrics_to_report = task_config["metrics"] if task_config.get("metrics", None) is not None else ["hit", "ndcg", "precision", "recall", "f1", "map", "mrr"] metrics = RankingMetrics(metrics_to_report) score_dict = metrics.evaluate(pred_dicts) formatted = {k: f"{v:.4f}" for k, v in score_dict.items()} score_dict["num_pred"] = len(pred_dicts) score_dict["num_data"] = len(gt_infos) print_master(f"Score of {dataset_name}:") print_master(formatted) print_master(f"Outputting final score to: {score_path}") with open(score_path, "w") as f: json.dump(score_dict, f, indent=4) with open(pred_path, "w") as f: for pred in pred_dicts: f.write(json.dumps(pred) + '\n') if __name__ == '__main__': main()