refactor: restructure the class

Signed-off-by: jupyterjazz <saba.sturua@jina.ai>

modeling_lora.py

@@ -12,7 +12,6 @@ from transformers import PretrainedConfig
 from .modeling_xlm_roberta import (
     XLMRobertaFlashConfig,
     XLMRobertaModel,
-    XLMRobertaPreTrainedModel,
 )
 
 
@@ -209,19 +208,13 @@ class LoRAParametrization(nn.Module):
             layer.current_task = task_idx
 
 
-class XLMRobertaLoRA(XLMRobertaPreTrainedModel):
+class XLMRobertaLoRA(XLMRobertaModel):
     def __init__(
         self,
         config: XLMRobertaFlashConfig,
-        roberta: Optional[XLMRobertaModel] = None,
     ):
         super().__init__(config)
 
-        if roberta is None:
-            self.roberta = XLMRobertaModel(config)
-        else:
-            self.roberta = roberta
-
         self._num_adaptations = len(config.lora_adaptations)
         self._rank = config.lora_rank
         self._dropout_p = config.lora_dropout_p
@@ -238,6 +231,7 @@ class XLMRobertaLoRA(XLMRobertaPreTrainedModel):
         # By default, we select the first LoRA
         self.current_task = 0
 
+
     @property
     def main_params_trainable(self):
         return self._main_params_trainable
@@ -273,15 +267,14 @@ class XLMRobertaLoRA(XLMRobertaPreTrainedModel):
         config = XLMRobertaFlashConfig.from_pretrained(
             pretrained_model_name_or_path, *model_args, **kwargs
         )
+
         if config.load_trained_adapters:
             return super().from_pretrained(
                 pretrained_model_name_or_path, *model_args, **kwargs
             )
         else:
-            roberta = XLMRobertaModel.from_pretrained(
-                pretrained_model_name_or_path, *model_args, **kwargs
-            )
-            return cls(config, roberta=roberta)
+            torch.set_default_dtype(torch.float16)
+            return cls(config)
 
     def _register_lora(self, num_adaptations, rank, dropout_p, alpha):
         self.apply(
@@ -320,7 +313,7 @@ class XLMRobertaLoRA(XLMRobertaPreTrainedModel):
     def forward(self, *args, lora_adaptation: Union[None, int] = -1, **kwargs):
         if lora_adaptation is None or lora_adaptation >= 0:
             self.current_task = lora_adaptation
-        return self.roberta(*args, **kwargs)
+        return super().forward(*args, **kwargs)
 
     def parameters(self, recurse: bool = True) -> Iterator[Parameter]:
         for _, param in self.named_parameters(recurse=recurse):