docs: add comments

Signed-off-by: jupyterjazz <[email protected]>

modeling_lora.py

@@ -162,6 +162,16 @@ class LoRAParametrization(nn.Module):
         dropout_p: float,
         alpha: float,
     ):
+        """
+        Registering LoRA adapters to all embedding and linear layers.
+
+        Additionally, we implement a custom forward function for LoRA parametrization.
+        This function modifies the layer's forward pass to optionally use task-specific
+        parameters. When a `task_id` is provided, it employs a LoRA parametrization
+        to modify the original weights according to the specific task. This allows
+        the layer to adapt dynamically to different tasks at runtime. If no `task_id`
+        is specified, the layer uses its original weights.
+        """
         if isinstance(layer, nn.Linear):
             parametrize.register_parametrization(
                 layer,
@@ -312,11 +322,11 @@ class XLMRobertaLoRA(XLMRobertaPreTrainedModel):
         config = XLMRobertaFlashConfig.from_pretrained(
             pretrained_model_name_or_path, *model_args, **kwargs
         )
-        if config.load_trained_adapters:
+        if config.load_trained_adapters: # checkpoint already contains LoRA adapters
             return super().from_pretrained(
                 pretrained_model_name_or_path, *model_args, **kwargs
             )
-        else:
+        else: # initializing new adapters
             roberta = XLMRobertaModel.from_pretrained(
                 pretrained_model_name_or_path, *model_args, **kwargs
             )
@@ -358,10 +368,12 @@ class XLMRobertaLoRA(XLMRobertaPreTrainedModel):
         **kwargs,
     ) -> Union[List[torch.Tensor], np.ndarray, torch.Tensor]:
         """
-        Computes sentence embeddings
+        Computes sentence embeddings.
 
+        sentences(`str` or `List[str]`):
+            Sentence or sentences to be encoded
         task_type(`str`, *optional*, defaults to `None`):
-            Specifies the task for which the encoding is intended. If `task_type` is not provide,
+            Specifies the task for which the encoding is intended. If `task_type` is not provided,
             all LoRA adapters are disabled, and the model reverts to its original,
             general-purpose weights.
         """