# 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 faiss import numpy as np import paddle import paddle.distributed as dist from tqdm import tqdm from paddlenlp.transformers import AutoConfig, AutoModel, AutoTokenizer class MiningNegativeSamples: def __init__( self, model_path, tokenizer_path, input_data_path, output_data_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.input_data_path = input_data_path self.output_data_path = output_data_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 mining(self): query_pos_dict = {} query_data_list = [] pos_data_list = [] temp_data_list = [] count = 0 with open(self.input_data_path, "r") as f: for line in tqdm(f): data = json.loads(line) query = data["query"] pos_passage = data["pos_passage"][0] if query not in query_pos_dict: temp_data_list.append(data) query_data_list.append(query) pos_data_list.append(pos_passage) query_pos_dict[query] = [pos_passage] else: # print('error1',query) count += 1 query_pos_dict[query].append(pos_passage) world_size = paddle.distributed.get_world_size() rank = paddle.distributed.get_rank() assert len(pos_data_list) == len(query_data_list) chunk_size = len(pos_data_list) // world_size if rank == world_size - 1: # The last process handles the remaining data pos_data_chunk = pos_data_list[rank * chunk_size :] query_data_chunk = query_data_list[rank * chunk_size :] else: pos_data_chunk = pos_data_list[rank * chunk_size : (rank + 1) * chunk_size] query_data_chunk = query_data_list[rank * chunk_size : (rank + 1) * 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(pos_data_chunk), batch_size), desc="Processing passage embeddings"): batch_start = batch batch_end = min(batch_start + batch_size, len(pos_data_chunk)) batch_texts = pos_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) index = faiss.IndexFlatIP(len(p_vecs[0])) p_vecs = np.asarray(p_vecs, dtype=np.float32) if paddle.is_compiled_with_cuda(): co = faiss.GpuMultipleClonerOptions() co.shard = True co.useFloat16 = False index = faiss.index_cpu_to_all_gpus(index, co=co) index.add(p_vecs) count = 0 batch_size = 16 output_data_list = [] for i in tqdm(range(0, len(q_vecs), batch_size)): batch_queries = query_data_list[i : i + batch_size] batch_q_vecs = q_vecs[i : i + batch_size] _, batch_ids = index.search(batch_q_vecs, k=10) for j, ids in enumerate(batch_ids): query = batch_queries[j] converted = [id for id in ids] neg_list = [] for id in converted: if pos_data_list[id] in query_pos_dict[query]: continue neg_list.append(pos_data_list[id]) # you can mining k negatives,there k==2 if len(neg_list) > 2: neg_list = neg_list[:2] assert query == temp_data_list[i + j]["query"] temp_data_list[i + j]["neg_passage"] = neg_list output_data_list.append(temp_data_list[i + j]) else: print("error2", query) count += 1 del index with open(self.output_data_path, "w", encoding="utf-8") as f: for item in tqdm(output_data_list): f.write(json.dumps(item, ensure_ascii=False)) f.write("\n") if __name__ == "__main__": input_data_path = "./toy_data/toy_source.json" output_data_path = "./toy_data/test_min_neg.json" model_path = "BAAI/bge-m3" tokenizer_path = "BAAI/bge-m3" test_mining = MiningNegativeSamples( model_path, tokenizer_path, input_data_path=input_data_path, output_data_path=output_data_path ) test_mining.mining()