wehub-resource-sync
290 lines
11 KiB
Python
290 lines
11 KiB
Python
# 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()
|