feat(model): add cross-encoder backend

Author: AlisinaDevelo

CHANGELOG.md

@@ -16,6 +16,7 @@
 - Calibration summary and probability-bin diagnostics in `quorabust-report`.
 - Holdout-selected `decision_threshold` metadata from `quorabust-train`.
 - `quorabust-validate-report` release gate for machine-readable model-card JSON.
+- Optional `--feature-backend cross-encoder` for modern transformer pair scoring.
 
 ### Fixed
 - Capped NumPy below 2.5 so Mypy can parse dependency stubs with the configured

README.md

@@ -15,7 +15,8 @@ buyer-facing use cases, method strategy, and production gaps.
 
 ## What it demonstrates
 
-- Pairwise text features with TF-IDF or optional sentence-transformer embeddings
+- Pairwise text features with TF-IDF, optional sentence-transformer embeddings, or
+  optional cross-encoder pair scoring
 - XGBoost training with optional holdout evaluation and early stopping
 - Saved artifacts that include lineage, feature schema, dataset checksum, and metrics
 - FastAPI inference with health/readiness checks, thresholded decisions, OpenAPI docs,
@@ -66,7 +67,7 @@ quorabust-train --csv data/raw/train.csv --out models/quorabust.pkl
 python -m quorabust --csv data/raw/train.csv --out models/quorabust.pkl   # equivalent
 ```
 
-Options: `--max-rows N`, `--eval-fraction 0.1` (default), `--eval-fraction 0` to train on all rows without a holdout, `--seed`, `--feature-backend {tfidf,embedding}`, `--embedding-model …`, `--thresholds`, `--threshold-metric {accuracy,precision,recall,f1}` for holdout-based decision-threshold selection, `--registry-dir` (JSONL registry), `--metadata-out` (JSON sidecar for reviewing artifact lineage without loading the pickle).
+Options: `--max-rows N`, `--eval-fraction 0.1` (default), `--eval-fraction 0` to train on all rows without a holdout, `--seed`, `--feature-backend {tfidf,embedding,cross-encoder}`, `--embedding-model …`, `--cross-encoder-model …`, `--thresholds`, `--threshold-metric {accuracy,precision,recall,f1}` for holdout-based decision-threshold selection, `--registry-dir` (JSONL registry), `--metadata-out` (JSON sidecar for reviewing artifact lineage without loading the pickle).
 
 ### Generate a model card
 

docs/ENTERPRISE.md

@@ -36,7 +36,10 @@ prefer a non-pickle format such as `skops` or ONNX in a future release.
 
 ## Scale and NLP
 
-See [SCALING.md](SCALING.md) for chunked CSV I/O, optional **embedding** training (`pip install ".[nlp]"`, `quorabust-train --feature-backend embedding`), and pointers to distributed XGBoost.
+See [SCALING.md](SCALING.md) for chunked CSV I/O, optional **embedding** training
+(`pip install ".[nlp]"`, `quorabust-train --feature-backend embedding`), optional
+**cross-encoder** pair scoring (`quorabust-train --feature-backend cross-encoder`), and
+pointers to distributed XGBoost.
 
 ## Registry and drift (lightweight)
 

docs/PRODUCT.md

@@ -36,8 +36,8 @@ The modern path is:
 1. **TF-IDF + XGBoost baseline** for speed, explainable feature values, and cheap serving.
 2. **Sentence-transformer embeddings + XGBoost** for better semantic recall when wording
    differs.
-3. **Cross-encoder reranker** for the highest-accuracy pair scoring when latency and
-   compute budget allow it.
+3. **Cross-encoder pair scoring + XGBoost** for the highest-accuracy pair scoring when
+   latency and compute budget allow it.
 
 Do not claim state-of-the-art quality from the checked-in smoke model. Use a real held-out
 dataset or customer-domain labels before making performance claims.

docs/SCALING.md

@@ -10,6 +10,12 @@ This document maps **ambitions** to what ships in-repo and what stays external.
 ## Modern NLP (embeddings)
 
 - **Optional extra** `pip install ".[nlp]"`**: `PairEmbeddingBuilder` in `quorabust.embedding_features` uses `sentence-transformers` to encode pairs and feeds cosine / L2 / pooling stats into the same XGBoost head. Training can select `--feature-backend embedding` in `quorabust-train`.
+- **Cross-encoder pair scoring**: `PairCrossEncoderBuilder` in
+  `quorabust.cross_encoder_features` uses a Sentence Transformers `CrossEncoder` to score
+  each pair directly, then feeds that score plus length stats into the same XGBoost head.
+  Select it with `--feature-backend cross-encoder`. This is the modern high-accuracy path
+  for pair scoring, but it is slower than TF-IDF or bi-encoder embeddings because every
+  pair must be passed through the transformer jointly.
 
 ## Online serving, SLOs, monitoring
 

src/quorabust/cli.py

@@ -81,15 +81,23 @@ def main(argv: list[str] | None = None) -> int:
     p.add_argument("--seed", type=int, default=42)
     p.add_argument(
         "--feature-backend",
-        choices=["tfidf", "embedding"],
+        choices=["tfidf", "embedding", "cross-encoder"],
         default="tfidf",
-        help="tfidf (default) or sentence-transformer embeddings (requires nlp extra)",
+        help=(
+            "tfidf (default), sentence-transformer embeddings, or cross-encoder pair scores "
+            "(requires nlp extra)"
+        ),
     )
     p.add_argument(
         "--embedding-model",
         default="sentence-transformers/all-MiniLM-L6-v2",
         help="When --feature-backend=embedding, SentenceTransformer model id",
     )
+    p.add_argument(
+        "--cross-encoder-model",
+        default="cross-encoder/quora-distilroberta-base",
+        help="When --feature-backend=cross-encoder, CrossEncoder model id",
+    )
     p.add_argument(
         "--registry-dir",
         type=Path,
@@ -146,6 +154,10 @@ def main(argv: list[str] | None = None) -> int:
         from quorabust.embedding_features import PairEmbeddingBuilder
 
         feature_builder = PairEmbeddingBuilder(model_name=args.embedding_model)
+    elif args.feature_backend == "cross-encoder":
+        from quorabust.cross_encoder_features import PairCrossEncoderBuilder
+
+        feature_builder = PairCrossEncoderBuilder(model_name=args.cross_encoder_model)
 
     builder, clf = train_duplicate_classifier(
         train_df,
@@ -154,10 +166,11 @@ def main(argv: list[str] | None = None) -> int:
         feature_builder=feature_builder,
     )
 
-    if args.feature_backend == "embedding":
-        feat_names = ["cos", "l2", "mad", "len_ratio", "len_sum"]
-    else:
-        feat_names = ["cos", "jaccard", "len_ratio", "abs_len_diff", "len_sum"]
+    feat_names = (
+        builder.feature_names()
+        if hasattr(builder, "feature_names")
+        else ["feature_0", "feature_1", "feature_2", "feature_3", "feature_4"]
+    )
 
     meta: dict[str, Any] = {
         "n_train": len(train_df),

src/quorabust/cross_encoder_features.py

@@ -0,0 +1,71 @@
+from __future__ import annotations
+
+import numpy as np
+import pandas as pd
+
+from quorabust.preprocess import clean_text
+
+try:
+    from sentence_transformers import CrossEncoder
+except ImportError:
+    CrossEncoder = None  # type: ignore[misc, assignment]
+
+
+class PairCrossEncoderBuilder:
+    """Cross-encoder pair scores plus simple length features (optional ``nlp`` extra)."""
+
+    def __init__(self, model_name: str = "cross-encoder/quora-distilroberta-base") -> None:
+        if CrossEncoder is None:
+            raise RuntimeError(
+                'Missing dependency: install with pip install "Quorabust[nlp]"',
+            )
+        self.model_name = model_name
+        self._model = CrossEncoder(model_name)
+        self._fitted = True
+
+    def fit(self, corpus: list[str] | None = None) -> PairCrossEncoderBuilder:
+        return self
+
+    def feature_names(self) -> list[str]:
+        return ["cross_score", "len_ratio", "abs_len_diff", "len_sum"]
+
+    def fit_from_frame(
+        self,
+        df: pd.DataFrame,
+        col_q1: str = "question1",
+        col_q2: str = "question2",
+    ) -> PairCrossEncoderBuilder:
+        return self
+
+    def transform_pairs(
+        self,
+        q1: list[str],
+        q2: list[str],
+    ) -> np.ndarray:
+        if len(q1) != len(q2):
+            raise ValueError("q1 and q2 must have the same length.")
+        t1 = [clean_text(x) for x in q1]
+        t2 = [clean_text(x) for x in q2]
+        scores = np.asarray(
+            self._model.predict(list(zip(t1, t2, strict=True)), show_progress_bar=False),
+            dtype=np.float64,
+        ).reshape(-1)
+
+        rows: list[list[float]] = []
+        for idx, (a, b) in enumerate(zip(t1, t2, strict=True)):
+            la, lb = len(a.split()), len(b.split())
+            max_len = max(la, lb, 1)
+            len_ratio = min(la, lb) / max_len
+            rows.append([float(scores[idx]), len_ratio, float(abs(la - lb)), float(la + lb)])
+        return np.asarray(rows, dtype=np.float64)
+
+    def transform_frame(
+        self,
+        df: pd.DataFrame,
+        col_q1: str = "question1",
+        col_q2: str = "question2",
+    ) -> np.ndarray:
+        return self.transform_pairs(
+            df[col_q1].astype(str).tolist(),
+            df[col_q2].astype(str).tolist(),
+        )

tests/test_cli.py

@@ -2,9 +2,18 @@
 
 import pandas as pd
 
+import quorabust.cross_encoder_features as cef
 from quorabust.cli import main
 
 
+class _FakeCrossEncoder:
+    def __init__(self, model_name: str) -> None:
+        self.model_name = model_name
+
+    def predict(self, pairs, show_progress_bar: bool = False):
+        return [0.9 if idx % 2 else 0.1 for idx, _pair in enumerate(pairs)]
+
+
 def _write_synthetic_csv(path, n: int = 30) -> None:
     df = pd.DataFrame(
         {
@@ -101,6 +110,38 @@ def test_cli_rejects_bad_threshold_grid(tmp_path):
     )
 
 
+def test_cli_trains_with_cross_encoder_backend(tmp_path, monkeypatch):
+    monkeypatch.setattr(cef, "CrossEncoder", _FakeCrossEncoder)
+    csv = tmp_path / "train.csv"
+    _write_synthetic_csv(csv)
+    out = tmp_path / "model.pkl"
+    meta = tmp_path / "model.meta.json"
+
+    assert (
+        main(
+            [
+                "--csv",
+                str(csv),
+                "--out",
+                str(out),
+                "--feature-backend",
+                "cross-encoder",
+                "--cross-encoder-model",
+                "fake-cross-encoder",
+                "--metadata-out",
+                str(meta),
+                "--eval-fraction",
+                "0",
+            ]
+        )
+        == 0
+    )
+
+    payload = json.loads(meta.read_text(encoding="utf-8"))
+    assert payload["feature_backend"] == "cross-encoder"
+    assert payload["feature_schema"] == ["cross_score", "len_ratio", "abs_len_diff", "len_sum"]
+
+
 def test_cli_rejects_bad_columns(tmp_path):
     csv = tmp_path / "bad.csv"
     pd.DataFrame({"a": [1]}).to_csv(csv, index=False)

tests/test_cross_encoder_features.py

@@ -0,0 +1,37 @@
+import numpy as np
+import pandas as pd
+import pytest
+
+import quorabust.cross_encoder_features as cef
+
+
+class _FakeCrossEncoder:
+    def __init__(self, model_name: str) -> None:
+        self.model_name = model_name
+
+    def predict(self, pairs, show_progress_bar: bool = False):
+        return np.asarray([0.9 if a.split()[0:1] == b.split()[0:1] else 0.1 for a, b in pairs])
+
+
+def test_pair_cross_encoder_builder_requires_sentence_transformers(monkeypatch):
+    monkeypatch.setattr(cef, "CrossEncoder", None)
+    with pytest.raises(RuntimeError, match="nlp"):
+        cef.PairCrossEncoderBuilder()
+
+
+def test_pair_cross_encoder_builder_shapes(monkeypatch):
+    monkeypatch.setattr(cef, "CrossEncoder", _FakeCrossEncoder)
+    builder = cef.PairCrossEncoderBuilder(model_name="fake-cross-encoder")
+    df = pd.DataFrame(
+        {
+            "question1": ["hello world", "foo bar"],
+            "question2": ["hello there", "baz qux"],
+        }
+    )
+
+    X = builder.fit_from_frame(df).transform_frame(df)
+
+    assert builder.feature_names() == ["cross_score", "len_ratio", "abs_len_diff", "len_sum"]
+    assert X.shape == (2, 4)
+    assert X[0, 0] == 0.9
+    assert X[1, 0] == 0.1