Upload model

modelling_cxrmate_ed.py

@@ -143,13 +143,19 @@ class CXRMateEDModel(transformers.LlavaForConditionalGeneration):
 
         # assert isinstance(self.config.time_delta_monotonic_inversion, bool)
 
-        with open(os.path.join(os.path.dirname(os.path.abspath(__file__)), 'tables.json'), 'r') as f:
+        path = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'tables.json')
+        path = path if os.path.exists(path) else hf_hub_download(repo_id='aehrc/cxrmate-ed', filename='tables.json')
+        with open(path, 'r') as f:
             self.tables = json.load(f)
         
-        with open(os.path.join(os.path.dirname(os.path.abspath(__file__)), 'lookup_tables.json'), 'r') as f:
+        path = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'lookup_tables.json')
+        path = path if os.path.exists(path) else hf_hub_download(repo_id='aehrc/cxrmate-ed', filename='lookup_tables.json')
+        with open(path, 'r') as f:
             self.luts = json.load(f)
-            
-        with open(os.path.join(os.path.dirname(os.path.abspath(__file__)), 'token_type_ids.json'), 'r') as f:
+        
+        path = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'token_type_ids.json')
+        path = path if os.path.exists(path) else hf_hub_download(repo_id='aehrc/cxrmate-ed', filename='token_type_ids.json')
+        with open(path, 'r') as f:
             self.token_type_to_token_type_id = json.load(f)
 
         self.tables = {k: self.tables[k] for k in self.config.tables_filter}
@@ -183,21 +189,6 @@ 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(
     #     cls,
@@ -1134,7 +1125,10 @@ class CXRMateEDModel(transformers.LlavaForConditionalGeneration):
 
         # Train set:
         if not test_set_only:
-            with open(os.path.join(os.path.dirname(os.path.abspath(__file__)), f'{study_id_split}_train_study_ids.json'), 'r') as f:
+            
+            path = os.path.join(os.path.dirname(os.path.abspath(__file__)), f'{study_id_split}_train_study_ids.json')
+            path = path if os.path.exists(path) else hf_hub_download(repo_id='aehrc/cxrmate-ed', filename=f'{study_id_split}_train_study_ids.json')
+            with open(path, 'r') as f:
                 study_ids = json.load(f)
             train_set = dataset['train']
             train_set_study_ids = train_set['study_id']
@@ -1149,7 +1143,9 @@ class CXRMateEDModel(transformers.LlavaForConditionalGeneration):
 
         # Validation set:
         if not test_set_only:
-            with open(os.path.join(os.path.dirname(os.path.abspath(__file__)), f'{study_id_split}_validate_study_ids.json'), 'r') as f:
+            path = os.path.join(os.path.dirname(os.path.abspath(__file__)), f'{study_id_split}_validate_study_ids.json')
+            path = path if os.path.exists(path) else hf_hub_download(repo_id='aehrc/cxrmate-ed', filename=f'{study_id_split}_validate_study_ids.json')
+            with open(path, 'r') as f:
                 study_ids = json.load(f)
             val_set = dataset['validate']
             val_set_study_ids = val_set['study_id']
@@ -1163,7 +1159,9 @@ class CXRMateEDModel(transformers.LlavaForConditionalGeneration):
             val_set = None
 
         # Test set:
-        with open(os.path.join(os.path.dirname(os.path.abspath(__file__)), f'{study_id_split}_test_study_ids.json'), 'r') as f:
+        path = os.path.join(os.path.dirname(os.path.abspath(__file__)), f'{study_id_split}_test_study_ids.json')
+        path = path if os.path.exists(path) else hf_hub_download(repo_id='aehrc/cxrmate-ed', filename=f'{study_id_split}_test_study_ids.json')
+        with open(path, 'r') as f:
             study_ids = json.load(f)
         test_set = dataset['test']
         test_set_study_ids = test_set['study_id']
@@ -1216,7 +1214,9 @@ class CXRMateEDModel(transformers.LlavaForConditionalGeneration):
         dataset = datasets.load_from_disk(dataset_path)
 
         # Train set:
-        with open(os.path.join(os.path.dirname(os.path.abspath(__file__)), f'mimic_cxr_jpg_train_study_ids.json'), 'r') as f:
+        path = os.path.join(os.path.dirname(os.path.abspath(__file__)), f'mimic_cxr_jpg_train_study_ids.json')
+        path = path if os.path.exists(path) else hf_hub_download(repo_id='aehrc/cxrmate-ed', filename='mimic_cxr_jpg_train_study_ids.json')
+        with open(path, 'r') as f:
             study_ids = json.load(f)
         train_set = dataset['train']
         train_set_study_ids = train_set['study_id']
@@ -1228,7 +1228,9 @@ class CXRMateEDModel(transformers.LlavaForConditionalGeneration):
         train_set = Subset(train_set, indices)
 
         # Validation set:
-        with open(os.path.join(os.path.dirname(os.path.abspath(__file__)), f'mimic_cxr_jpg_validate_study_ids.json'), 'r') as f:
+        path = os.path.join(os.path.dirname(os.path.abspath(__file__)), f'mimic_cxr_jpg_validate_study_ids.json')
+        path = path if os.path.exists(path) else hf_hub_download(repo_id='aehrc/cxrmate-ed', filename='mimic_cxr_jpg_validate_study_ids.json')
+        with open(path, 'r') as f:
             study_ids = json.load(f)
         val_set = dataset['validate']
         val_set_study_ids = val_set['study_id']
@@ -1240,7 +1242,9 @@ class CXRMateEDModel(transformers.LlavaForConditionalGeneration):
         val_set = Subset(val_set, indices)
 
         # Test set:
-        with open(os.path.join(os.path.dirname(os.path.abspath(__file__)), f'mimic_cxr_jpg_test_study_ids.json'), 'r') as f:
+        path = os.path.join(os.path.dirname(os.path.abspath(__file__)), f'mimic_cxr_jpg_test_study_ids.json')
+        path = path if os.path.exists(path) else hf_hub_download(repo_id='aehrc/cxrmate-ed', filename='mimic_cxr_jpg_test_study_ids.json')
+        with open(path, 'r') as f:
             study_ids = json.load(f)
         test_set = dataset['test']
         test_set_study_ids = test_set['study_id']