model

model.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):
@@ -226,9 +226,19 @@ class MultitaskProteinModel(PreTrainedModel):
             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']
-        model.load_state_dict(state_dict, strict=False)
+        #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)