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()