wehub-resource-sync
625 lines
25 KiB
Python
625 lines
25 KiB
Python
import os
|
|
import math
|
|
import random
|
|
import logging
|
|
import datasets
|
|
import numpy as np
|
|
import torch.distributed as dist
|
|
from dataclasses import dataclass
|
|
from torch.utils.data import Dataset
|
|
from transformers import (
|
|
PreTrainedTokenizer,
|
|
DataCollatorWithPadding,
|
|
TrainerCallback,
|
|
TrainerState,
|
|
TrainerControl
|
|
)
|
|
|
|
from .AbsArguments import AbsEmbedderDataArguments, AbsEmbedderTrainingArguments
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
class AbsEmbedderTrainDataset(Dataset):
|
|
"""Abstract class for training dataset.
|
|
|
|
Args:
|
|
args (AbsEmbedderDataArguments): Data arguments.
|
|
tokenizer (PreTrainedTokenizer): Tokenizer to use.
|
|
"""
|
|
def __init__(
|
|
self,
|
|
args: AbsEmbedderDataArguments,
|
|
tokenizer: PreTrainedTokenizer
|
|
):
|
|
self.args = args
|
|
self.tokenizer = tokenizer
|
|
self.shuffle_ratio = args.shuffle_ratio
|
|
|
|
train_datasets = []
|
|
for data_dir in args.train_data:
|
|
if not os.path.isdir(data_dir):
|
|
if not (data_dir.endswith('.json') or data_dir.endswith('.jsonl')): continue
|
|
temp_dataset = self._load_dataset(data_dir)
|
|
if len(temp_dataset) == 0: continue
|
|
train_datasets.append(temp_dataset)
|
|
else:
|
|
for file in os.listdir(data_dir):
|
|
if not (file.endswith('.json') or file.endswith('.jsonl')): continue
|
|
temp_dataset = self._load_dataset(os.path.join(data_dir, file))
|
|
if len(temp_dataset) == 0: continue
|
|
train_datasets.append(temp_dataset)
|
|
self.dataset = datasets.concatenate_datasets(train_datasets)
|
|
|
|
def _load_dataset(self, file_path: str):
|
|
"""Load dataset from path.
|
|
|
|
Args:
|
|
file_path (str): Path to load the datasets from.
|
|
|
|
Raises:
|
|
ValueError: `pos_scores` and `neg_scores` not found in the features of training data
|
|
|
|
Returns:
|
|
datasets.Dataset: Loaded HF dataset.
|
|
"""
|
|
safe_rank = dist.get_rank() if dist.is_initialized() else 0
|
|
if safe_rank == 0:
|
|
logger.info(f'loading data from {file_path} ...')
|
|
|
|
temp_dataset = datasets.load_dataset('json', data_files=file_path, split='train', cache_dir=self.args.cache_path)
|
|
if len(temp_dataset) > self.args.max_example_num_per_dataset:
|
|
temp_dataset = temp_dataset.select(random.sample(list(range(len(temp_dataset))), self.args.max_example_num_per_dataset))
|
|
if not self.args.knowledge_distillation:
|
|
if 'pos_scores' in temp_dataset.column_names:
|
|
temp_dataset = temp_dataset.remove_columns(['pos_scores'])
|
|
if 'neg_scores' in temp_dataset.column_names:
|
|
temp_dataset = temp_dataset.remove_columns(['neg_scores'])
|
|
else:
|
|
if 'pos_scores' not in temp_dataset.column_names or 'neg_scores' not in temp_dataset.column_names:
|
|
raise ValueError(f"`pos_scores` and `neg_scores` not found in the features of training data in {file_path}, which is necessary when using knowledge distillation.")
|
|
return temp_dataset
|
|
|
|
def _shuffle_text(self, text):
|
|
"""shuffle the input text.
|
|
|
|
Args:
|
|
text (str): Input text.
|
|
|
|
Returns:
|
|
str: Shuffled text.
|
|
"""
|
|
if self.shuffle_ratio > 0 and len(text) > 100 and random.random() < self.shuffle_ratio:
|
|
split_text = []
|
|
chunk_size = len(text)//3 + 1
|
|
for i in range(0, len(text), chunk_size):
|
|
split_text.append(text[i:i+chunk_size])
|
|
random.shuffle(split_text)
|
|
return " ".join(split_text)
|
|
else:
|
|
return text
|
|
|
|
def __len__(self):
|
|
return len(self.dataset)
|
|
|
|
def __getitem__(self, item):
|
|
data = self.dataset[item]
|
|
train_group_size = self.args.train_group_size
|
|
|
|
query = data['query']
|
|
if self.args.query_instruction_for_retrieval is not None:
|
|
query = self.args.query_instruction_format.format(
|
|
data['prompt'] if 'prompt' in data else self.args.query_instruction_for_retrieval,
|
|
query
|
|
)
|
|
|
|
passages = []
|
|
teacher_scores = []
|
|
|
|
assert isinstance(data['pos'], list) and isinstance(data['neg'], list)
|
|
|
|
pos_idx = random.choice(list(range(len(data['pos']))))
|
|
passages.append(self._shuffle_text(data['pos'][pos_idx]))
|
|
|
|
neg_all_idx = list(range(len(data['neg'])))
|
|
if len(data['neg']) < train_group_size - 1:
|
|
num = math.ceil((train_group_size - 1) / len(data['neg']))
|
|
neg_idxs = random.sample(neg_all_idx * num, train_group_size - 1)
|
|
else:
|
|
neg_idxs = random.sample(neg_all_idx, self.args.train_group_size - 1)
|
|
for neg_idx in neg_idxs:
|
|
passages.append(data['neg'][neg_idx])
|
|
|
|
if self.args.knowledge_distillation:
|
|
assert isinstance(data['pos_scores'], list) and isinstance(data['neg_scores'], list)
|
|
teacher_scores.append(data['pos_scores'][pos_idx])
|
|
for neg_idx in neg_idxs:
|
|
teacher_scores.append(data['neg_scores'][neg_idx])
|
|
if not all(isinstance(score, (int, float)) for score in teacher_scores):
|
|
raise ValueError(f"pos_score or neg_score must be digit")
|
|
else:
|
|
teacher_scores = None
|
|
|
|
if self.args.passage_instruction_for_retrieval is not None:
|
|
passages = [
|
|
self.args.passage_instruction_format.format(
|
|
self.args.passage_instruction_for_retrieval, p
|
|
)
|
|
for p in passages
|
|
]
|
|
|
|
return query, passages, teacher_scores
|
|
|
|
@dataclass
|
|
class AbsEmbedderCollator(DataCollatorWithPadding):
|
|
"""
|
|
The abstract embedder collator.
|
|
"""
|
|
query_max_len: int = 32
|
|
passage_max_len: int = 128
|
|
sub_batch_size: int = -1
|
|
|
|
def __call__(self, features):
|
|
queries = [f[0] for f in features]
|
|
passages = [f[1] for f in features]
|
|
teacher_scores = [f[2] for f in features]
|
|
if teacher_scores[0] is None:
|
|
teacher_scores = None
|
|
elif isinstance(teacher_scores[0], list):
|
|
teacher_scores = sum(teacher_scores, [])
|
|
|
|
if isinstance(queries[0], list):
|
|
queries = sum(queries, [])
|
|
if isinstance(passages[0], list):
|
|
passages = sum(passages, [])
|
|
|
|
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.query_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
|
|
))
|
|
return {
|
|
"queries": q_collated,
|
|
"passages": d_collated,
|
|
"teacher_scores": teacher_scores,
|
|
"no_in_batch_neg_flag": False
|
|
}
|
|
|
|
|
|
class AbsEmbedderSameDatasetTrainDataset(AbsEmbedderTrainDataset):
|
|
"""Abstract class for training dataset that samples batches from same dataset.
|
|
|
|
Args:
|
|
args (AbsEmbedderDataArguments): Data arguments.
|
|
default_batch_size (int): The default batch size for training.
|
|
seed (int): Random seed.
|
|
tokenizer (PreTrainedTokenizer): Tokenizer to use.
|
|
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: AbsEmbedderDataArguments,
|
|
default_batch_size: int,
|
|
seed: int,
|
|
tokenizer: PreTrainedTokenizer,
|
|
process_index: int=0,
|
|
num_processes: int=1
|
|
):
|
|
self.args = args
|
|
self.shuffle_ratio = args.shuffle_ratio
|
|
self.defaut_batch_size = default_batch_size
|
|
self.deterministic_generator = np.random.default_rng(seed)
|
|
self.tokenizer = tokenizer
|
|
self.process_index = process_index
|
|
self.num_processes = num_processes
|
|
|
|
self.step = 0
|
|
|
|
train_datasets = []
|
|
each_data_idxs = []
|
|
batch_size_idxs = []
|
|
no_in_batch_neg_flags = []
|
|
cur_all_num = 0
|
|
|
|
small_threshold = args.small_threshold
|
|
drop_threshold = args.drop_threshold
|
|
|
|
for data_dir in args.train_data:
|
|
if not os.path.isdir(data_dir):
|
|
# Add `no_in_batch_neg` **suffix** to `data_dir` to indicate that this dataset does not use in-batch negatives
|
|
no_in_batch_neg_flag = data_dir.split('.')[-2].endswith('no_in_batch_neg')
|
|
if not (data_dir.endswith('.json') or data_dir.endswith('.jsonl')): continue
|
|
temp_dataset = self._load_dataset(data_dir)
|
|
|
|
if len(temp_dataset) == 0 or len(temp_dataset) < small_threshold: continue
|
|
else:
|
|
train_datasets.append(temp_dataset)
|
|
each_data_idxs.append(np.arange(len(temp_dataset)) + cur_all_num)
|
|
cur_all_num += len(temp_dataset)
|
|
batch_size_idxs.append(self._get_file_batch_size(temp_dataset, default_batch_size))
|
|
no_in_batch_neg_flags.append(no_in_batch_neg_flag)
|
|
|
|
else:
|
|
small_datasets = []
|
|
small_batch_size = math.inf
|
|
|
|
# Add `no_in_batch_neg` **suffix** to `data_dir` to indicate that this dataset does not use in-batch negatives
|
|
no_in_batch_neg_flag = data_dir.endswith('no_in_batch_neg')
|
|
for file in os.listdir(data_dir):
|
|
if not (file.endswith('.json') or file.endswith('.jsonl')): continue
|
|
temp_dataset = self._load_dataset(os.path.join(data_dir, file))
|
|
|
|
if len(temp_dataset) == 0: continue
|
|
elif len(temp_dataset) < small_threshold:
|
|
small_datasets.append(temp_dataset)
|
|
small_batch_size = min(small_batch_size, self._get_file_batch_size(temp_dataset, default_batch_size))
|
|
else:
|
|
train_datasets.append(temp_dataset)
|
|
each_data_idxs.append(np.arange(len(temp_dataset)) + cur_all_num)
|
|
cur_all_num += len(temp_dataset)
|
|
batch_size_idxs.append(self._get_file_batch_size(temp_dataset, default_batch_size))
|
|
no_in_batch_neg_flags.append(no_in_batch_neg_flag)
|
|
|
|
if len(small_datasets) > 0:
|
|
small_dataset = datasets.concatenate_datasets(small_datasets)
|
|
if len(small_dataset) >= drop_threshold:
|
|
train_datasets.append(small_dataset)
|
|
each_data_idxs.append(np.arange(len(small_dataset)) + cur_all_num)
|
|
cur_all_num += len(small_dataset)
|
|
batch_size_idxs.append(small_batch_size)
|
|
no_in_batch_neg_flags.append(no_in_batch_neg_flag)
|
|
|
|
self.dataset = datasets.concatenate_datasets(train_datasets)
|
|
self.each_data_idxs = each_data_idxs
|
|
self.datasets_inxs = np.arange(len(each_data_idxs))
|
|
self.batch_size_idxs = batch_size_idxs
|
|
self.no_in_batch_neg_flags = no_in_batch_neg_flags
|
|
|
|
self.refresh_epoch()
|
|
|
|
def _load_dataset(self, file_path: str):
|
|
"""Load datset from given path.
|
|
|
|
Args:
|
|
file_path (str): The path to load or download from HF hub.
|
|
|
|
Returns:
|
|
datasets.Dataset: The loaded dataset.
|
|
"""
|
|
safe_rank = dist.get_rank() if dist.is_initialized() else 0
|
|
if safe_rank == 0:
|
|
logger.info(f'loading data from {file_path} ...')
|
|
|
|
temp_dataset = datasets.load_dataset('json', data_files=file_path, split='train', cache_dir=self.args.cache_path)
|
|
if len(temp_dataset) > self.args.max_example_num_per_dataset:
|
|
temp_dataset = temp_dataset.select(random.sample(list(range(len(temp_dataset))), self.args.max_example_num_per_dataset))
|
|
if not self.args.knowledge_distillation:
|
|
if 'pos_scores' in temp_dataset.column_names:
|
|
temp_dataset = temp_dataset.remove_columns(['pos_scores'])
|
|
if 'neg_scores' in temp_dataset.column_names:
|
|
temp_dataset = temp_dataset.remove_columns(['neg_scores'])
|
|
return temp_dataset
|
|
|
|
@staticmethod
|
|
def _get_file_batch_size(temp_dataset: datasets.Dataset, default_batch_size: int):
|
|
"""Get the appropriate batch size for the dataset.
|
|
|
|
Args:
|
|
temp_dataset (datasets.Dataset): Loaded :data:`datasets.Dataset` object.
|
|
default_batch_size (int): The default batch size to use if not specified in the dataset.
|
|
|
|
Returns:
|
|
int: The final batch size to use.
|
|
"""
|
|
if 'batch_size' in temp_dataset.column_names:
|
|
return temp_dataset['batch_size'][0]
|
|
if 'type' in temp_dataset.column_names:
|
|
data_type = temp_dataset['type'][0]
|
|
if 'symmetric' in data_type:
|
|
return default_batch_size // 2 # make the symmetric data have smaller batch size
|
|
return default_batch_size
|
|
|
|
def refresh_epoch(self):
|
|
"""
|
|
Refresh data for epoch.
|
|
"""
|
|
logger.info(f'-- Rank {self.process_index}: refresh data --')
|
|
self.deterministic_generator.shuffle(self.datasets_inxs)
|
|
|
|
batch_datas = []
|
|
for dataset_inx in self.datasets_inxs:
|
|
self.deterministic_generator.shuffle(self.each_data_idxs[dataset_inx])
|
|
cur_batch_size = self.batch_size_idxs[dataset_inx]*self.num_processes
|
|
no_in_batch_neg_flag = self.no_in_batch_neg_flags[dataset_inx]
|
|
for start_index in range(0, len(self.each_data_idxs[dataset_inx]), cur_batch_size):
|
|
# judge the last batch's length
|
|
if len(self.each_data_idxs[dataset_inx]) - start_index < cur_batch_size:
|
|
break
|
|
batch_datas.append((
|
|
self.each_data_idxs[dataset_inx][start_index:start_index+cur_batch_size],
|
|
no_in_batch_neg_flag
|
|
))
|
|
self.deterministic_generator.shuffle(batch_datas)
|
|
self.batch_datas = batch_datas
|
|
self.step = 0
|
|
|
|
def __len__(self):
|
|
return len(self.batch_datas) * self.num_processes
|
|
|
|
def __getitem__(self, _):
|
|
batch_indices, no_in_batch_neg_flag = self.batch_datas[self.step] # extend here
|
|
cur_batch_size = int(len(batch_indices) / self.num_processes)
|
|
batch_indices = batch_indices[self.process_index * cur_batch_size: (self.process_index + 1) * cur_batch_size]
|
|
batch_data = self.dataset[batch_indices]
|
|
self.step += 1
|
|
queries, passages, teacher_scores = self._create_batch_data(batch_raw_data=batch_data)
|
|
return queries, passages, teacher_scores, no_in_batch_neg_flag
|
|
|
|
def _get_train_group_size(self, batch_raw_data):
|
|
"""Get the training group size and data type.
|
|
|
|
Args:
|
|
batch_raw_data (datasets.Dataset): One batch of raw data.
|
|
|
|
Returns:
|
|
int: The training group size.
|
|
str: The type of data for the task.
|
|
"""
|
|
if 'type' in batch_raw_data:
|
|
data_type = batch_raw_data['type'][0]
|
|
if data_type in ['only_1neg']:
|
|
return 2, data_type
|
|
elif data_type in ['symmetric_class']:
|
|
return min(len(batch_raw_data['neg'][0]) + 1, self.args.train_group_size), data_type
|
|
else:
|
|
return self.args.train_group_size, data_type
|
|
elif 'train_group_size' in batch_raw_data:
|
|
train_group_size = batch_raw_data['train_group_size'][0]
|
|
if isinstance(train_group_size, int) and train_group_size > 0:
|
|
return train_group_size, None
|
|
else:
|
|
return self.args.train_group_size, None
|
|
return self.args.train_group_size, None
|
|
|
|
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)
|
|
|
|
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
|
|
]
|
|
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)
|
|
|
|
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.query_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
|
|
}
|
|
|
|
|
|
class EmbedderTrainerCallbackForDataRefresh(TrainerCallback):
|
|
"""
|
|
Callback class to inspect the state of the training loop and take decision.
|
|
"""
|
|
def __init__(self, train_dataset: AbsEmbedderSameDatasetTrainDataset):
|
|
self.train_dataset = train_dataset
|
|
|
|
def on_epoch_end(
|
|
self,
|
|
args: AbsEmbedderTrainingArguments,
|
|
state: TrainerState,
|
|
control: TrainerControl,
|
|
**kwargs
|
|
):
|
|
"""
|
|
Event called at the end of an epoch.
|
|
"""
|
|
self.train_dataset.refresh_epoch()
|