clarify classifier warning

automodel.py

@@ -137,13 +137,23 @@ class ClinicalMosaicForSequenceClassification(BertPreTrainedModel):
         
         missing_keys, unexpected_keys = model.load_state_dict(state_dict, strict=False)
 
-        if len(missing_keys) > 0:
+        # Calculate classifier parameters
+        num_classifier_params = config.hidden_size * config.num_labels + config.num_labels
+        classifier_keys = {"classifier.weight", "classifier.bias"}
+        
+        # Check if only the classification layer is missing
+        if set(missing_keys) == classifier_keys:
+            print(
+                f"Checkpoint does not contain the classification layer "
+                f"({config.hidden_size}x{config.num_labels} + {config.num_labels} = {num_classifier_params} params). "
+                "It will be randomly initialized."
+            )
+        elif len(missing_keys) > 0:
             logger.warning(
-                f"Found these missing keys in the checkpoint: {', '.join(missing_keys)}"
+                f"Checkpoint is missing {len(missing_keys)} parameters, including possibly critical ones: "
+                f"{', '.join(missing_keys)}"
             )
 
-            logger.warning(f"the number of which is equal to {len(missing_keys)}")
-
         if len(unexpected_keys) > 0:
             logger.warning(
                 f"Found these unexpected keys in the checkpoint: {', '.join(unexpected_keys)}",