Fix remote train_config loading

foundation_bert.py

@@ -7,6 +7,7 @@ from pathlib import Path
 # from ..utils.yaml_util import MyLoader
 from dataclasses import dataclass
 from transformers import ModernBertModel, ModernBertConfig, PretrainedConfig
+from transformers.utils import cached_file
 from typing import Optional, Union
 
 # import yaml
@@ -176,11 +177,22 @@ class FoundationBert(ModernBertModel):
         """
             Modification to correctly handle loading extraneous parameters for GBert
         """
-        if 'checkpoint' in pretrained_model_name_or_path:
-            model_config = Path(pretrained_model_name_or_path).parent / 'train_config.yaml'
+        path = Path(pretrained_model_name_or_path)
+        if 'checkpoint' in str(pretrained_model_name_or_path):
+            model_config = path.parent / 'train_config.yaml'
+        elif path.is_dir():
+            model_config = path / 'train_config.yaml'
         else:
-            model_config = Path(pretrained_model_name_or_path) / 'train_config.yaml'
-            
+            model_config = cached_file(
+                pretrained_model_name_or_path,
+                'train_config.yaml',
+                cache_dir=cache_dir,
+                force_download=force_download,
+                local_files_only=local_files_only,
+                token=token,
+                revision=revision,
+            )
+        
         with open(model_config, 'r') as f:
             config = yaml.load(f, Loader=MyLoader)
         kwargs['modalities'] = config['modalities']