feat: added docstrings

modeling_lora.py

@@ -65,6 +65,8 @@ class LoRAParametrization(nn.Module):
         fan_in_fan_out = layer_type == "embedding"
         self.swap = (lambda x: (x[1], x[0])) if fan_in_fan_out else (lambda x: x)
 
+        # For the officially "correct" LoRA initialization, check here: https://github.com/microsoft/LoRA
+        # TODO: Ensure that the initialization here is correct
         if layer_type == "linear":
             self.lora_A = nn.Parameter(
                 initialized_weights((rank, fan_in), num_adaptions, init="kaiming")
@@ -225,7 +227,15 @@ class BertLoRA(BertPreTrainedModel):
         return self._main_params_trainable
 
     @main_params_trainable.setter
-    def main_params_trainable(self, val):
+    def main_params_trainable(self, val: bool):
+        """Whether the main parameters (i.e. those that are not LoRA) should be trainable.
+
+        This method sets the `requires_grad_` attribute of the main weights
+        and controls which parameters are returned in `self.parameters()`.
+
+        :param val: Whether or not to make the parameters trainable.
+        :return: None
+        """
         self._main_params_trainable = val
         for name, param in super().named_parameters():
             if "lora" not in name:
@@ -259,7 +269,13 @@ class BertLoRA(BertPreTrainedModel):
         use_safetensors: bool = None,
         **kwargs,
     ):
-        # TODO: choose between from_bert and super().from_pretrained
+        """
+        TODO: choose between from_bert and super().from_pretrained
+
+        We want to be able to load both a pretrained BertModel, and a trained
+        BertLoRA via this method. To this end, we need to check which of these
+        models we are expected to load.
+        """
         return cls.from_bert(pretrained_model_name_or_path)
 
     def _register_lora(self, num_adaptions=1, rank=4, lora_dropout_p=0.0, lora_alpha=1):
@@ -275,14 +291,26 @@ class BertLoRA(BertPreTrainedModel):
 
     @property
     def current_task(self):
+        """ Which LoRA is currently selected
+        :return: Integer or None (when LoRA is disabled)
+        """
         return self._task_idx
 
     @current_task.setter
     def current_task(self, task_idx: Union[None, int]):
+        """Set the LoRA that is to be used.
+
+        The LoRA is specified by `task_idx`, which may be an integer >= 0,
+        indexing the available LoRAs. If it is None, no LoRA is used.
+
+        :param task_idx: Which LoRA to use
+        :return:
+        """
         if self._is_merged:
             raise Exception('LoRA has been merged, cannot select new task')
         assert task_idx is None or 0 <= task_idx < self._num_adaptions
         if self._task_idx != task_idx:
+            # In this case, we need to update the LoRAs everywhere
             self._task_idx = task_idx
             self.apply(
                 partial(LoRAParametrization.select_task_for_layer, task_idx=task_idx)