Fix import error and add source_files to config

config.json

@@ -7,6 +7,13 @@
   "architectures": [
     "FoundationBert"
   ],
+  "source_files": [
+      "foundation_bert.py",
+      "utils/__init__.py",
+      "utils/masked_data_modeling_loss.py",
+      "utils/yaml_util.py",
+      "train_config.yaml"
+    ],
   "attention_probs_dropout_prob": 0.1,
   "classifier_dropout": null,
   "hidden_act": "gelu",

foundation_bert.py

@@ -1,7 +1,13 @@
-import torch
+import sys
 import os
-import yaml
 from pathlib import Path
+
+current_dir = os.path.dirname(os.path.abspath(__file__))
+if current_dir not in sys.path:
+    sys.path.append(current_dir)
+
+import torch
+import yaml
 from utils.masked_data_modeling_loss import MaskedDataLossWithSoftmax
 # from ..utils.contrastive_loss import ContrastiveLoss
 from utils.yaml_util import MyLoader