Update modeling_matryoshka.py

modeling_matryoshka.py

@@ -90,17 +90,20 @@ class MatryoshkaBertForSentenceSimilarity(PreTrainedModel):
         self.base_model = AutoModel.from_pretrained(config.base_model_name)
         
         print("Loading PEFT model...")
-        self.peft_model = PeftModel.from_pretrained(self.base_model, config.name_or_path)
+        # Fix the name_or_path attribute access
+        repo_path = getattr(config, 'name_or_path', getattr(config, '_name_or_path', 'Abdalrahmankamel/matryoshka-arabert'))
+        self.peft_model = PeftModel.from_pretrained(self.base_model, repo_path)
         
         print("Creating wrapper...")
         self.matryoshka_wrapper = MatryoshkaWrapper(self.peft_model, dims=config.dims)
 
         # Load wrapper weights
         try:
-            wrapper_weights_path = hf_hub_download(repo_id=config.name_or_path, filename="matryoshka_wrapper.pt")
+            repo_path = getattr(config, 'name_or_path', getattr(config, '_name_or_path', 'Abdalrahmankamel/matryoshka-arabert'))
+            wrapper_weights_path = hf_hub_download(repo_id=repo_path, filename="matryoshka_wrapper.pt")
             state_dict = torch.load(wrapper_weights_path, map_location='cpu')
             self.matryoshka_wrapper.load_state_dict(state_dict, strict=False)
-            print(f"✓ Loaded matryoshka_wrapper.pt from {config.name_or_path}")
+            print(f"✓ Loaded matryoshka_wrapper.pt from {repo_path}")
         except Exception as e:
             print(f"⚠ Could not load matryoshka_wrapper.pt: {e}")
 
@@ -118,11 +121,6 @@ class MatryoshkaBertForSentenceSimilarity(PreTrainedModel):
     def get_embedding(self, text, tokenizer, dim="256"):
         return self.matryoshka_wrapper.get_embedding(text, tokenizer, dim)
 
-# Register the model and config
-AutoConfig.register("matryoshka-arabert", MatryoshkaConfig)
-AutoModel.register(MatryoshkaConfig, MatryoshkaBertForSentenceSimilarity)
-
-# Add helper function for easy loading
 def load_matryoshka_model(repo_id="Abdalrahmankamel/matryoshka-arabert"):
     """
     Helper function to load the model easily