Z User
v3.1.0: Part B architecture migration — eliminate Hugging Face inference credit dependency
6c516b6 | """Inference runtime for the email-risk model (PyTorch / transformers). | |
| Strictly optional: `load()` returning False (or `predict()` returning None) | |
| means "no ML this run" — the deterministic pipeline continues unchanged. | |
| Imports are lazy so the API boots with or without torch installed, and the test | |
| suite never requires a model download. | |
| The deployed Space image DOES contain CPU torch (Dockerfile installs it from | |
| the PyTorch CPU index) and ships the bundled model at `app/ml/model`. Do not | |
| strip either: the loader fails open, so removing them silently disables the | |
| advisory layer instead of raising. | |
| CPU inference is offloaded to a worker thread (`predict_async`, `load_async`) | |
| because running it inline stalls every concurrent request — measured at 42.6s | |
| cold start before this. | |
| """ | |
| from __future__ import annotations | |
| import asyncio | |
| import threading | |
| from typing import Callable, Optional | |
| from ..config import settings | |
| def default_pipeline_factory(repo: str, batch_size: int): | |
| """Production pipeline builder: transformers (PyTorch). | |
| The ONNX export produced confidently-wrong labels under two transformers | |
| versions while the PyTorch checkpoint scored 0.999 held-out - so inference | |
| runs on the trained weights directly. transformers is imported lazily here: | |
| a module-level import can fail under some import contexts (importlib_metadata | |
| 'Version' bug on Python 3.14), which must degrade to "no ML" only when | |
| inference is actually attempted, not kill the capability at import time. | |
| The repo is a PUBLIC HF Hub download (free, unmetered); no token is used — | |
| Part B removed every HF token from the deployment surface. | |
| """ | |
| import transformers | |
| tokenizer = transformers.AutoTokenizer.from_pretrained(repo) | |
| model = transformers.AutoModelForSequenceClassification.from_pretrained(repo) | |
| return transformers.pipeline( | |
| "text-classification", | |
| model=model, | |
| tokenizer=tokenizer, | |
| device="cpu", | |
| truncation=True, | |
| max_length=64, | |
| ) | |
| class EmailRiskModel: | |
| """Lazily-loaded ONNX text-classification wrapper. Fail-open by design.""" | |
| def __init__( | |
| self, | |
| repo: str, | |
| batch_size: int = 64, | |
| pipeline_factory: Optional[Callable[[str, int], object]] = None, | |
| ): | |
| self.repo = repo | |
| self.batch_size = batch_size | |
| self._factory = pipeline_factory or default_pipeline_factory | |
| self._pipeline = None | |
| self._loaded = False | |
| self.last_error: Optional[str] = None | |
| self._lock = threading.Lock() | |
| def loaded(self) -> bool: | |
| return self._loaded | |
| def load(self) -> bool: | |
| """Load the model. Returns False on ANY failure (fail-open design).""" | |
| with self._lock: | |
| if self._loaded: | |
| return True | |
| try: | |
| self._pipeline = self._factory( | |
| self.repo, self.batch_size | |
| ) | |
| self._loaded = True | |
| return True | |
| except Exception as exc: | |
| self._pipeline = None | |
| self._loaded = False | |
| self.last_error = f"{type(exc).__name__}: {exc}" | |
| return False | |
| async def load_async(self) -> bool: | |
| """Load the model without blocking the event loop. | |
| Loading torch + a checkpoint takes seconds; doing it inline in an async | |
| handler stalls every concurrent request (measured: 42.6s cold start | |
| before this). | |
| """ | |
| return await asyncio.to_thread(self.load) | |
| def predict(self, texts: list[str]) -> Optional[list[tuple[str, float]]]: | |
| """Classify texts. Returns (LABEL, confidence) tuples, or None on failure. | |
| The label is uppercased and guaranteed to be one of the training labels' | |
| names only if the model provides it; unknown labels pass through as-is | |
| for the caller to ignore. | |
| """ | |
| if not self._loaded or self._pipeline is None or not texts: | |
| return None | |
| try: | |
| preds = self._pipeline(texts, batch_size=self.batch_size) | |
| out: list[tuple[str, float]] = [] | |
| for batch in preds: | |
| # transformers may return a flat list or one list per element. | |
| if isinstance(batch, list): | |
| preds_nested = batch | |
| else: | |
| preds_nested = [batch] | |
| for p in preds_nested: | |
| # An unexpected payload shape must fail open, not raise: | |
| # the batch path calls predict() outside its own try. | |
| if not isinstance(p, dict): | |
| raise TypeError(f"unexpected prediction: {type(p).__name__}") | |
| label = str(p.get("label", "")).upper() | |
| try: | |
| score = float(p.get("score", 0.0)) | |
| except (TypeError, ValueError): | |
| score = 0.0 | |
| out.append((label, score)) | |
| return out if len(out) == len(texts) else None | |
| except Exception as exc: | |
| self.last_error = f"predict: {type(exc).__name__}: {exc}" | |
| return None | |
| _model: Optional[EmailRiskModel] = None | |
| _model_lock = threading.Lock() | |
| async def predict_async(model, texts: list[str]): | |
| """Classify texts without blocking the event loop. | |
| Works with either an `EmailRiskModel` (CPU inference in a worker thread) or | |
| any injected test double exposing a sync `predict` (called directly). | |
| """ | |
| if model is None or not texts: | |
| return None | |
| predict = getattr(model, "predict", None) | |
| if predict is None: | |
| return None | |
| if isinstance(model, EmailRiskModel): | |
| return await asyncio.to_thread(predict, texts) | |
| return predict(texts) | |
| async def get_model_async() -> Optional[EmailRiskModel]: | |
| """`get_model()` with the (slow) first load kept off the event loop.""" | |
| global _model | |
| if not settings.ml_disposable_enabled or not settings.ml_model_repo: | |
| return None | |
| if _model is not None: | |
| return _model | |
| await asyncio.to_thread(get_model) | |
| return _model | |
| def get_model() -> Optional[EmailRiskModel]: | |
| """Process-wide model, built AND loaded when flag + repo are configured.""" | |
| global _model | |
| if not settings.ml_disposable_enabled or not settings.ml_model_repo: | |
| return None | |
| with _model_lock: | |
| if _model is None: | |
| _model = EmailRiskModel( | |
| repo=settings.ml_model_repo, | |
| batch_size=settings.ml_batch_size, | |
| ) | |
| # Load at construction: without this the model fail-opens to None | |
| # on every predict and the feature is silently dead (found live). | |
| _model.load() | |
| return _model |