Spaces:
Running
Running
| """ | |
| app/classifiers/cascade_classifier.py | |
| ======================================= | |
| Two-stage cascade classifier: sklearn fast path → ONNX/HF slow path. | |
| For inputs where sklearn is highly confident (prob ≥ high_threshold or | |
| prob ≤ low_threshold), the verdict is returned immediately without invoking | |
| the heavier ONNX/HF model. Only inputs in the ambiguous "grey zone" are | |
| escalated. | |
| Measured on the external benchmark with the default 0.05/0.95 band (see | |
| docs/BENCHMARKS.md): ~77% of prompts are settled by sklearn alone and ~23% are | |
| escalated, with recall/FPR matching the transformer on its own. Hard, | |
| out-of-distribution prompts escalate far more often (~60% of the holdout set). | |
| The block threshold is the slow model's tuned threshold: grey-zone prompts are | |
| scored by the slow model alone, and fast-path sklearn scores sit outside the | |
| [low, high] band (beyond any threshold in it), so one threshold serves both. | |
| The thresholds are tunable via CASCADE_SK_LOW_THRESHOLD / | |
| CASCADE_SK_HIGH_THRESHOLD environment variables (see config.py). | |
| Predict return dict is identical to individual classifiers plus a diagnostic | |
| "stage" key ("sklearn" | "onnx") for audit logging / metrics. | |
| """ | |
| from __future__ import annotations | |
| import logging | |
| from typing import Any | |
| logger = logging.getLogger(__name__) | |
| class CascadeClassifier: | |
| """ | |
| Combines a fast sklearn classifier with a slower ONNX/HF classifier. | |
| Parameters | |
| ---------- | |
| sklearn_clf : SklearnClassifier | |
| Fast first-stage classifier (~1.5 ms/request). | |
| slow_clf : ONNXClassifier | HFClassifier | |
| Accurate slow-path classifier (~18 ms on GPU, ~40 ms INT8 on CPU). | |
| low_threshold : float | |
| sklearn malicious_prob ≤ this → benign fast path (default 0.05). | |
| high_threshold : float | |
| sklearn malicious_prob ≥ this → malicious fast path (default 0.95). | |
| Both classifiers must be loaded before calling predict(). | |
| """ | |
| def __init__( | |
| self, | |
| sklearn_clf, | |
| slow_clf, | |
| low_threshold: float = 0.05, | |
| high_threshold: float = 0.95, | |
| slow_clf_label: str = "onnx", | |
| ) -> None: | |
| self.sklearn_clf = sklearn_clf | |
| self.slow_clf = slow_clf | |
| self.low_threshold = low_threshold | |
| self.high_threshold = high_threshold | |
| self.slow_clf_label = slow_clf_label | |
| def threshold(self) -> float | None: | |
| """Tuned threshold of the slow model (falls back to sklearn's).""" | |
| for clf in (self.slow_clf, self.sklearn_clf): | |
| tuned = getattr(clf, "threshold", None) | |
| if isinstance(tuned, (int, float)): | |
| return float(tuned) | |
| return None | |
| def is_loaded(self) -> bool: | |
| return self.sklearn_clf.is_loaded() and self.slow_clf.is_loaded() | |
| async def predict(self, text: str) -> dict[str, Any]: | |
| """ | |
| Run the two-stage cascade. | |
| Returns: | |
| { | |
| "label": "benign" | "malicious", | |
| "malicious_prob": float, | |
| "benign_prob": float, | |
| "stage": "sklearn" | "onnx", # diagnostic only | |
| } | |
| """ | |
| if not self.is_loaded(): | |
| raise RuntimeError("CascadeClassifier: one or both classifiers are not loaded.") | |
| # ── Stage 1: fast sklearn check ──────────────────────────────────────── | |
| sk_result = await self.sklearn_clf.predict(text) | |
| sk_prob: float = sk_result["malicious_prob"] | |
| if sk_prob >= self.high_threshold: | |
| logger.debug( | |
| "Cascade fast-path MALICIOUS (sklearn=%.3f ≥ %.2f)", | |
| sk_prob, | |
| self.high_threshold, | |
| ) | |
| return {**sk_result, "stage": "sklearn"} | |
| if sk_prob <= self.low_threshold: | |
| logger.debug( | |
| "Cascade fast-path benign (sklearn=%.3f ≤ %.2f)", | |
| sk_prob, | |
| self.low_threshold, | |
| ) | |
| return {**sk_result, "stage": "sklearn"} | |
| # ── Stage 2: grey zone → escalate to ONNX/HF ────────────────────────── | |
| logger.debug( | |
| "Cascade escalating to ONNX (sklearn=%.3f in grey zone [%.2f, %.2f])", | |
| sk_prob, | |
| self.low_threshold, | |
| self.high_threshold, | |
| ) | |
| slow_result = await self.slow_clf.predict(text) | |
| # The slow model decides the grey zone on its own score. Its threshold | |
| # is the cascade's threshold, and the two models' scores are on | |
| # different scales (an INT8 graph's tuned threshold can be ~0.3 while | |
| # sklearn's is ~0.55), so mixing them — e.g. max(sk, slow) — would let | |
| # middling sklearn scores block. Strong sklearn signals are already | |
| # honoured by the high_threshold fast path above. | |
| return {**slow_result, "stage": self.slow_clf_label} | |