Upload model

modelling_cxrmate_ed.py

@@ -7,6 +7,7 @@ from typing import Optional, Tuple, Union
 import datasets
 import torch
 import transformers
+from huggingface_hub import hf_hub_download
 from torch.nn import CrossEntropyLoss
 from torch.utils.data import Subset
 from torchvision.io import decode_image
@@ -182,6 +183,20 @@ class CXRMateEDModel(transformers.LlavaForConditionalGeneration):
         
         self.post_init()
 
+    @classmethod
+    def from_pretrained(cls, pretrained_model_name_or_path, **kwargs):
+
+        hf_hub_download(repo_id=pretrained_model_name_or_path, filename='tables.json')
+        hf_hub_download(repo_id=pretrained_model_name_or_path, filename='token_type_ids.json')
+        hf_hub_download(repo_id=pretrained_model_name_or_path, filename='lookup_tables.json')
+        hf_hub_download(repo_id=pretrained_model_name_or_path, filename='mimic_cxr_jpg_train_study_ids.json')
+        hf_hub_download(repo_id=pretrained_model_name_or_path, filename='mimic_cxr_jpg_validate_study_ids.json')
+        hf_hub_download(repo_id=pretrained_model_name_or_path, filename='mimic_cxr_jpg_test_study_ids.json')
+        hf_hub_download(repo_id=pretrained_model_name_or_path, filename='mimic_iv_ed_mimic_cxr_jpg_train_study_ids.json')
+        hf_hub_download(repo_id=pretrained_model_name_or_path, filename='mimic_iv_ed_mimic_cxr_jpg_validate_study_ids.json')
+        hf_hub_download(repo_id=pretrained_model_name_or_path, filename='mimic_iv_ed_mimic_cxr_jpg_test_study_ids.json')
+        
+        return super().from_pretrained(pretrained_model_name_or_path, **kwargs)
 
     # @classmethod
     # def from_encoder_decoder_pretrained(