Fix import error and add source_files to config

foundation_bert.py

@@ -143,17 +143,13 @@ class FoundationBert(BertModel):
     ):
         from huggingface_hub import hf_hub_download
         
-        # 1. 如果是远程加载，pretrained_model_name_or_path 就是 REPO_ID
-        # 我们显式地请求下载 train_config.yaml
         try:
-            # 这一步会检查缓存，如果没有则从云端下载并返回本地绝对路径
             model_config = hf_hub_download(
                 repo_id=pretrained_model_name_or_path,
                 filename="train_config.yaml",
                 revision=kwargs.get("revision", "main")
             )
         except Exception as e:
-            # 备选方案：如果本地路径已存在（例如 Snigdaa 的用法）
             model_config = os.path.join(pretrained_model_name_or_path, "train_config.yaml")
 
         # print(f"✅ Successfully located config at: {model_config}")

test_load.py

@@ -0,0 +1,45 @@
+import torch
+import sys
+from transformers import AutoModel, AutoConfig
+
+# Repository ID on Hugging Face Hub
+REPO_ID = "StarNetLaboratory/mosaic"
+
+# Since custom modeling code (foundation_bert.py) is used, 
+# trust_remote_code must be set to True.
+TRUST_CODE = True
+
+print(f"--- 1. Attempting to load configuration ---")
+try:
+    # Attempt to load config to verify config.json is present and readable
+    config = AutoConfig.from_pretrained(REPO_ID, trust_remote_code=TRUST_CODE)
+    print(f"✅ Config loaded successfully: {config.architectures}")
+except Exception as e:
+    print(f"❌ Config loading failed: {e}")
+    sys.exit(1)
+
+print(f"\n--- 2. Attempting to load model ---")
+try:
+    # Detect device (GPU if available, else CPU)
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    
+    # This call triggers transformers to download and execute foundation_bert.py
+    # and load weights from model.safetensors
+    model = AutoModel.from_pretrained(
+        REPO_ID, 
+        config=config, 
+        trust_remote_code=TRUST_CODE, 
+        torch_dtype=torch.float32 # Match the dtype used during training/local testing
+    ).to(device)
+    
+    model.eval()
+    
+    # Calculate and print total parameters to verify the architecture
+    total_params = sum(p.numel() for p in model.parameters())
+    print(f"✅ Model loaded successfully! Total parameters: {total_params:,}")
+
+except Exception as e:
+    print(f"❌ Model loading failed.")
+    print(f"Check file integrity and remote code logic (foundation_bert.py).")
+    print(f"Error details: {e}")
+    sys.exit(1)