# Copyright (c) 2025 PaddlePaddle Authors. All Rights Reserved. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. import json import numpy as np import paddle import paddle.distributed as dist from tqdm import tqdm from paddlenlp.transformers import AutoConfig, AutoModel, AutoTokenizer class Embedding_Evaluation: def __init__( self, model_path, tokenizer_path, query_pos_passage_path, neg_passage_path, template="{text}", dimension=1024, max_src_len=8192, normalize=True, dtype=None, ): # initialize the tokenizer self.tokenizer = AutoTokenizer.from_pretrained( tokenizer_path, padding_side="right", truncation_side="right", ) self.config = AutoConfig.from_pretrained(model_path) self.config.embedding_negatives_cross_device = False self.dtype = dtype if dtype else self.config.dtype # Initialize the distributed environment dist.init_parallel_env() world_size = dist.get_world_size() if world_size > 1: print(f"Running in multi-GPU mode with {world_size} GPUs.") else: print("Running in single-GPU or CPU mode.") # Initialize the embedding model self.model = AutoModel.from_pretrained( model_path, config=self.config, dtype=self.dtype, low_cpu_mem_usage=False ) self.model.eval() self.query_pos_passage_path = query_pos_passage_path self.neg_passage_path = neg_passage_path self.template = template self.dimension = dimension self.max_src_len = max_src_len self.normalize = normalize def _preprocess(self, texts): """Pre-process inputs.""" template_prefix, template_suffix = self.template.split("{text}") prefix_tokens = self.tokenizer(template_prefix, add_special_tokens=False).input_ids suffix_tokens = self.tokenizer(template_suffix, add_special_tokens=False).input_ids # If the template does not contain a suffix token, add the EOS token if len(suffix_tokens) == 0: suffix_tokens = [self.tokenizer.eos_token_id] # If the template does not contain a prefix token, add the BOS token if len(prefix_tokens) == 0: prefix_tokens = [self.tokenizer.bos_token_id] available_len = self.max_src_len - len(prefix_tokens) - len(suffix_tokens) truncated_token_ids = self._batch_truncate_and_tokenize(texts, available_len) complete_token_ids = [prefix_tokens + tid + suffix_tokens for tid in truncated_token_ids] position_ids = [list(range(len(cid))) for cid in complete_token_ids] max_len = max([len(cid) for cid in complete_token_ids]) embedding_indices = [[idx, len(cid) - 1] for idx, cid in enumerate(complete_token_ids)] inputs = self.tokenizer.pad( { "input_ids": complete_token_ids, "position_ids": position_ids, "embedding_indices": embedding_indices, }, padding="max_length", return_attention_mask=True, max_length=max_len, return_tensors="pd", ) return inputs def _batch_truncate_and_tokenize(self, texts, available_len): """Tokenize the batch of texts.""" batch_tokenized = self.tokenizer( texts, add_special_tokens=False, padding=False, truncation=True, max_length=available_len ) truncated_token_ids = [token_ids for token_ids in batch_tokenized["input_ids"]] return truncated_token_ids def _forward(self, inputs, dimension): """Run model forward.""" input_type = type(inputs["input_ids"]) outputs = self.model(**inputs) if isinstance(outputs, input_type): hidden_states = outputs else: hidden_states = outputs[0] last_hidden_state = hidden_states[:, 0] if dimension > self.config.hidden_size: raise ValueError( f"Dimension ({dimension}) cannot be greater than hidden_size ({self.config.hidden_size})." ) elif dimension != self.config.hidden_size: last_hidden_state = last_hidden_state[:, :dimension] if self.normalize: last_hidden_state = paddle.nn.functional.normalize(last_hidden_state, axis=-1) last_hidden_state = last_hidden_state.astype("float16").tolist() return last_hidden_state @paddle.no_grad() def get_embedding(self, texts, dimension=None): """Get inference sequence.""" if dimension is None: dimension = self.dimension inputs = self._preprocess(texts) if self.config.model_type in ["xlm-roberta"]: del inputs["embedding_indices"] del inputs["position_ids"] outputs = self._forward(inputs, dimension) return outputs def evaluate(self): query_data_list = [] pos_passage_data_list = [] with open(self.query_pos_passage_path, "r") as f: for line in f: single_data = json.loads(line) query_data_list.append(single_data["query"]) pos_passage_data_list.append(single_data["pos_passage"][0]) neg_passage_data_list = [] with open(self.neg_passage_path, "r") as f: for line in f: single_data = json.loads(line) neg_passage_data_list.append(single_data["neg_passage"][0]) passage_data_list = pos_passage_data_list + neg_passage_data_list world_size = paddle.distributed.get_world_size() rank = paddle.distributed.get_rank() query_chunk_size = len(query_data_list) // world_size passage_chunk_size = len(passage_data_list) // world_size if rank == world_size - 1: # The last process handles the remaining data query_data_chunk = query_data_list[rank * query_chunk_size :] passage_data_chunk = passage_data_list[rank * passage_chunk_size :] else: query_data_chunk = query_data_list[rank * query_chunk_size : (rank + 1) * query_chunk_size] passage_data_chunk = passage_data_list[rank * passage_chunk_size : (rank + 1) * passage_chunk_size] batch_size = 4 # Adjust batch size according to your hardware and needs local_p_vecs = [] local_q_vecs = [] # Use tqdm to iterate over query_data_chunk and get embeddings in batches for batch in tqdm(range(0, len(passage_data_chunk), batch_size), desc="Processing passage embeddings"): batch_start = batch batch_end = min(batch_start + batch_size, len(passage_data_chunk)) batch_texts = passage_data_chunk[batch_start:batch_end] # Call get_embedding to obtain embeddings for the current batch batch_embeddings = self.get_embedding(batch_texts) local_p_vecs.extend(batch_embeddings) for batch in tqdm(range(0, len(query_data_chunk), batch_size), desc="Processing query embeddings"): batch_start = batch batch_end = min(batch_start + batch_size, len(query_data_chunk)) batch_texts = query_data_chunk[batch_start:batch_end] batch_embeddings = self.get_embedding(batch_texts) local_q_vecs.extend(batch_embeddings) local_p_vecs_file = f"local_p_vecs_rank_{rank}.npy" local_q_vecs_file = f"local_q_vecs_rank_{rank}.npy" np.save(local_p_vecs_file, local_p_vecs) np.save(local_q_vecs_file, local_q_vecs) dist.barrier() # Ensure all cards have reached this point before continuing if rank == 0: all_p_vecs_list = [] all_q_vecs_list = [] world_size = paddle.distributed.get_world_size() for i in range(world_size): local_p_vecs_file = f"local_p_vecs_rank_{i}.npy" local_q_vecs_file = f"local_q_vecs_rank_{i}.npy" # Load the embedding vector file from each process local_p_vecs = np.load(local_p_vecs_file) local_q_vecs = np.load(local_q_vecs_file) all_p_vecs_list.append(local_p_vecs) all_q_vecs_list.append(local_q_vecs) all_q_vecs = [] for q_vecs in all_q_vecs_list: all_q_vecs.extend(q_vecs) q_vecs = np.asarray(all_q_vecs, dtype=np.float32) all_p_vecs = [] for p_vecs in all_p_vecs_list: all_p_vecs.extend(p_vecs) p_vecs = np.asarray(all_p_vecs, dtype=np.float32) query_embedding_tensor = paddle.to_tensor(q_vecs, dtype=self.dtype) passage_embedding_tensor = paddle.to_tensor(p_vecs, dtype=self.dtype) similarity_matrix = self.calculate_cosine_similarity_matrix( query_embedding_tensor, passage_embedding_tensor ) query_num = len(query_data_list) true_answers = [i for i in range(query_num)] hit_count_10, hit_coun_5, hit_count_3, hit_count_1 = 0, 0, 0, 0 reciprocal_rank_sum_10, reciprocal_rank_sum_5, reciprocal_rank_sum_3 = 0, 0, 0 ndcg_10, ndcg_5, ndcg_3 = 0.0, 0.0, 0.0 for i in range(query_num): similarities = similarity_matrix[i] # get the sorted indices sorted_indices = paddle.argsort(-similarities) # find the index of the true answer true_answer_index = true_answers[i] rank = paddle.where(sorted_indices == true_answer_index)[0][0] + 1 # rank starts from 1 if rank <= 10: hit_count_10 += 1 reciprocal_rank_sum_10 += 1.0 / rank if rank <= 5: hit_coun_5 += 1 reciprocal_rank_sum_5 += 1.0 / rank if rank <= 3: hit_count_3 += 1 reciprocal_rank_sum_3 += 1.0 / rank if rank <= 1: hit_count_1 += 1 relevance_scores = [0] * 10 if rank <= 10: relevance_scores[rank - 1] = 1 ndcg_10 += self.calculate_ndcg(relevance_scores[:10], k=10) ndcg_5 += self.calculate_ndcg(relevance_scores[:5], k=5) ndcg_3 += self.calculate_ndcg(relevance_scores[:3], k=3) print(f"Hit rate when recall Top 10: ({hit_count_10*100./query_num:.2f}%)\n") print(f"Hit rate when recall Top 5: ({hit_coun_5*100./query_num:.2f}%)\n") print(f"Hit rate when recall Top 3: ({hit_count_3*100./query_num:.2f}%)\n") print(f"Hit rate when recall Top 1: ({hit_count_1*100./query_num:.2f}%)\n") print(f"MRR when recall Top 10: ({reciprocal_rank_sum_10.item() / query_num:.4f})\n") print(f"MRR when recall Top 5: ({reciprocal_rank_sum_5.item() / query_num:.4f})\n") print(f"MRR when recall Top 3: ({reciprocal_rank_sum_3.item() / query_num:.4f})\n") print(f"NDCG@10: ({ndcg_10/ query_num:.4f})\n") print(f"NDCG@5: ({ndcg_5/ query_num:.4f})\n") print(f"NDCG@3: ({ndcg_3/ query_num:.4f})\n") eval_result_dict = { "hit_rate@10": hit_count_10 / query_num, "hit_rate@5": hit_coun_5 / query_num, "hit_rate@3": hit_count_3 / query_num, "hit_rate@1": hit_count_1 / query_num, "mrr@10": reciprocal_rank_sum_10.item() / query_num, "mrr@5": reciprocal_rank_sum_5.item() / query_num, "mrr@3": reciprocal_rank_sum_3.item() / query_num, "ndcg@10": ndcg_10 / query_num, "ndcg@5": ndcg_5 / query_num, "ndcg@3": ndcg_3 / query_num, } return eval_result_dict def calculate_cosine_similarity_matrix(self, query_matrix, answer_matrix): """Calculate the cosine similarity between two matrices by processing query vectors one by one.""" num_queries = query_matrix.shape[0] num_answers = answer_matrix.shape[0] # Precompute the norms of answer vectors to save computation answer_norms = paddle.linalg.norm(answer_matrix, axis=1, keepdim=True) # Initialize the similarity matrix with zeros similarity_matrix = paddle.zeros((num_queries, num_answers)) # Process each query vector one by one for i in tqdm(range(num_queries)): query_vector = query_matrix[i : i + 1] # Extract the i-th query vector # Calculate the norm of the query vector query_norm = paddle.linalg.norm(query_vector, axis=1, keepdim=True) # Calculate the dot product between the query vector and all answer vectors dot_product = paddle.matmul(query_vector, answer_matrix, transpose_y=True) # Calculate the cosine similarity for the i-th query vector similarity_vector = dot_product / (query_norm * answer_norms.transpose((1, 0))) # Update the similarity matrix with the computed similarity vector similarity_matrix[i] = similarity_vector return similarity_matrix def calculate_ndcg(self, relevance_scores, k): """Calculate NDCG@k for a given set of relevance scores""" # Calculate DCG dcg = sum((rel) / np.log2(i + 2) for i, rel in enumerate(relevance_scores[:k])) # Calculate IDCG (Ideal DCG) ideal_relevance_scores = sorted(relevance_scores, reverse=True) idcg = sum((rel) / np.log2(i + 2) for i, rel in enumerate(ideal_relevance_scores[:k])) # Avoid division by zero if idcg == 0: return 0.0 # Calculate NDCG ndcg = dcg / idcg return ndcg if __name__ == "__main__": model_path = "BAAI/bge-m3" tokenizer_path = "BAAI/bge-m3" query_pos_passage_path = "./toy_data/toy_dev.json" neg_passage_path = "./toy_data/toy_dev_neg.json" eval = Embedding_Evaluation(model_path, tokenizer_path, query_pos_passage_path, neg_passage_path) print(eval.evaluate())