ft base & m3

Author: ZiyiXia

FlagEmbedding/abc/finetune/embedder/AbsModeling.py

@@ -34,7 +34,7 @@ class AbsEmbedderModel(ABC, nn.Module):
         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): Knowledge distillation type. Defaults to ``"kl_div"``.
+        kd_loss_type (str, optional): Type of knowledge distillation loss. Defaults to ``"kl_div"``.
     """
     def __init__(
         self,

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

@@ -9,6 +9,19 @@
 
 
 class BiEncoderOnlyEmbedderModel(AbsEmbedderModel):
+    """Embedder class for encoder 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 '`cls`'.
+        normalize_embeddings (bool, optional): If True, normalize the embedding vector. Defaults to ``False``.
+    """
     TRANSFORMER_CLS = AutoModel
 
     def __init__(
@@ -35,6 +48,14 @@ 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 +91,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 +126,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/encoder_only/base/runner.py

@@ -13,7 +13,15 @@
 
 
 class EncoderOnlyEmbedderRunner(AbsEmbedderRunner):
+    """
+    Finetune Runner for base embedding models.
+    """
     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.model_name_or_path,
             cache_dir=self.model_args.cache_dir,
@@ -58,6 +66,11 @@ def load_tokenizer_and_model(self) -> Tuple[PreTrainedTokenizer, AbsEmbedderMode
         return tokenizer, model
 
     def load_trainer(self) -> EncoderOnlyEmbedderTrainer:
+        """Load the trainer.
+
+        Returns:
+            EncoderOnlyEmbedderTrainer: Loaded trainer instance.
+        """
         trainer = EncoderOnlyEmbedderTrainer(
             model=self.model,
             args=self.training_args,

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

@@ -9,7 +9,18 @@
 
 
 class EncoderOnlyEmbedderTrainer(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/encoder_only/m3/arguments.py

@@ -8,11 +8,17 @@
 
 @dataclass
 class EncoderOnlyEmbedderM3ModelArguments(AbsEmbedderModelArguments):
+    """
+    Model argument class for M3.
+    """
     colbert_dim: int = field(default=-1, metadata={"help": "Dim of colbert linear"})
 
 
 @dataclass
 class EncoderOnlyEmbedderM3TrainingArguments(AbsEmbedderTrainingArguments):
+    """
+    Training argument class for M3.
+    """
     unified_finetuning: bool = field(default=False, metadata={"help": "use unify fine-tuning"})
     use_self_distill: bool = field(default=False, metadata={"help": "use self-distill when using unify fine-tuning"})
     fix_encoder: bool = field(default=False, metadata={"help": "Freeze the parameters of encoder"})