healthcheck: fine-tune with a real warmup and 3 epochs

CHANGELOG.md

@@ -1,5 +1,15 @@
 # Changelog
 
+## 0.3.1
+
+`cacheverifier healthcheck` now actually fine-tunes. It passed no
+`warmup_steps` to `CrossEncoder.fit`, whose default of 10000 kept a typical
+run below a few percent of its learning rate, so `auc_tuned` came back equal
+(or nearly equal) to `auc_baseline`. Warmup is now 10% of total steps and the
+default is 3 epochs (`--epochs`), the same as the hosted service since
+verifier-core 0.2.0. On a 500-row later holdout this took LmArena AUC from
+0.903 to 0.932 and AmazonHelp from 0.551 to 0.700. Runs take about 3x as long.
+
 ## 0.3.0
 
 Make the request-path defaults safe. `verify()` and `verify_batch()` now use a

cacheverifier/_healthcheck/_finetune.py

@@ -43,6 +43,13 @@
 _MIN_QUERY_TOKENS = 32
 
 TRAIN_BATCH_SIZE = 4
+TRAIN_EPOCHS = 3
+TRAIN_WARMUP_FRACTION = 0.1
+"""Same as the hosted service (verifier-core 0.2.0). `CrossEncoder.fit`
+defaults to `warmup_steps=10000`; without an explicit value a typical run
+(a few hundred steps) never got past a few percent of its learning rate and
+the "fine-tuned" model was a near monotone shift of the base model, so
+`auc_tuned` came back equal to `auc_baseline`."""
 MIN_TRAIN_EXAMPLES = 20
 """Hard floor -- below this there aren't enough rows for one meaningful
 epoch plus a held-out split."""
@@ -183,7 +190,7 @@ def run_healthcheck(
     output_dir: str,
     *,
     base_model: str = DEFAULT_BASE_MODEL,
-    epochs: int = 1,
+    epochs: int = TRAIN_EPOCHS,
 ) -> HealthCheckResult:
     """Fine-tune `base_model` on a chronological prefix of `examples` and
     measure held-out AUC for the stock vs. fine-tuned model.
@@ -233,7 +240,8 @@ def run_healthcheck(
     # whether or not torch's stubs are installed (base CI has no torch).
     loader: DataLoader[Any] = DataLoader(cast(Any, train_examples), shuffle=True, batch_size=TRAIN_BATCH_SIZE)
     t0 = time.time()
-    tuned.fit(train_dataloader=loader, epochs=epochs, show_progress_bar=False)
+    warmup_steps = max(1, int(TRAIN_WARMUP_FRACTION * len(loader) * epochs))
+    tuned.fit(train_dataloader=loader, epochs=epochs, warmup_steps=warmup_steps, show_progress_bar=False)
     train_time_seconds = time.time() - t0
     tuned.save(output_dir)
 

cacheverifier/_healthcheck/cli.py

@@ -48,8 +48,8 @@ def add_subparser(subparsers: argparse._SubParsersAction) -> None:
     p.add_argument(
         "--epochs",
         type=int,
-        default=1,
-        help="fine-tuning epochs (default: 1, matching the hosted service)",
+        default=3,
+        help="fine-tuning epochs (default: 3, matching the hosted service)",
     )
     p.add_argument(
         "--keep-model",

pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
 
 [project]
 name = "cacheverifier"
-version = "0.3.0"
+version = "0.3.1"
 description = "Python client for the hosted CacheVerifier semantic-cache verification API"
 readme = "README.md"
 requires-python = ">=3.9"

tests/test_healthcheck.py

@@ -98,3 +98,28 @@ def test_run_healthcheck_end_to_end(tmp_path):
     assert result.n_train + result.n_calibrate + result.n_test == 36
     assert 0.0 <= result.auc_tuned <= 1.0
     assert (tmp_path / "model").is_dir()
+
+
+@pytest.mark.skipif(not HAS_TORCH, reason="needs the healthcheck extra (torch)")
+def test_run_healthcheck_warms_up_over_a_tenth_of_its_steps(tmp_path, monkeypatch):
+    # CrossEncoder.fit defaults to warmup_steps=10000, which kept a
+    # few-hundred-step run from ever reaching its learning rate.
+    from sentence_transformers import CrossEncoder
+
+    from cacheverifier._healthcheck import _finetune
+    from cacheverifier._healthcheck._finetune import GrayZoneExample, run_healthcheck
+
+    seen: dict = {}
+    real_fit = CrossEncoder.fit
+
+    def spy_fit(self, *args, **kwargs):
+        seen.update(kwargs, steps_per_epoch=len(kwargs["train_dataloader"]))
+        return real_fit(self, *args, **kwargs)
+
+    monkeypatch.setattr(CrossEncoder, "fit", spy_fit)
+    examples = [GrayZoneExample(f"question {i}", f"answer {i % 3}", i % 2 == 0) for i in range(40)]
+    run_healthcheck(examples, str(tmp_path / "model"))
+
+    assert seen["epochs"] == _finetune.TRAIN_EPOCHS == 3
+    expected = max(1, int(_finetune.TRAIN_WARMUP_FRACTION * seen["steps_per_epoch"] * seen["epochs"]))
+    assert seen["warmup_steps"] == expected