Upload 3 files

__init__.py

@@ -0,0 +1,19 @@
+from transformers import PretrainedConfig
+from typing import List
+
+class MetaLoRAConfig(PretrainedConfig):
+    model_type = "metalora"
+
+    def __init__(
+        self,
+        mlora_layers:List[int],
+        base_size:int,
+        embd_model:str,
+        llm_tokenizer:str,
+        **kwargs,
+    ):
+        self.mlora_layers = mlora_layers
+        self.base_size = base_size
+        self.embd_model = embd_model
+        self.llm_tokenizer = llm_tokenizer
+        super().__init__(**kwargs)

configuration_metalora.py

@@ -0,0 +1,20 @@
+from transformers import PretrainedConfig
+from typing import List
+
+class MetaLoRAConfig(PretrainedConfig):
+    model_type = "MLoRAModel"
+    
+    def __init__(
+        self,
+        mlora_layers: List[int]=["init"],
+        base_size:int=384,
+        embd_model:str="init",
+        llm_tokenizer:str="init",
+        **kwargs,
+    ):
+        self.auto_map = {"AutoModel": "Arthur-LAGACHERIE/MetaLoRA-code--modeling_metalora.MLoRAModel", "AutoConfig":"Arthur-LAGACHERIE/MetaLoRA-code--configuration_metalora.MetaLoRAConfig"}
+        self.mlora_layers = mlora_layers
+        self.base_size = base_size
+        self.embd_model = embd_model
+        self.llm_tokenizer = llm_tokenizer
+        super().__init__(**kwargs)

modeling_metalora.py

@@ -16,6 +16,7 @@ from transformers import Cache, DynamicCache, StaticCache
 from transformers.processing_utils import Unpack
 from transformers.modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast
 
+
 class MLoRASingleton:
     _instance = None  # Stocke l'instance unique
 
@@ -197,28 +198,27 @@ class MLoRAModel(nn.Module, GenerationMixin):
         self.singleton.set_tensor(torch.zeros(1, 1).to("cpu"))
 
     @classmethod
-    def from_pretrained(cls, repo, token=None):
+    def from_pretrained(cls, repo, token=None, print_=False):
         
         with tempfile.TemporaryDirectory() as tmp:
             # load the repo and the config
             snapshot_download(repo_id=repo, local_dir=f"{tmp}/repo/", token=token)
-            config = open(f"{tmp}/repo/config.json")
-            config = json.load(config)
+            config = MetaLoRAConfig.from_pretrained(repo, token=token)
             
             ### load model
             llm_config = AutoConfig.from_pretrained(f"{tmp}/repo/model/config.json")
             print("load LLM")
             llm = AutoModelForCausalLM.from_config(llm_config)
-            llm = get_mlora_model(llm, config["mlora_layers"], config["base_size"], lm_head="lm_head.weight" in config["mlora_layers"])
+            llm = get_mlora_model(llm, config.mlora_layers, config.base_size, lm_head="lm_head.weight" in config.mlora_layers, print_=print_)
             sfts.load_model(llm, f"{tmp}/repo/model/model.safetensors")
             print("LLM loaded")
 
             ### load tokenizer
-            tokenizer = AutoTokenizer.from_pretrained(config["llm_tokenizer"])
+            tokenizer = AutoTokenizer.from_pretrained(config.llm_tokenizer)
             
             ### load embd_model
             print("load Embd model")
-            embd_model = EmbdModel(AutoTokenizer.from_pretrained(config["embd_model"]), AutoModel.from_pretrained(config["embd_model"]))
+            embd_model = EmbdModel(AutoTokenizer.from_pretrained(config.embd_model), AutoModel.from_pretrained(config.embd_model))
             print("Embd model loaded")
 
             ### create instance
@@ -261,7 +261,7 @@ class MLoRAModel(nn.Module, GenerationMixin):
             ### push config
             config_path = saved_path / "config.json"
             with open(config_path, "w") as config_file:
-                json.dump(self.config, config_file)
+                json.dump(self.config.to_dict(), config_file)
 
             api.upload_file(
                 path_or_fileobj=config_path,