FIX: [DEV-14438] ModernBert Accuracy new predict batch size default (#854)

* FIX: [DEV-14438] ModernBert Accuracy new predict batch size default

* FIX: being defensive

Author: benleetownsend

finetune/base.py

@@ -187,6 +187,19 @@ def resolve_config(self, no_fp16=False, **kwargs):
                 raise ValueError("There is no auto setting for {}".format(ak))
             config[ak] = overrides[ak]
 
+        override_predict_batch = overrides.get("predict_batch_size")
+        current_predict_batch = config.get("predict_batch_size")
+        if (
+            override_predict_batch is not None
+            and current_predict_batch is not None
+            and override_predict_batch < current_predict_batch
+            and hasattr(self, "config")
+        ):
+            LOGGER.info(
+                f"Overriding loaded predict batch size from {current_predict_batch} to the new default of {override_predict_batch}"
+            )
+            self.config["predict_batch_size"] = override_predict_batch
+
         if hasattr(self, "input_pipeline"):
             self.input_pipeline.config = config
 

finetune/base_models/modern_bert/model.py

@@ -73,7 +73,7 @@ def get_optimal_params(cls, config):
                 "n_epochs": base_n_epochs,
                 "batch_size": 8,
                 "chunk_context": None,
-                "predict_batch_size": 16,
+                "predict_batch_size": 8,
                 "mixed_precision": True,
                 "float_16_predict": True,
                 "lr": base_learning_rate,

tests/test_backwards_compat.py

@@ -1,5 +1,6 @@
 import glob
 import os
+from pathlib import Path
 
 import numpy as np
 import pytest
@@ -45,6 +46,20 @@ def test_roberta_default_no_change(get_untrained_sequence_labeler):
     assert model.config.collapse_whitespace == True
 
 
+def test_table_model_predict_batch_size_override(monkeypatch):
+    model_path = (
+        Path(__file__).parent
+        / "backwards_compat_bundles"
+        / "ModernBertModel_sequence_labeling_crf_True.jl"
+    )
+    model = SequenceLabeler.load(model_path, key="model")
+    try:
+        assert model.config.optimize_for.lower() == "accuracy"
+        assert model.config.predict_batch_size == 8
+    finally:
+        model.close()
+
+
 # TODO: eventually clean this up and push these files to s3.
 BUNDLES = glob.glob(os.path.join("/Finetune/tests/backwards_compat_bundles/*.jl"))