abs ft reranker

Author: ZiyiXia

FlagEmbedding/abc/finetune/embedder/AbsTrainer.py

@@ -27,7 +27,7 @@ def compute_loss(self, model, inputs, return_outputs=False, **kwargs):
                 returns only the loss.
         
         Returns:
-            Union[torch.Tensor, tuple(torch.Tensor, ModelOutput)]: The computed loss. If ``return_outputs`` is ``True``, 
+            Union[torch.Tensor, tuple(torch.Tensor, EmbedderOutput)]: The computed loss. If ``return_outputs`` is ``True``, 
                 also returns the model's outputs in a tuple ``(loss, outputs)``.
         """
 

FlagEmbedding/abc/finetune/reranker/AbsArguments.py

@@ -8,7 +8,7 @@
 @dataclass
 class AbsRerankerModelArguments:
     """
-    Abstract class for model arguments.
+    Abstract class for reranker model arguments.
     """
 
     model_name_or_path: str = field(
@@ -46,6 +46,9 @@ class AbsRerankerModelArguments:
 
 @dataclass
 class AbsRerankerDataArguments:
+    """
+    Abstract class for reranker data arguments.
+    """
     train_data: str = field(
         default=None, metadata={
             "help": "One or more paths to training data. `query: str`, `pos: List[str]`, `neg: List[str]` are required in the training data.",

FlagEmbedding/abc/finetune/reranker/AbsDataset.py

@@ -21,6 +21,12 @@
 
 
 class AbsRerankerTrainDataset(Dataset):
+    """Abstract class for reranker training dataset.
+
+    Args:
+        args (AbsRerankerDataArguments): Data arguments.
+        tokenizer (PreTrainedTokenizer): Tokenizer to use.
+    """
     def __init__(
         self,
         args: AbsRerankerDataArguments,
@@ -47,6 +53,17 @@ def __init__(
         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.
+        """
         if dist.get_rank() == 0:
             logger.info(f'loading data from {file_path} ...')
 
@@ -64,6 +81,14 @@ def _load_dataset(self, file_path: str):
         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
@@ -78,6 +103,15 @@ 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(
@@ -143,6 +177,9 @@ def __getitem__(self, item):
 
 @dataclass
 class AbsRerankerCollator(DataCollatorWithPadding):
+    """
+    The abstract reranker collator.
+    """
     query_max_len: int = 32
     passage_max_len: int = 128
 
@@ -171,6 +208,12 @@ def __call__(self, features) -> list[BatchEncoding]:
         }
 
 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,

FlagEmbedding/abc/finetune/reranker/AbsModeling.py

@@ -18,6 +18,13 @@ class RerankerOutput(ModelOutput):
 
 
 class AbsRerankerModel(ABC, nn.Module):
+    """Abstract class of embedding model for training.
+
+    Args:
+        base_model: The base model to train on.
+        tokenizer (AutoTokenizer, optional): The tokenizer to use. Defaults to ``None``.
+        train_batch_size (int, optional): Batch size used for training. Defaults to ``4``.
+    """
     def __init__(
         self,
         base_model: None,
@@ -38,16 +45,36 @@ def __init__(
         self.yes_loc = self.tokenizer('Yes', add_special_tokens=False)['input_ids'][-1]
 
     def gradient_checkpointing_enable(self, **kwargs):
+        """
+        Activates gradient checkpointing for the current model.
+        """
         self.model.gradient_checkpointing_enable(**kwargs)
 
     def enable_input_require_grads(self, **kwargs):
+        """
+        Enables the gradients for the input embeddings.
+        """
         self.model.enable_input_require_grads(**kwargs)
 
     @abstractmethod
     def encode(self, features):
+        """Abstract method of encode.
+
+        Args:
+            features (dict): Teatures to pass to the model.
+        """
         pass
 
     def forward(self, pair: Union[Dict[str, Tensor], List[Dict[str, Tensor]]] = None, teacher_scores: Optional[Tensor] = None):
+        """The computation performed at every call.
+
+        Args:
+            pair (Union[Dict[str, Tensor], List[Dict[str, Tensor]]], optional): The query-document pair. Defaults to ``None``.
+            teacher_scores (Optional[Tensor], optional): Teacher scores of knowledge distillation. Defaults to None.
+
+        Returns:
+            RerankerOutput: Output of reranker model.
+        """
         ranker_logits = self.encode(pair) # (batch_size * num, dim)
         if teacher_scores is not None:
             teacher_scores = torch.Tensor(teacher_scores)
@@ -72,9 +99,23 @@ def forward(self, pair: Union[Dict[str, Tensor], List[Dict[str, Tensor]]] = None
         )
 
     def compute_loss(self, scores, target):
+        """Compute the loss.
+
+        Args:
+            scores (torch.Tensor): Computed scores.
+            target (torch.Tensor): The target value.
+
+        Returns:
+            torch.Tensor: The computed loss.
+        """
         return self.cross_entropy(scores, target)
 
     def save(self, output_dir: str):
+        """Save the model.
+
+        Args:
+            output_dir (str): Directory for saving the model.
+        """
         # self.model.save_pretrained(output_dir)
         state_dict = self.model.state_dict()
         state_dict = type(state_dict)(
@@ -84,5 +125,8 @@ def save(self, output_dir: str):
         self.model.save_pretrained(output_dir, state_dict=state_dict)
 
     def save_pretrained(self, *args, **kwargs):
+        """
+        Save the tokenizer and model.
+        """
         self.tokenizer.save_pretrained(*args, **kwargs)
         return self.model.save_pretrained(*args, **kwargs)

FlagEmbedding/abc/finetune/reranker/AbsRunner.py

@@ -22,6 +22,13 @@
 
 
 class AbsRerankerRunner(ABC):
+    """Abstract class to run reranker model fine-tuning.
+
+    Args:
+        model_args (AbsRerankerModelArguments): Model arguments
+        data_args (AbsRerankerDataArguments): Data arguments.
+        training_args (AbsRerankerTrainingArguments): Training arguments.
+    """
     def __init__(
         self,
         model_args: AbsRerankerModelArguments,
@@ -70,13 +77,28 @@ def __init__(
 
     @abstractmethod
     def load_tokenizer_and_model(self) -> Tuple[PreTrainedTokenizer, AbsRerankerModel]:
+        """Abstract method to load the tokenizer and model.
+
+        Returns:
+            Tuple[PreTrainedTokenizer, AbsRerankerModel]: Loaded tokenizer and model instances.
+        """
         pass
 
     @abstractmethod
     def load_trainer(self) -> AbsRerankerTrainer:
+        """Abstract method to load the trainer.
+
+        Returns:
+            AbsRerankerTrainer: The loaded trainer instance.
+        """
         pass
 
     def load_train_dataset(self) -> AbsRerankerTrainDataset:
+        """Loads the training dataset based on data arguments.
+
+        Returns:
+            AbsRerankerTrainDataset: The loaded dataset instance.
+        """
         if self.model_args.model_type == 'encoder':
             train_dataset = AbsRerankerTrainDataset(
                 args=self.data_args,
@@ -90,6 +112,11 @@ def load_train_dataset(self) -> AbsRerankerTrainDataset:
         return train_dataset
 
     def load_data_collator(self) -> AbsRerankerCollator:
+        """Loads the appropriate data collator.
+
+        Returns:
+            AbsRerankerCollator: Loaded data collator.
+        """
         if self.model_args.model_type == 'encoder':
             RerankerCollator = AbsRerankerCollator
         else:
@@ -106,6 +133,9 @@ def load_data_collator(self) -> AbsRerankerCollator:
         return data_collator
 
     def run(self):
+        """
+        Executes the training process.
+        """
         Path(self.training_args.output_dir).mkdir(parents=True, exist_ok=True)
 
         # Training

FlagEmbedding/abc/finetune/reranker/AbsTrainer.py

@@ -7,6 +7,9 @@
 
 
 class AbsRerankerTrainer(ABC, Trainer):
+    """
+    Abstract class for the trainer of reranker.
+    """
     @abstractmethod
     def _save(self, output_dir: Optional[str] = None, state_dict=None):
         pass
@@ -16,6 +19,16 @@ def compute_loss(self, model, inputs, return_outputs=False, **kwargs):
         How the loss is computed by Trainer. By default, all models return the loss in the first element.
 
         Subclass and override for custom behavior.
+        
+        Args:
+            model (AbsRerankerModel): The model being trained.
+            inputs (dict): A dictionary of input tensors to be passed to the model.
+            return_outputs (bool, optional): If ``True``, returns both the loss and the model's outputs. Otherwise,
+                returns only the loss. Defaults to ``False``.
+        
+        Returns:
+            Union[torch.Tensor, tuple(torch.Tensor, RerankerOutput)]: The computed loss. If ``return_outputs`` is ``True``, 
+                also returns the model's outputs in a tuple ``(loss, outputs)``.
         """
 
         outputs = model(**inputs)