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, BatchEncoding, DataCollatorForSeq2Seq ) from typing import List from .AbsArguments import AbsRerankerDataArguments logger = logging.getLogger(__name__) class AbsRerankerTrainDataset(Dataset): """Abstract class for reranker training dataset. Args: args (AbsRerankerDataArguments): Data arguments. tokenizer (PreTrainedTokenizer): Tokenizer to use. """ def __init__( self, args: AbsRerankerDataArguments, tokenizer: PreTrainedTokenizer ): self.args = args self.tokenizer = tokenizer 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) self.max_length = self.args.query_max_len + self.args.passage_max_len 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.args.shuffle_ratio > 0 and len(text) > 100 and random.random() < self.args.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 create_one_example(self, qry_encoding: str, doc_encoding: str): """Creates a single input example by encoding and preparing a query and document pair for the model. Args: qry_encoding (str): Query to be encoded. doc_encoding (str): Document to be encoded. Returns: dict: A dictionary containing tokenized and prepared inputs, ready for model consumption. """ qry_inputs = self.tokenizer.encode(qry_encoding, truncation=True, max_length=self.args.query_max_len + self.args.passage_max_len // 4, add_special_tokens=False) doc_inputs = self.tokenizer.encode(doc_encoding, truncation=True, max_length=self.args.passage_max_len + self.args.query_max_len // 2, add_special_tokens=False) item = self.tokenizer.prepare_for_model( qry_inputs, doc_inputs, truncation='only_second', max_length=self.args.query_max_len + self.args.passage_max_len, padding=False, ) return item 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_rerank is not None: query = self.args.query_instruction_format.format( data['query_prompt'] if 'query_prompt' in data else self.args.query_instruction_for_rerank, 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_rerank is not None: passages = [ self.args.passage_instruction_format.format( data['passage_prompt'] if 'passage_prompt' in data else self.args.passage_instruction_for_rerank, p ) for p in passages ] batch_data = [] for passage in passages: batch_data.append(self.create_one_example(query, passage)) return batch_data, teacher_scores @dataclass class AbsRerankerCollator(DataCollatorWithPadding): """ The abstract reranker collator. """ query_max_len: int = 32 passage_max_len: int = 128 def __call__(self, features) -> List[BatchEncoding]: teacher_scores = [f[1] for f in features] if teacher_scores[0] is None: teacher_scores = None elif isinstance(teacher_scores[0], list): teacher_scores = sum(teacher_scores, []) features = [f[0] for f in features] if isinstance(features[0], list): features = sum(features, []) collated = self.tokenizer.pad( features, padding=self.padding, max_length=self.query_max_len + self.passage_max_len, pad_to_multiple_of=self.pad_to_multiple_of, return_tensors=self.return_tensors, ) return { "pair": collated, "teacher_scores": teacher_scores, } class AbsLLMRerankerTrainDataset(AbsRerankerTrainDataset): """Abstract class for LLM reranker training dataset. Args: args (AbsRerankerDataArguments): Data arguments. tokenizer (PreTrainedTokenizer): Tokenizer to use. """ def __init__( self, args: AbsRerankerDataArguments, tokenizer: PreTrainedTokenizer ): super().__init__(args, tokenizer) sep = self.args.sep_token self.sep_inputs = self.tokenizer( sep, return_tensors=None, add_special_tokens=False )['input_ids'] def __getitem__(self, item) -> List[BatchEncoding]: data = self.dataset[item] train_group_size = self.args.train_group_size query = data['query'] if self.args.query_instruction_for_rerank is not None: query = self.args.query_instruction_format.format( data['query_prompt'] if 'query_prompt' in data else self.args.query_instruction_for_rerank, 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_rerank is not None: passages = [ self.args.passage_instruction_format.format( data['passage_prompt'] if 'passage_prompt' in data else self.args.passage_instruction_for_rerank, p ) for p in passages ] prompt = self.dataset[item].get('prompt', "Given a query A and a passage B, determine whether the passage contains an answer to the query by providing a prediction of either 'Yes' or 'No'.") query_inputs = self.tokenizer( query, return_tensors=None, max_length=self.args.query_max_len + self.args.passage_max_len // 4, truncation=True, add_special_tokens=False ) prompt_inputs = self.tokenizer( prompt, return_tensors=None, add_special_tokens=False )['input_ids'] max_length = self.max_length - len(prompt_inputs) - len(self.sep_inputs) passages_inputs = [] for i, passage in enumerate(passages): passage_inputs = self.tokenizer( passage, return_tensors=None, max_length=self.args.passage_max_len + self.args.query_max_len // 2, truncation=True, add_special_tokens=False ) if self.tokenizer.bos_token_id is not None and self.tokenizer.bos_token_id != self.tokenizer.pad_token_id: item = self.tokenizer.prepare_for_model( [self.tokenizer.bos_token_id] + query_inputs['input_ids'], self.sep_inputs + passage_inputs['input_ids'], truncation='only_second', max_length=max_length, padding=False, return_attention_mask=False, return_token_type_ids=False, add_special_tokens=False ) else: item = self.tokenizer.prepare_for_model( query_inputs['input_ids'], self.sep_inputs + passage_inputs['input_ids'], truncation='only_second', max_length=max_length, padding=False, return_attention_mask=False, return_token_type_ids=False, add_special_tokens=False ) passage_inputs['input_ids'] = item['input_ids'] + self.sep_inputs + prompt_inputs passage_inputs['attention_mask'] = [1] * len(passage_inputs['input_ids']) # passage_inputs['labels'] = passage_inputs['input_ids'].copy() # passage_inputs['labels'] = [-100] * (len(passage_inputs['input_ids']) - 1) + passage_inputs['labels'][(len(passage_inputs['input_ids']) - 1):] passage_inputs.pop('token_type_ids') if 'token_type_ids' in passage_inputs.keys() else None if 'position_ids' in passage_inputs.keys(): passage_inputs['position_ids'] = list(range(len(passage_inputs['input_ids']))) passages_inputs.append(passage_inputs) return passages_inputs, teacher_scores @dataclass class AbsLLMRerankerCollator(DataCollatorForSeq2Seq): """ Wrapper that does conversion from List[Tuple[encode_qry, encode_psg]] to List[qry], List[psg] and pass batch separately to the actual collator. Abstract out data detail for the model. """ query_max_len: int = 32 passage_max_len: int = 128 def __call__(self, features, return_tensors='pt'): if return_tensors is None: return_tensors = self.return_tensors teacher_scores = [f[1] for f in features] if teacher_scores[0] is None: teacher_scores = None elif isinstance(teacher_scores[0], list): teacher_scores = sum(teacher_scores, []) features = [f[0] for f in features] if isinstance(features[0], list): features = sum(features, []) labels = [feature["labels"] for feature in features] if "labels" in features[0].keys() else None # We have to pad the labels before calling `tokenizer.pad` as this method won't pad them and needs them of the # same length to return tensors. if labels is not None: max_label_length = max(len(l) for l in labels) # print(max_label_length) if self.pad_to_multiple_of is not None: max_label_length = ( (max_label_length + self.pad_to_multiple_of - 1) // self.pad_to_multiple_of * self.pad_to_multiple_of ) padding_side = self.tokenizer.padding_side for feature in features: remainder = [self.label_pad_token_id] * (max_label_length - len(feature["labels"])) if isinstance(feature["labels"], list): feature["labels"] = ( feature["labels"] + remainder if padding_side == "right" else remainder + feature["labels"] ) elif padding_side == "right": feature["labels"] = np.concatenate([feature["labels"], remainder]).astype(np.int64) else: feature["labels"] = np.concatenate([remainder, feature["labels"]]).astype(np.int64) collated = self.tokenizer.pad( features, padding=self.padding, max_length=self.query_max_len + self.passage_max_len, return_tensors=return_tensors, pad_to_multiple_of=self.pad_to_multiple_of, ) return { "pair": collated, "teacher_scores": teacher_scores, }