Update modeling_metalatte.py

modeling_metalatte.py

@@ -14,7 +14,7 @@ import gc
 from torch.optim.lr_scheduler import _LRScheduler
 from transformers import EsmModel, PreTrainedModel
 from configuration import MetaLATTEConfig
-
+from urllib.parse import urljoin
 seed_everything(42)
    
 class GELU(nn.Module):
@@ -218,6 +218,26 @@ class MultitaskProteinModel(PreTrainedModel):
 
         # Initialize weights and apply final processing
         self.post_init()
+
+    @classmethod
+    def from_pretrained(cls, pretrained_model_name_or_path, *model_args, **kwargs):
+        config = kwargs.pop("config", None)
+        if config is None:
+            config = MetaLATTEConfig.from_pretrained(pretrained_model_name_or_path)
+        
+        model = cls(config)        
+        #state_dict = torch.load(f"{pretrained_model_name_or_path}/pytorch_model.bin", map_location=torch.device('cpu'))['state_dict']
+        try:
+            state_dict_url = urljoin(f"https://huggingface.co/{pretrained_model_name_or_path}/resolve/main/", "pytorch_model.bin")
+            state_dict = torch.hub.load_state_dict_from_url(
+                state_dict_url,
+                map_location=torch.device('cpu')
+            )['state_dict']
+            model.load_state_dict(state_dict, strict=False)
+        except Exception as e:
+            raise RuntimeError(f"Error loading state_dict from {pretrained_model_name_or_path}/pytorch_model.bin: {e}")
+
+        return model
         
     def forward(self, input_ids, attention_mask=None):
         outputs = self.esm_model(input_ids=input_ids, attention_mask=attention_mask, output_hidden_states=True)