feat: make num of loras part of the config

configuration_bert.py

@@ -86,6 +86,7 @@ class JinaBertConfig(PretrainedConfig):
         use_qk_norm=True,
         emb_pooler=None,
         classifier_dropout=None,
+        num_loras=5,
         **kwargs,
     ):
         assert 'position_embedding_type' not in kwargs
@@ -118,4 +119,5 @@ class JinaBertConfig(PretrainedConfig):
         self.use_flash_attn = use_flash_attn
         self.use_qk_norm = use_qk_norm
         self.emb_pooler = emb_pooler
-        self.classifier_dropout = classifier_dropout
+        self.classifier_dropout = classifier_dropout
+        self.num_loras = num_loras

modeling_lora.py

@@ -201,14 +201,14 @@ class LoRAParametrization(nn.Module):
 
 
 class BertLoRA(BertPreTrainedModel):
-    def __init__(self, config: JinaBertConfig, bert: Optional[BertModel] = None, add_pooling_layer=True, num_adaptions=1):
+    def __init__(self, config: JinaBertConfig, bert: Optional[BertModel] = None, add_pooling_layer=True):
         super().__init__(config)
         if bert is None:
             self.bert = BertModel(config, add_pooling_layer=add_pooling_layer)
         else:
             self.bert = bert
-        self._num_adaptions = num_adaptions
-        self._register_lora(num_adaptions)
+        self._num_adaptions = config.num_loras
+        self._register_lora(self._num_adaptions)
         self.main_params_trainable = False
         self.current_task = 0