cacheverifier 0.3.1__tar.gz → 0.4.0__tar.gz

{cacheverifier-0.3.1 → cacheverifier-0.4.0}/CHANGELOG.md

@@ -1,5 +1,18 @@
 # Changelog
 
+## 0.4.0
+
+`cacheverifier healthcheck --base-model` picks the model the local Health
+Check starts from: `ms_marco` (default, unchanged), `nli`, `multilingual` or
+`multilingual_small` — the same choices as the hosted service's `base_model`
+— or any Hugging Face model id / local model directory. Chinese and other
+non-English traffic should use a multilingual base; the CLI now warns when
+most of the input is CJK text and the base is English-only. The NLI base's
+3-class head is scored `P(entailment) - P(contradiction)` for the stock
+baseline and re-initialized to a single logit for fine-tuning, matching the
+hosted service. `--emit-summary` records the base model (a local directory
+is recorded only as "local directory").
+
 ## 0.3.1
 
 `cacheverifier healthcheck` now actually fine-tunes. It passed no

{cacheverifier-0.3.1 → cacheverifier-0.4.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.5
 Name: cacheverifier
-Version: 0.3.1
+Version: 0.4.0
 Summary: Python client for the hosted CacheVerifier semantic-cache verification API
 Project-URL: Homepage, https://www.cacheverifier.com
 Project-URL: Documentation, https://www.cacheverifier.com/docs
@@ -177,6 +177,28 @@ Nothing is sent anywhere — the base model downloads once from Hugging Face, th
 fully offline. `--emit-summary` writes an aggregate-only JSON file (AUCs, counts, rates —
 no query or answer text) that's safe to share for a human read.
 
+**Choosing the base model.** `--base-model` takes the same values as the hosted
+service's `base_model` fine-tune parameter, so local and hosted numbers stay comparable:
+
+| `--base-model` | model | use for |
+|---|---|---|
+| `ms_marco` (default) | `cross-encoder/ms-marco-MiniLM-L6-v2` | English traffic |
+| `nli` | `cross-encoder/nli-MiniLM2-L6-H768` | English; entailment pretraining, stronger on negation / entity swaps |
+| `multilingual` | `BAAI/bge-reranker-base` | Chinese and other non-English traffic (~6 GB RAM to fine-tune) |
+| `multilingual_small` | `cross-encoder/mmarco-mMiniLMv2-L12-H384-v1` | same, smaller (~3 GB RAM) |
+
+The English bases read most Chinese characters as unknown tokens, so on Chinese
+traffic their result doesn't mean anything — the CLI warns when it sees mostly CJK text
+on an English base. You can also pass any Hugging Face model id or a local model
+directory. If Hugging Face is unreachable from your network, set `HF_ENDPOINT` to a
+mirror, or download the model elsewhere and pass its directory:
+
+```bash
+cacheverifier healthcheck traffic.jsonl --base-model multilingual
+HF_ENDPOINT=https://hf-mirror.com cacheverifier healthcheck traffic.jsonl --base-model multilingual_small
+cacheverifier healthcheck traffic.jsonl --base-model ./models/bge-reranker-base
+```
+
 ```
 results
 ------------------------------------------------------------------

{cacheverifier-0.3.1 → cacheverifier-0.4.0}/README.md

@@ -145,6 +145,28 @@ Nothing is sent anywhere — the base model downloads once from Hugging Face, th
 fully offline. `--emit-summary` writes an aggregate-only JSON file (AUCs, counts, rates —
 no query or answer text) that's safe to share for a human read.
 
+**Choosing the base model.** `--base-model` takes the same values as the hosted
+service's `base_model` fine-tune parameter, so local and hosted numbers stay comparable:
+
+| `--base-model` | model | use for |
+|---|---|---|
+| `ms_marco` (default) | `cross-encoder/ms-marco-MiniLM-L6-v2` | English traffic |
+| `nli` | `cross-encoder/nli-MiniLM2-L6-H768` | English; entailment pretraining, stronger on negation / entity swaps |
+| `multilingual` | `BAAI/bge-reranker-base` | Chinese and other non-English traffic (~6 GB RAM to fine-tune) |
+| `multilingual_small` | `cross-encoder/mmarco-mMiniLMv2-L12-H384-v1` | same, smaller (~3 GB RAM) |
+
+The English bases read most Chinese characters as unknown tokens, so on Chinese
+traffic their result doesn't mean anything — the CLI warns when it sees mostly CJK text
+on an English base. You can also pass any Hugging Face model id or a local model
+directory. If Hugging Face is unreachable from your network, set `HF_ENDPOINT` to a
+mirror, or download the model elsewhere and pass its directory:
+
+```bash
+cacheverifier healthcheck traffic.jsonl --base-model multilingual
+HF_ENDPOINT=https://hf-mirror.com cacheverifier healthcheck traffic.jsonl --base-model multilingual_small
+cacheverifier healthcheck traffic.jsonl --base-model ./models/bge-reranker-base
+```
+
 ```
 results
 ------------------------------------------------------------------

{cacheverifier-0.3.1 → cacheverifier-0.4.0}/cacheverifier/__init__.py

@@ -18,5 +18,5 @@ Research behind it: https://github.com/imxinchengyou/CacheVerifier
 
 from cacheverifier.client import CacheVerifier, CacheVerifierError, VerifyResult
 
-__version__ = "0.3.0"
+__version__ = "0.4.0"
 __all__ = ["CacheVerifier", "CacheVerifierError", "VerifyResult", "__version__"]

cacheverifier-0.4.0/cacheverifier/_healthcheck/_base_models.py

@@ -0,0 +1,62 @@
+"""Base-model choices for `cacheverifier healthcheck --base-model`.
+
+Kept free of torch / sentence-transformers so `cacheverifier healthcheck
+--help` can list the choices without the `healthcheck` extra installed.
+The aliases and model ids match the hosted service's `base_model` fine-tune
+parameter, so a local run is comparable to `POST /v1/finetune/dry-run`
+with the same value.
+"""
+
+from __future__ import annotations
+
+import re
+from pathlib import Path
+
+DEFAULT_BASE_MODEL = "cross-encoder/ms-marco-MiniLM-L6-v2"
+"""The stock off-the-shelf verifier every tenant starts on -- the same base
+model the hosted service fine-tunes from. English only."""
+
+BASE_MODEL_CHOICES: dict[str, str] = {
+    "ms_marco": DEFAULT_BASE_MODEL,
+    # SNLI+MultiNLI entailment pretraining, same size as ms_marco. Ships a
+    # 3-class head; see _finetune for how it is scored and fine-tuned.
+    "nli": "cross-encoder/nli-MiniLM2-L6-H768",
+    # Non-English (e.g. Chinese) traffic: the English bases turn most
+    # Chinese characters into [UNK] at tokenization.
+    "multilingual": "BAAI/bge-reranker-base",
+    "multilingual_small": "cross-encoder/mmarco-mMiniLMv2-L12-H384-v1",
+}
+
+ENGLISH_ONLY_BASE_MODELS = frozenset({BASE_MODEL_CHOICES["ms_marco"], BASE_MODEL_CHOICES["nli"]})
+
+_CJK_RE = re.compile(r"[぀-ヿ㐀-䶿一-鿿가-힯]")
+
+CJK_WARNING_RATIO = 0.3
+"""Share of non-whitespace characters that are CJK above which an
+English-only base model gets a warning."""
+
+
+def resolve_base_model(value: str) -> str:
+    """An alias from `BASE_MODEL_CHOICES`, a Hugging Face model id
+    ("org/name"), or a local model directory (e.g. one downloaded from a
+    mirror or ModelScope). Raises ValueError for anything else, so a typo'd
+    alias fails immediately instead of after a failed download."""
+    if value in BASE_MODEL_CHOICES:
+        return BASE_MODEL_CHOICES[value]
+    # `Path("")` is the current directory -- never treat an empty value as
+    # a local model.
+    if value.strip() and Path(value).expanduser().is_dir():
+        return str(Path(value).expanduser())
+    if "/" in value and not value.startswith("/"):
+        return value
+    raise ValueError(
+        f"unknown base model {value!r}: use one of {', '.join(BASE_MODEL_CHOICES)}, "
+        "a Hugging Face model id (org/name), or a local model directory"
+    )
+
+
+def cjk_ratio(texts: list[str]) -> float:
+    chars = "".join("".join(t.split()) for t in texts)
+    if not chars:
+        return 0.0
+    return len(_CJK_RE.findall(chars)) / len(chars)

{cacheverifier-0.3.1 → cacheverifier-0.4.0}/cacheverifier/_healthcheck/_finetune.py

@@ -24,15 +24,17 @@ import torch
 from sentence_transformers import CrossEncoder, InputExample
 from torch.utils.data import DataLoader
 
-# All four imports above are from the `healthcheck` extra. This module is
+from cacheverifier._healthcheck._base_models import DEFAULT_BASE_MODEL
+
+# The four third-party imports above are from the `healthcheck` extra. This module is
 # only ever imported from inside `cacheverifier._healthcheck.cli.run()`,
 # which catches the ImportError and points the user at
 # `pip install "cacheverifier[healthcheck]"` -- so a plain
 # `import cacheverifier` / `cacheverifier --help` never pays this cost.
 
-DEFAULT_BASE_MODEL = "cross-encoder/ms-marco-MiniLM-L6-v2"
-"""The stock off-the-shelf verifier every tenant starts on -- the same base
-model the hosted service fine-tunes from."""
+NLI_LABEL_ORDER = ("contradiction", "entailment", "neutral")
+"""`cross-encoder/nli-*` models' class order, used when the loaded model
+doesn't expose an `id2label` naming them."""
 
 MAX_SEQUENCE_LENGTH = 128
 """Token cap applied identically at train and score time. Matches the
@@ -185,6 +187,28 @@ def _diagnose_ceiling(scores: np.ndarray, labels: np.ndarray) -> str | None:
     return "still_improvable"
 
 
+def _nli_column_indices(model: Any) -> tuple[int, int]:
+    """(entailment_idx, contradiction_idx) for a 3-class NLI cross-encoder."""
+    id2label = getattr(getattr(model, "config", None), "id2label", None)
+    if isinstance(id2label, dict) and len(id2label) == 3:
+        lut = {str(v).lower(): int(k) for k, v in id2label.items()}
+        if "entailment" in lut and "contradiction" in lut:
+            return lut["entailment"], lut["contradiction"]
+    return NLI_LABEL_ORDER.index("entailment"), NLI_LABEL_ORDER.index("contradiction")
+
+
+def score_pairs(model: Any, pairs: list[tuple[str, str]], batch_size: int = 32) -> np.ndarray:
+    """Higher means "more likely an acceptable answer". A 1-logit model
+    (ms-marco, the multilingual rerankers, anything fine-tuned here) returns
+    its raw logit; a 3-class NLI model returns P(entailment) - P(contradiction)
+    -- the same scoring the hosted service uses for the stock NLI baseline."""
+    if int(model.config.num_labels) == 3:
+        ent, contra = _nli_column_indices(model)
+        probs = np.asarray(model.predict(pairs, batch_size=batch_size, show_progress_bar=False, apply_softmax=True))
+        return probs[:, ent] - probs[:, contra]
+    return np.asarray(model.predict(pairs, batch_size=batch_size, show_progress_bar=False))
+
+
 def run_healthcheck(
     examples: list[GrayZoneExample],
     output_dir: str,
@@ -221,11 +245,21 @@ def run_healthcheck(
 
     baseline = CrossEncoder(base_model, device=device, max_length=MAX_SEQUENCE_LENGTH)
     test_pairs_scored = [smart_truncate_pair(q, a, baseline.tokenizer) for q, a in test_pairs]
-    baseline_scores = np.array(baseline.predict(test_pairs_scored, batch_size=32, show_progress_bar=False))
+    baseline_scores = score_pairs(baseline, test_pairs_scored)
     auc_baseline = roc_auc(baseline_scores, test_labels)
     del baseline
 
-    tuned = CrossEncoder(base_model, device=device, max_length=MAX_SEQUENCE_LENGTH)
+    # num_labels=1 forces a single-logit head, matching the hosted service: a
+    # no-op for 1-logit bases, a deliberate head reinitialization (encoder
+    # body kept) for the 3-class NLI base. ignore_mismatched_sizes lets that
+    # 3->1 reinit happen instead of raising.
+    tuned = CrossEncoder(
+        base_model,
+        num_labels=1,
+        device=device,
+        max_length=MAX_SEQUENCE_LENGTH,
+        automodel_args={"ignore_mismatched_sizes": True},
+    )
     train_examples = [
         InputExample(
             texts=list(smart_truncate_pair(e.query, e.candidate_answer, tuned.tokenizer)),
@@ -246,10 +280,10 @@ def run_healthcheck(
     tuned.save(output_dir)
 
     calib_scored = [smart_truncate_pair(q, a, tuned.tokenizer) for q, a in calib_pairs]
-    calib_scores = np.array(tuned.predict(calib_scored, batch_size=32, show_progress_bar=False))
+    calib_scores = score_pairs(tuned, calib_scored)
     threshold = select_threshold(calib_scores, calib_labels)
 
-    tuned_scores = np.array(tuned.predict(test_pairs_scored, batch_size=32, show_progress_bar=False))
+    tuned_scores = score_pairs(tuned, test_pairs_scored)
     auc_tuned = roc_auc(tuned_scores, test_labels)
     ceiling_status = _diagnose_ceiling(tuned_scores, test_labels)
 
@@ -259,7 +293,7 @@ def run_healthcheck(
     if threshold is not None:
         threshold_hit_rate, threshold_error_rate = _hit_and_error_rate(tuned_scores, test_labels, threshold)
         train_pairs_scored = [smart_truncate_pair(e.query, e.candidate_answer, tuned.tokenizer) for e in train_rows]
-        train_scores = np.array(tuned.predict(train_pairs_scored, batch_size=32, show_progress_bar=False))
+        train_scores = score_pairs(tuned, train_pairs_scored)
         train_labels = np.array([1 if e.was_correct else 0 for e in train_rows])
         train_pred = (train_scores >= threshold).astype(int)
         train_label_disagreement_rate = float(np.mean(train_pred != train_labels))

{cacheverifier-0.3.1 → cacheverifier-0.4.0}/cacheverifier/_healthcheck/cli.py

@@ -13,6 +13,14 @@ import sys
 from pathlib import Path
 from typing import Any
 
+from cacheverifier._healthcheck._base_models import (
+    BASE_MODEL_CHOICES,
+    CJK_WARNING_RATIO,
+    ENGLISH_ONLY_BASE_MODELS,
+    cjk_ratio,
+    resolve_base_model,
+)
+
 _MISSING_EXTRA_HINT = (
     "the healthcheck extra is not installed -- run:\n"
     '    pip install "cacheverifier[healthcheck]"\n'
@@ -51,6 +59,17 @@ def add_subparser(subparsers: argparse._SubParsersAction) -> None:
         default=3,
         help="fine-tuning epochs (default: 3, matching the hosted service)",
     )
+    p.add_argument(
+        "--base-model",
+        default="ms_marco",
+        metavar="MODEL",
+        help=(
+            f"model to start from: {', '.join(BASE_MODEL_CHOICES)} (same choices as the hosted "
+            "service's base_model), a Hugging Face model id, or a local model directory. "
+            "Default ms_marco is English-only -- use multilingual or multilingual_small for "
+            "Chinese or other non-English traffic (they need several GB of RAM to fine-tune)"
+        ),
+    )
     p.add_argument(
         "--keep-model",
         action="store_true",
@@ -125,6 +144,12 @@ def run(args: argparse.Namespace) -> int:
         print(f"error: {e}", file=sys.stderr)
         return 2
 
+    try:
+        base_model = resolve_base_model(args.base_model)
+    except ValueError as e:
+        print(f"error: {e}", file=sys.stderr)
+        return 2
+
     n = len(examples)
     print(f"\nCacheVerifier local Health Check  --  {input_path}")
     print("=" * 66)
@@ -133,6 +158,17 @@ def run(args: argparse.Namespace) -> int:
     if n_stale:
         print(f"stale, excluded:      {n_stale}")
 
+    print(f"base model:           {base_model}")
+    if base_model in ENGLISH_ONLY_BASE_MODELS:
+        ratio = cjk_ratio([t for e in examples for t in (e.query, e.candidate_answer)])
+        if ratio > CJK_WARNING_RATIO:
+            print(
+                f"\nWarning: {ratio:.0%} of your text is Chinese/Japanese/Korean, but {base_model} is an "
+                "English-only model\n         (it reads most CJK characters as unknown tokens), so this "
+                "result won't reflect your\n         traffic. Re-run with --base-model multilingual "
+                "(or multilingual_small)."
+            )
+
     if n < _finetune.MIN_TRAIN_EXAMPLES:
         print(f"\nNeed at least {_finetune.MIN_TRAIN_EXAMPLES} non-stale rows to run; have {n}.", file=sys.stderr)
         return 2
@@ -149,7 +185,7 @@ def run(args: argparse.Namespace) -> int:
     print(f"\nfine-tuning locally (nothing sent) -> {out_dir}")
     print("this takes a few minutes on CPU; the base model downloads once on first run...\n")
 
-    result = _finetune.run_healthcheck(examples, out_dir, base_model=_finetune.DEFAULT_BASE_MODEL, epochs=args.epochs)
+    result = _finetune.run_healthcheck(examples, out_dir, base_model=base_model, epochs=args.epochs)
 
     print("results")
     print("-" * 66)
@@ -175,6 +211,9 @@ def run(args: argparse.Namespace) -> int:
 
     if args.emit_summary is not None:
         summary = {
+            # A local directory's path can carry a username -- this file is
+            # meant to be shareable, so record only that it was local.
+            "base_model": "local directory" if Path(base_model).is_dir() else base_model,
             "n_rows_read": len(rows),
             "n_usable": n,
             "n_stale_excluded": n_stale,

{cacheverifier-0.3.1 → cacheverifier-0.4.0}/pyproject.toml

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

{cacheverifier-0.3.1 → cacheverifier-0.4.0}/tests/test_healthcheck.py

@@ -123,3 +123,83 @@ def test_run_healthcheck_warms_up_over_a_tenth_of_its_steps(tmp_path, monkeypatc
     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
+
+
+class TestBaseModel:
+    def test_aliases_match_the_hosted_service(self):
+        from cacheverifier._healthcheck._base_models import BASE_MODEL_CHOICES, resolve_base_model
+
+        assert set(BASE_MODEL_CHOICES) == {"ms_marco", "nli", "multilingual", "multilingual_small"}
+        assert resolve_base_model("ms_marco") == "cross-encoder/ms-marco-MiniLM-L6-v2"
+        assert resolve_base_model("multilingual") == "BAAI/bge-reranker-base"
+
+    def test_hf_id_and_local_dir_pass_through(self, tmp_path):
+        from cacheverifier._healthcheck._base_models import resolve_base_model
+
+        assert resolve_base_model("someorg/some-reranker") == "someorg/some-reranker"
+        assert resolve_base_model(str(tmp_path)) == str(tmp_path)
+
+    @pytest.mark.parametrize("bad", ["multilingal", "/no/such/model/dir", ""])
+    def test_typos_and_missing_paths_are_rejected(self, bad):
+        from cacheverifier._healthcheck._base_models import resolve_base_model
+
+        with pytest.raises(ValueError, match="unknown base model"):
+            resolve_base_model(bad)
+
+    def test_cjk_ratio(self):
+        from cacheverifier._healthcheck._base_models import cjk_ratio
+
+        assert cjk_ratio(["how do I cancel"]) == 0.0
+        assert cjk_ratio(["怎么取消订阅"]) == 1.0
+        assert cjk_ratio(["怎么取消 subscription"]) == 4 / 16
+        assert cjk_ratio([]) == 0.0
+
+    def test_help_lists_base_model_choices(self, capsys):
+        with pytest.raises(SystemExit):
+            main(["healthcheck", "--help"])
+        out = capsys.readouterr().out
+        assert "--base-model" in out and "multilingual_small" in out
+
+    @pytest.mark.skipif(not HAS_TORCH, reason="run() imports the healthcheck extra first")
+    def test_unknown_base_model_is_an_input_error(self, capsys, tmp_path):
+        f = tmp_path / "t.jsonl"
+        f.write_text(json.dumps({"query": "q", "candidate_answer": "a", "was_correct": True}) + "\n")
+        assert main(["healthcheck", str(f), "--base-model", "multilingal"]) == 2
+        assert "unknown base model" in capsys.readouterr().err
+
+    @pytest.mark.skipif(not HAS_TORCH, reason="run() imports the healthcheck extra first")
+    def test_chinese_data_on_english_model_warns(self, capsys, tmp_path):
+        # Below MIN_TRAIN_EXAMPLES, so run() stops before any download/training.
+        f = tmp_path / "t.jsonl"
+        f.write_text(
+            "\n".join(
+                json.dumps({"query": "怎么取消自动续费", "candidate_answer": "进入设置取消", "was_correct": True}, ensure_ascii=False)
+                for _ in range(3)
+            )
+        )
+        assert main(["healthcheck", str(f)]) == 2
+        out = capsys.readouterr().out
+        assert "English-only" in out and "--base-model multilingual" in out
+
+        assert main(["healthcheck", str(f), "--base-model", "multilingual_small"]) == 2
+        assert "English-only" not in capsys.readouterr().out
+
+
+@pytest.mark.skipif(not HAS_TORCH, reason="needs the healthcheck extra (torch, sentence-transformers)")
+def test_run_healthcheck_nli_base_reinitializes_to_single_logit(tmp_path):
+    """The NLI base ships a 3-class head: the stock baseline is scored
+    P(entailment) - P(contradiction), and fine-tuning reinitializes it to the
+    same 1-logit head as every other base (matching the hosted service)."""
+    from sentence_transformers import CrossEncoder
+
+    from cacheverifier._healthcheck._base_models import BASE_MODEL_CHOICES
+    from cacheverifier._healthcheck._finetune import GrayZoneExample, run_healthcheck
+
+    examples = [
+        GrayZoneExample("how do I cancel my plan", "Settings > Billing > Cancel." if i % 2 == 0 else "Settings > Pause.", i % 2 == 0)
+        for i in range(24)
+    ]
+    result = run_healthcheck(examples, str(tmp_path / "model"), base_model=BASE_MODEL_CHOICES["nli"], epochs=1)
+    assert 0.0 <= result.auc_baseline <= 1.0
+    assert 0.0 <= result.auc_tuned <= 1.0
+    assert CrossEncoder(str(tmp_path / "model")).config.num_labels == 1