feat: Allow LoRA to be merged into weights

modeling_lora.py

@@ -199,6 +199,12 @@ class LoRAParametrization(nn.Module):
         if isinstance(layer, LoRAParametrization):
             layer.current_task = task_idx
 
+    @classmethod
+    def merge_lora_into_layer(cls, layer: nn.Module):
+        if hasattr(layer, "parametrizations"):
+            for attr_name in layer.parametrizations.keys():
+                parametrize.remove_parametrizations(layer, attr_name, leave_parametrized=True)
+
 
 class BertLoRA(BertPreTrainedModel):
     def __init__(self, config: JinaBertConfig, bert: Optional[BertModel] = None, add_pooling_layer=True):
@@ -207,6 +213,7 @@ class BertLoRA(BertPreTrainedModel):
             self.bert = BertModel(config, add_pooling_layer=add_pooling_layer)
         else:
             self.bert = bert
+        self._is_merged = False
         self._num_adaptions = config.num_loras
         self._register_lora(self._num_adaptions)
         self.main_params_trainable = False
@@ -230,6 +237,13 @@ class BertLoRA(BertPreTrainedModel):
         config = JinaBertConfig.from_pretrained(*args, **kwargs)
         return cls(config, bert=bert, num_adaptions=num_adaptions)
 
+    def merge_lora(self):
+        """Merges currently selected LoRA into main weights."""
+        if self._is_merged:
+            raise Exception('LoRA has already been merged, cannot merge again')
+        self._is_merged = True
+        self.apply(LoRAParametrization.merge_lora_into_layer)
+
     @classmethod
     def from_pretrained(
         cls,
@@ -265,6 +279,8 @@ class BertLoRA(BertPreTrainedModel):
 
     @current_task.setter
     def current_task(self, task_idx: Union[None, int]):
+        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:
             self._task_idx = task_idx