wehub-resource-sync
292 lines
12 KiB
Python
292 lines
12 KiB
Python
import math
|
|
import random
|
|
import logging
|
|
from dataclasses import dataclass
|
|
from transformers import (
|
|
PreTrainedTokenizer,
|
|
DataCollatorWithPadding,
|
|
)
|
|
|
|
from FlagEmbedding.abc.finetune.embedder import AbsEmbedderSameDatasetTrainDataset
|
|
|
|
from .arguments import DecoderOnlyEmbedderICLDataArguments
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
class DecoderOnlyEmbedderICLSameDatasetTrainDataset(AbsEmbedderSameDatasetTrainDataset):
|
|
"""Dataset class for icl model.
|
|
|
|
Args:
|
|
args (DecoderOnlyEmbedderICLDataArguments): Data argument class for icl model.
|
|
default_batch_size (int): The default batch size.
|
|
seed (int): Random seed to use.
|
|
tokenizer (PreTrainedTokenizer): Tokenzier.
|
|
process_index (int, optional): Current process index. Defaults to 0.
|
|
num_processes (int, optional): Total number of processes. Defaults to 1.
|
|
"""
|
|
def __init__(
|
|
self,
|
|
args: DecoderOnlyEmbedderICLDataArguments,
|
|
default_batch_size: int,
|
|
seed: int,
|
|
tokenizer: PreTrainedTokenizer,
|
|
process_index: int=0,
|
|
num_processes: int=1
|
|
):
|
|
super().__init__(
|
|
args=args,
|
|
default_batch_size=default_batch_size,
|
|
seed=seed,
|
|
tokenizer=tokenizer,
|
|
process_index=process_index,
|
|
num_processes=num_processes
|
|
)
|
|
self.args: DecoderOnlyEmbedderICLDataArguments
|
|
|
|
self.suffix = self.tokenizer(f"{self.args.icl_suffix_str}{self.tokenizer.eos_token}", add_special_tokens=False)['input_ids']
|
|
|
|
self.prefix = self.tokenizer(f"{self.tokenizer.bos_token}", add_special_tokens=False)['input_ids']
|
|
|
|
def _create_batch_data(self, batch_raw_data):
|
|
"""Create a comple batch of data with queries, documents and teacher scores.
|
|
|
|
Args:
|
|
batch_raw_data (datasets.Dataset): One batch of raw data.
|
|
|
|
Returns:
|
|
List[str]: Queries with instruction format.
|
|
List[str]: Documents with instruction format.
|
|
List[float]: Teacher scores for model distillation.
|
|
"""
|
|
queries, passages, teacher_scores = [], [], []
|
|
|
|
train_group_size, data_type = self._get_train_group_size(batch_raw_data)
|
|
|
|
icl_pairs = []
|
|
|
|
for i in range(len(batch_raw_data['query'])):
|
|
if data_type is not None:
|
|
assert batch_raw_data['type'][i] == data_type, f"Data type is not consistent in the same batch"
|
|
|
|
queries.append(
|
|
self.args.query_instruction_format.format(
|
|
batch_raw_data['prompt'][i] if 'prompt' in batch_raw_data else self.args.query_instruction_for_retrieval,
|
|
batch_raw_data['query'][i]
|
|
)
|
|
)
|
|
tmp_passages = []
|
|
pos_idx = random.choice(list(range(len(batch_raw_data['pos'][i]))))
|
|
pos = self._shuffle_text(batch_raw_data['pos'][i][pos_idx])
|
|
tmp_passages.append(pos)
|
|
|
|
neg_all_idx = list(range(len(batch_raw_data['neg'][i])))
|
|
if len(batch_raw_data['neg'][i]) < train_group_size - 1:
|
|
num = math.ceil((train_group_size - 1) / len(batch_raw_data['neg'][i]))
|
|
neg_idxs = random.sample(neg_all_idx * num, train_group_size - 1)
|
|
else:
|
|
neg_idxs = random.sample(neg_all_idx, train_group_size - 1)
|
|
for neg_idx in neg_idxs:
|
|
tmp_passages.append(batch_raw_data['neg'][i][neg_idx])
|
|
|
|
if self.args.knowledge_distillation:
|
|
if 'pos_scores' in batch_raw_data and batch_raw_data['pos_scores'][i] is not None:
|
|
teacher_scores.append(batch_raw_data['pos_scores'][i][pos_idx])
|
|
for neg_idx in neg_idxs:
|
|
if 'neg_scores' in batch_raw_data and batch_raw_data['neg_scores'][i] is not None:
|
|
teacher_scores.append(batch_raw_data['neg_scores'][i][neg_idx])
|
|
else:
|
|
teacher_scores = None
|
|
|
|
if data_type is not None and data_type in ['symmetric_sts', 'symmetric_clustering']:
|
|
tmp_passages = [
|
|
self.args.query_instruction_format.format(
|
|
batch_raw_data['prompt'][i] if 'prompt' in batch_raw_data else self.args.query_instruction_for_retrieval,
|
|
p
|
|
) for p in tmp_passages
|
|
]
|
|
tmp_passages = self.tokenizer.batch_decode(
|
|
self.tokenizer(
|
|
tmp_passages,
|
|
max_length=self.args.passage_max_len - 1 - len(self.suffix),
|
|
truncation=True,
|
|
add_special_tokens=False,
|
|
)['input_ids']
|
|
)
|
|
for j in range(len(tmp_passages)):
|
|
tmp_passages[j] += self.args.icl_suffix_str
|
|
else:
|
|
if self.args.passage_instruction_for_retrieval is not None:
|
|
tmp_passages = [
|
|
self.args.passage_instruction_format.format(
|
|
self.args.passage_instruction_for_retrieval, p
|
|
) for p in tmp_passages
|
|
]
|
|
|
|
passages.extend(tmp_passages)
|
|
|
|
if teacher_scores is not None:
|
|
if len(teacher_scores) > 0 and len(passages) > 0:
|
|
assert len(teacher_scores) == len(passages)
|
|
|
|
# add icl pairs
|
|
if self.args.retrieval_use_examples or (
|
|
data_type in ['symmetric_sts', 'symmetric_clustering', 'symmetric_class']
|
|
):
|
|
if data_type == 'symmetric_clustering':
|
|
icl_pairs.append((
|
|
self.tokenizer.decode(
|
|
self.tokenizer(
|
|
queries[-1],
|
|
add_special_tokens=False
|
|
)['input_ids'][:self.args.example_query_max_len]
|
|
),
|
|
self.tokenizer.decode(
|
|
self.tokenizer(
|
|
batch_raw_data['category'][i], # use category as example
|
|
add_special_tokens=False
|
|
)['input_ids'][:self.args.example_passage_max_len]
|
|
)
|
|
))
|
|
else:
|
|
icl_pairs.append((
|
|
self.tokenizer.decode(
|
|
self.tokenizer(
|
|
queries[-1],
|
|
add_special_tokens=False
|
|
)['input_ids'][:self.args.example_query_max_len]
|
|
),
|
|
self.tokenizer.decode(
|
|
self.tokenizer(
|
|
pos,
|
|
add_special_tokens=False
|
|
)['input_ids'][:self.args.example_passage_max_len]
|
|
)
|
|
))
|
|
else:
|
|
icl_pairs = []
|
|
|
|
# handle queries
|
|
for i in range(len(queries)):
|
|
choices = random.choice([0, 1, 2, 3, 4, 5])
|
|
if choices > 0 and len(icl_pairs) > 0:
|
|
prefix_ids = random.sample(list(range(len(icl_pairs))), min(choices + 1, len(icl_pairs)))
|
|
if i in prefix_ids:
|
|
prefix_ids.remove(i)
|
|
prefix_ids = prefix_ids[:choices]
|
|
|
|
prefix = ''
|
|
for idx in prefix_ids:
|
|
tmp = prefix + self.args.icl_suffix_str.join(icl_pairs[idx]) + '\n\n'
|
|
if len(self.tokenizer(tmp)['input_ids']) > self.args.query_max_len - 512:
|
|
break
|
|
prefix = tmp
|
|
else:
|
|
prefix = ''
|
|
|
|
queries[i] = prefix + queries[i]
|
|
queries[i] = self.tokenizer.decode(
|
|
self.tokenizer(
|
|
queries[i],
|
|
max_length=self.args.query_max_len - len(self.prefix) - len(self.suffix),
|
|
truncation=True,
|
|
add_special_tokens=False
|
|
)['input_ids']
|
|
) + self.args.icl_suffix_str
|
|
|
|
return queries, passages, teacher_scores
|
|
|
|
|
|
@dataclass
|
|
class AbsEmbedderSameDatasetCollator(DataCollatorWithPadding):
|
|
"""
|
|
EmbedCollator for SameDataset.
|
|
Note that after using this collator, the training_args should be set as:
|
|
|
|
``training_args.per_device_train_batch_size = 1``
|
|
|
|
``training_args.dataloader_num_workers = 0 # avoid multi-processing``
|
|
"""
|
|
query_max_len: int = 32
|
|
passage_max_len: int = 128
|
|
sub_batch_size: int = -1
|
|
|
|
def __call__(self, features):
|
|
queries = features[0][0]
|
|
passages = features[0][1]
|
|
teacher_scores = features[0][2]
|
|
no_in_batch_neg_flag = features[0][3]
|
|
|
|
queries_inputs = self.tokenizer(
|
|
queries,
|
|
truncation=True,
|
|
max_length=self.query_max_len,
|
|
return_tensors=None
|
|
)
|
|
passages_inputs = self.tokenizer(
|
|
passages,
|
|
truncation=True,
|
|
max_length=self.passage_max_len,
|
|
return_tensors=None
|
|
)
|
|
|
|
if self.sub_batch_size is None or self.sub_batch_size <= 0:
|
|
q_collated = self.tokenizer.pad(
|
|
queries_inputs,
|
|
padding=self.padding,
|
|
max_length=self.query_max_len,
|
|
pad_to_multiple_of=self.pad_to_multiple_of,
|
|
return_tensors=self.return_tensors,
|
|
)
|
|
|
|
d_collated = self.tokenizer.pad(
|
|
passages_inputs,
|
|
padding=self.padding,
|
|
max_length=self.passage_max_len,
|
|
pad_to_multiple_of=self.pad_to_multiple_of,
|
|
return_tensors=self.return_tensors,
|
|
)
|
|
else:
|
|
batch_size = self.sub_batch_size
|
|
|
|
q_collated = []
|
|
for i in range(0, len(queries_inputs['attention_mask']), batch_size):
|
|
start = i
|
|
end = min(len(queries_inputs['attention_mask']), i + batch_size)
|
|
sub_features = {}
|
|
for k, v in queries_inputs.items():
|
|
sub_features[k] = v[start:end]
|
|
q_collated.append(self.tokenizer.pad(
|
|
sub_features,
|
|
padding=self.padding,
|
|
max_length=self.passage_max_len,
|
|
pad_to_multiple_of=self.pad_to_multiple_of,
|
|
return_tensors=self.return_tensors,
|
|
))
|
|
|
|
d_collated = []
|
|
for i in range(0, len(passages_inputs['attention_mask']), batch_size):
|
|
start = i
|
|
end = min(len(passages_inputs['attention_mask']), i + batch_size)
|
|
sub_features = {}
|
|
|
|
for k, v in passages_inputs.items():
|
|
sub_features[k] = v[start:end]
|
|
d_collated.append(self.tokenizer.pad(
|
|
sub_features,
|
|
padding=self.padding,
|
|
max_length=self.passage_max_len,
|
|
pad_to_multiple_of=self.pad_to_multiple_of,
|
|
return_tensors=self.return_tensors,
|
|
))
|
|
|
|
if isinstance(teacher_scores, list) and len(teacher_scores) == 0:
|
|
teacher_scores = None
|
|
|
|
return {
|
|
"queries": q_collated,
|
|
"passages": d_collated,
|
|
"teacher_scores": teacher_scores,
|
|
"no_in_batch_neg_flag": no_in_batch_neg_flag
|
|
}
|