Merge pull request #1200 from ZiyiXia/master

Docstring

Author: ZiyiXia

FlagEmbedding/abc/evaluation/data_loader.py

@@ -374,7 +374,7 @@ def _download_gz_file(self, download_url: str, save_dir: str):
             save_dir (str): Path to the directory to save the gzip file.
 
         Raises:
-            FileNotFoundError: _description_
+            FileNotFoundError
 
         Returns:
             str: The path to the file after unzip.

FlagEmbedding/abc/finetune/embedder/AbsArguments.py

@@ -38,6 +38,9 @@ class AbsEmbedderModelArguments:
 
 @dataclass
 class AbsEmbedderDataArguments:
+    """
+    Abstract class for 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/embedder/AbsDataset.py

@@ -21,6 +21,12 @@
 
 
 class AbsEmbedderTrainDataset(Dataset):
+    """Abstract class for training dataset.
+
+    Args:
+        args (AbsEmbedderDataArguments): Data arguments.
+        tokenizer (PreTrainedTokenizer): Tokenizer to use.
+    """
     def __init__(
         self,
         args: AbsEmbedderDataArguments,
@@ -46,6 +52,17 @@ def __init__(
         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.
+        """
         if dist.get_rank() == 0:
             logger.info(f'loading data from {file_path} ...')
 
@@ -63,6 +80,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.shuffle_ratio > 0 and len(text) > 100 and random.random() < self.shuffle_ratio:
             split_text = []
             chunk_size = len(text)//3 + 1
@@ -126,6 +151,9 @@ def __getitem__(self, item):
 
 @dataclass
 class AbsEmbedderCollator(DataCollatorWithPadding):
+    """
+    The abstract embedder collator.
+    """
     query_max_len: int = 32
     passage_max_len: int = 128
     sub_batch_size: int = -1
@@ -214,6 +242,16 @@ def __call__(self, features):
 
 
 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,
@@ -296,6 +334,14 @@ def __init__(
         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.
+        """
         if dist.get_rank() == 0:
             logger.info(f'loading data from {file_path} ...')
 
@@ -311,6 +357,15 @@ def _load_dataset(self, file_path: str):
 
     @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:
@@ -320,6 +375,9 @@ def _get_file_batch_size(temp_dataset: datasets.Dataset, default_batch_size: int
         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)
 
@@ -353,6 +411,15 @@ def __getitem__(self, _):
         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']:
@@ -362,6 +429,16 @@ def _get_train_group_size(self, batch_raw_data):
         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)
@@ -426,10 +503,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
@@ -516,6 +595,9 @@ def __call__(self, features):
 
 
 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
 

FlagEmbedding/abc/finetune/embedder/AbsModeling.py

@@ -15,13 +15,27 @@
 
 @dataclass
 class EmbedderOutput(ModelOutput):
+    """
+    Output information returned by the model.
+    """
     q_reps: Optional[Tensor] = None
     p_reps: Optional[Tensor] = None
     loss: Optional[Tensor] = None
     scores: Optional[Tensor] = None
 
 
 class AbsEmbedderModel(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``.
+        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"``.
+    """
     def __init__(
         self,
         base_model,
@@ -48,21 +62,53 @@ def __init__(
 
     @abstractmethod
     def encode(self, features):
+        """Abstract method encode and get the embedding.
+
+        Args:
+            features (Union[list, dict]): Features feed to the model.
+        """
         pass
 
     @abstractmethod
     def compute_loss(self, scores, target):
+        """Abstract method compute the loss.
+
+        Args:
+            scores (torch.Tensor): Computed score.
+            target (torch.Tensor): The target value.
+        """
         pass
 
     @abstractmethod
     def compute_score(self, q_reps, p_reps):
+        """Abstract method to compute the score.
+
+        Args:
+            q_reps (torch.Tensor): Queries representations.
+            p_reps (torch.Tensor): Passages rerpresentations.
+        """
         pass
 
     @abstractmethod
     def save(self, output_dir: str):
+        """Abstract method to save the model.
+
+        Args:
+            output_dir (str): Directory for saving the model.
+        """
         pass
 
     def get_local_score(self, q_reps, p_reps, all_scores):
+        """Get the local score of queries and passages.
+
+        Args:
+            q_reps (torch.Tensor): Queries representations.
+            p_reps (torch.Tensor): Passages rerpresentations.
+            all_scores (torch.Tensor): All the query-passage scores computed.
+
+        Returns:
+            torch.Tensor: Local scores to compute loss.
+        """
         group_size = p_reps.size(0) // q_reps.size(0)
         indices = torch.arange(0, q_reps.size(0), device=q_reps.device) * group_size
         specific_scores = []
@@ -73,6 +119,17 @@ def get_local_score(self, q_reps, p_reps, all_scores):
         return torch.stack(specific_scores, dim=1).view(q_reps.size(0), -1)
 
     def compute_local_score(self, q_reps, p_reps, compute_score_func=None, **kwargs):
+        """Compute the local score of queries and passages.
+
+        Args:
+            q_reps (torch.Tensor): Queries representations.
+            p_reps (torch.Tensor): Passages rerpresentations.
+            compute_score_func (function, optional): Function to compute score. Defaults to ``None``, which will use the
+                :meth:`self.compute_score`.
+
+        Returns:
+            torch.Tensor: Local scores to compute loss.
+        """
         if compute_score_func is None:
             all_scores = self.compute_score(q_reps, p_reps)
         else:
@@ -181,6 +238,17 @@ def forward(
         teacher_scores: Union[None, List[float]] = None,
         no_in_batch_neg_flag: bool = False,
     ):
+        """The computation performed at every call.
+
+        Args:
+            queries (Union[Dict[str, Tensor], List[Dict[str, Tensor]]], optional): Input queries. Defaults to ``None``.
+            passages (Union[Dict[str, Tensor], List[Dict[str, Tensor]]], optional): Input passages. Defaults to ``None``.
+            teacher_scores (Union[None, List[float]], optional): Teacher scores for distillation. Defaults to ``None``.
+            no_in_batch_neg_flag (bool, optional): If True, use no in-batch negatives and no cross-device negatives. Defaults to ``False``.
+
+        Returns:
+            EmbedderOutput: Output of the forward call of model.
+        """
         q_reps = self.encode(queries) # (batch_size, dim)
         p_reps = self.encode(passages) # (batch_size * group_size, dim)
 
@@ -210,6 +278,20 @@ def forward(
 
     @staticmethod
     def distill_loss(kd_loss_type, teacher_targets, student_scores, group_size=None):
+        """Compute the distillation loss.
+
+        Args:
+            kd_loss_type (str): Type of knowledge distillation loss, supports "kl_div" and "m3_kd_loss".
+            teacher_targets (torch.Tensor): Targets from the teacher model.
+            student_scores (torch.Tensor): Score of student model.
+            group_size (int, optional): Number of groups for . Defaults to ``None``.
+
+        Raises:
+            ValueError: Invalid kd_loss_type
+
+        Returns:
+            torch.Tensor: A scalar of computed distillation loss.
+        """
         if kd_loss_type == 'kl_div':
             # teacher_targets: (batch_size, group_size) / (world_size * batch_size, group_size)
             # student_scores: (batch_size, group_size) / (world_size * batch_size, group_size)
@@ -236,6 +318,15 @@ def distill_loss(kd_loss_type, teacher_targets, student_scores, group_size=None)
             raise ValueError(f"Invalid kd_loss_type: {kd_loss_type}")
 
     def _dist_gather_tensor(self, t: Optional[torch.Tensor]):
+        """Gather a tensor from all processes in a distributed setting.
+
+        Args:
+            t (Optional[torch.Tensor]): The input tensor to be gathered. If `None`, no gathering is performed.
+
+        Returns:
+            Union[torch.Tensor, None]: A concatenated tensor from all processes if ``t`` is not ``None``, 
+                otherwise returns ``None``.
+        """
         if t is None:
             return None
         t = t.contiguous()