ft embedder docstring

Author: ZiyiXia

FlagEmbedding/finetune/embedder/decoder_only/base/arguments.py

@@ -10,6 +10,9 @@ def default_target_modules() -> List[int]:
 
 @dataclass
 class DecoderOnlyEmbedderModelArguments(AbsEmbedderModelArguments):
+    """
+    Model argument class for decoder only base model.
+    """
     peft_model_path: str = field(
         default='', metadata={"help": "The peft model checkpoint for initialization."}
     )

FlagEmbedding/finetune/embedder/decoder_only/base/load_model.py

@@ -11,6 +11,14 @@
 
 
 def find_largest_checkpoint(checkpoint_dir):
+    """Find the largest checkpoint from directory.
+
+    Args:
+        checkpoint_dir (str): Directory to the checkpoint.
+
+    Returns:
+        str: Directory to the checkpoint, None no matching found.
+    """
     checkpoint_pattern = re.compile(r'checkpoint-(\d+)')
     max_number = -1
     max_checkpoint_file = None
@@ -28,6 +36,17 @@ def find_largest_checkpoint(checkpoint_dir):
 
 
 def get_model(model_args: DecoderOnlyEmbedderModelArguments, output_dir: str, resize: bool, resize_tokens: int):
+    """Get the model.
+
+    Args:
+        model_args (DecoderOnlyEmbedderModelArguments): Model arguments instance.
+        output_dir (str): Directory to save the model.
+        resize (bool): Whether to resize the number of tokens.
+        resize_tokens (int): The new token size.
+
+    Returns:
+        transformers.PreTrainedModel or PeftModel: The loaded model.
+    """
     if model_args.config_name:
         config = AutoConfig.from_pretrained(
             model_args.config_name,
@@ -99,6 +118,13 @@ def get_model(model_args: DecoderOnlyEmbedderModelArguments, output_dir: str, re
 
 
 def save_merged_model(model_args: DecoderOnlyEmbedderModelArguments, output_dir: str):
+    """
+    Loads a model with specified configurations, merges it with PEFT layers if available.
+
+    Args:
+        model_args (DecoderOnlyEmbedderModelArguments): Model arguments instance.
+        output_dir (str): Directory to save the model.
+    """
     if model_args.config_name:
         config = AutoConfig.from_pretrained(
             model_args.config_name,

FlagEmbedding/finetune/embedder/decoder_only/base/modeling.py

@@ -9,6 +9,19 @@
 
 
 class BiDecoderOnlyEmbedderModel(AbsEmbedderModel):
+    """Embedder model class for decoder only model.
+
+    Args:
+        base_model (AutoModel): The base model to train on.
+        tokenizer (AutoTokenizer, optional): The tokenizer to use. Defaults to ``None``.
+        negatives_cross_device (bool, optional): If True, will compute cross devices negative loss. Defaults to ``False``.
+        temperature (float, optional): Temperature to control the scale of scores. Defaults to ``1.0``.
+        sub_batch_size (int, optional): Sub-batch size during encoding. If negative, will not split to sub-batch.
+            Defaults to ``-1``.
+        kd_loss_type (str, optional): Type of knowledge distillation loss. Defaults to ``'kl_div'``.
+        sentence_pooling_method (str, optional): Pooling method to get sentence embedding. Defaults to ``'last_token'``.
+        normalize_embeddings (bool, optional): If True, normalize the embedding vector. Defaults to ``False``.
+    """
     TRANSFORMER_CLS = AutoModel
 
     def __init__(
@@ -35,6 +48,15 @@ def __init__(
         self.cross_entropy = torch.nn.CrossEntropyLoss(reduction='mean')
 
     def encode(self, features):
+        """
+        Encode and get the embedding.
+
+        Args:
+            features (Union[list, dict]): Features feed to the model.
+
+        Returns:
+            torch.Tensor: The embedding vectors.
+        """
         if features is None:
             return None
         if not isinstance(features, list):
@@ -70,6 +92,18 @@ def encode(self, features):
             return all_p_reps.contiguous()
 
     def _sentence_embedding(self, last_hidden_state, attention_mask):
+        """Use the pooling method to get the sentence embedding.
+
+        Args:
+            last_hidden_state (torch.Tensor): The model output's last hidden state.
+            attention_mask (torch.Tensor): Mask out padding tokens during pooling.
+
+        Raises:
+            NotImplementedError: Specified pooling method not implemented.
+
+        Returns:
+            torch.Tensor: The sentence embeddings.
+        """
         if self.sentence_pooling_method == "cls":
             return last_hidden_state[:, 0]
         elif self.sentence_pooling_method == "mean":
@@ -93,25 +127,63 @@ def _sentence_embedding(self, last_hidden_state, attention_mask):
             raise NotImplementedError(f"pooling method {self.sentence_pooling_method} not implemented")
 
     def compute_score(self, q_reps, p_reps):
+        """Computes the scores between query and passage representations.
+
+        Args:
+            q_reps (torch.Tensor): Query representations.
+            p_reps (torch.Tensor): Passage representations.
+
+        Returns:
+            torch.Tensor: The computed scores, adjusted by temperature.
+        """
         scores = self._compute_similarity(q_reps, p_reps) / self.temperature
         scores = scores.view(q_reps.size(0), -1)
         return scores
 
     def _compute_similarity(self, q_reps, p_reps):
+        """Computes the similarity between query and passage representations using inner product.
+
+        Args:
+            q_reps (torch.Tensor): Query representations.
+            p_reps (torch.Tensor): Passage representations.
+
+        Returns:
+            torch.Tensor: The computed similarity matrix.
+        """
         if len(p_reps.size()) == 2:
             return torch.matmul(q_reps, p_reps.transpose(0, 1))
         return torch.matmul(q_reps, p_reps.transpose(-2, -1))
 
     def compute_loss(self, scores, target):
+        """Compute the loss using cross entropy.
+
+        Args:
+            scores (torch.Tensor): Computed score.
+            target (torch.Tensor): The target value.
+
+        Returns:
+            torch.Tensor: The computed cross entropy loss.
+        """
         return self.cross_entropy(scores, target)
 
     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)
 
     def save(self, output_dir: str):
+        """Save the model to the directory.
+
+        Args:
+            output_dir (str): Directory for saving the model.
+        """
         state_dict = self.model.state_dict()
         state_dict = type(state_dict)(
             {k: v.clone().cpu()

FlagEmbedding/finetune/embedder/decoder_only/base/runner.py

@@ -15,6 +15,13 @@
 
 
 class DecoderOnlyEmbedderRunner(AbsEmbedderRunner):
+    """Runner class for decoder only embedding model.
+
+    Args:
+        model_args (DecoderOnlyEmbedderModelArguments): Model arguments instance.
+        data_args (AbsEmbedderDataArguments): Data arguments instance.
+        training_args (AbsEmbedderTrainingArguments): Trainer arguments.
+    """
     def __init__(
         self,
         model_args: DecoderOnlyEmbedderModelArguments,
@@ -24,6 +31,11 @@ def __init__(
         super().__init__(model_args, data_args, training_args)
 
     def load_tokenizer_and_model(self) -> Tuple[PreTrainedTokenizer, AbsEmbedderModel]:
+        """Load tokenizer and model.
+
+        Returns:
+            Tuple[PreTrainedTokenizer, AbsEmbedderModel]: Tokenizer and model instances.
+        """
         tokenizer = AutoTokenizer.from_pretrained(
             self.model_args.tokenizer_name if self.model_args.tokenizer_name else self.model_args.model_name_or_path,
             token=self.model_args.token,
@@ -83,6 +95,11 @@ def load_tokenizer_and_model(self) -> Tuple[PreTrainedTokenizer, AbsEmbedderMode
         return tokenizer, model
 
     def load_trainer(self) -> DecoderOnlyEmbedderTrainer:
+        """Load the trainer.
+
+        Returns:
+            EncoderOnlyEmbedderTrainer: Loaded trainer instance.
+        """
         trainer = DecoderOnlyEmbedderTrainer(
             model=self.model,
             args=self.training_args,
@@ -95,6 +112,9 @@ def load_trainer(self) -> DecoderOnlyEmbedderTrainer:
         return trainer
 
     def run(self):
+        """
+        Run the finetune.
+        """
         Path(self.training_args.output_dir).mkdir(parents=True, exist_ok=True)
 
         # Training

FlagEmbedding/finetune/embedder/decoder_only/base/trainer.py

@@ -9,7 +9,18 @@
 
 
 class DecoderOnlyEmbedderTrainer(AbsEmbedderTrainer):
+    """
+    Trainer class for base encoder models.
+    """
     def _save(self, output_dir: Optional[str] = None, state_dict=None):
+        """Save the model to directory.
+
+        Args:
+            output_dir (Optional[str], optional): Output directory to save the model. Defaults to ``None``.
+
+        Raises:
+            NotImplementedError
+        """
         output_dir = output_dir if output_dir is not None else self.args.output_dir
         os.makedirs(output_dir, exist_ok=True)
         logger.info("Saving model checkpoint to %s", output_dir)

FlagEmbedding/finetune/embedder/decoder_only/icl/arguments.py

@@ -13,6 +13,9 @@ def default_target_modules() -> List[int]:
 
 @dataclass
 class DecoderOnlyEmbedderICLModelArguments(AbsEmbedderModelArguments):
+    """
+    Model argument class for decoder only icl model.
+    """
     peft_model_path: str = field(
         default='', metadata={"help": "The peft model checkpoint for initialization."}
     )
@@ -73,6 +76,9 @@ class DecoderOnlyEmbedderICLModelArguments(AbsEmbedderModelArguments):
 
 @dataclass
 class DecoderOnlyEmbedderICLDataArguments(AbsEmbedderDataArguments):
+    """
+    Data argument class for decoder only icl model.
+    """
     example_query_max_len: int = field(
         default=64,
         metadata={"help": "The max length of example query."}

FlagEmbedding/finetune/embedder/decoder_only/icl/dataset.py

@@ -15,6 +15,16 @@
 
 
 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,
@@ -39,6 +49,16 @@ def __init__(
         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)
@@ -179,10 +199,12 @@ def _create_batch_data(self, batch_raw_data):
 @dataclass
 class AbsEmbedderSameDatasetCollator(DataCollatorWithPadding):
     """
-    EmbedCollator for SameDataset
+    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
+    
+    ``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

FlagEmbedding/finetune/embedder/decoder_only/icl/load_model.py

@@ -11,6 +11,14 @@
 
 
 def find_largest_checkpoint(checkpoint_dir):
+    """Find the largest checkpoint from directory.
+
+    Args:
+        checkpoint_dir (str): Directory to the checkpoint.
+
+    Returns:
+        str: Directory to the checkpoint, None no matching found.
+    """
     checkpoint_pattern = re.compile(r'checkpoint-(\d+)')
     max_number = -1
     max_checkpoint_file = None
@@ -28,6 +36,17 @@ def find_largest_checkpoint(checkpoint_dir):
 
 
 def get_model(model_args: DecoderOnlyEmbedderICLModelArguments, output_dir: str, resize: bool, resize_tokens: int):
+    """Get the model.
+
+    Args:
+        model_args (DecoderOnlyEmbedderModelArguments): Model arguments instance.
+        output_dir (str): Directory to save the model.
+        resize (bool): Whether to resize the number of tokens.
+        resize_tokens (int): The new token size.
+
+    Returns:
+        transformers.PreTrainedModel or PeftModel: The loaded model.
+    """
     if model_args.config_name:
         config = AutoConfig.from_pretrained(
             model_args.config_name,
@@ -99,6 +118,13 @@ def get_model(model_args: DecoderOnlyEmbedderICLModelArguments, output_dir: str,
 
 
 def save_merged_model(model_args: DecoderOnlyEmbedderICLModelArguments, output_dir: str):
+    """
+    Loads a model with specified configurations, merges it with PEFT layers if available.
+
+    Args:
+        model_args (DecoderOnlyEmbedderModelArguments): Model arguments instance.
+        output_dir (str): Directory to save the model.
+    """
     if model_args.config_name:
         config = AutoConfig.from_pretrained(
             model_args.config_name,