"""Pure send-gate classifier: tokenizer -> ONNX -> sigmoid -> threshold. No regex. No canonicalisation. No word lists. No model-view routing. The classifier sees the raw message plus its UTF-8 bytes for the character CNN. """ from __future__ import annotations from dataclasses import dataclass from pathlib import Path import numpy as np import onnxruntime as ort from tokenizers import Tokenizer from sklearn.feature_extraction.text import HashingVectorizer DEFAULT_THRESHOLD = 0.94 DEFAULT_NGRAM_THRESHOLD = 0.95 NGRAM_FEATURES = 2 ** 18 MAX_LEN = 128 CHAR_LEN = 256 @dataclass class Result: blocked: bool score: float reason: str | None class SendGate: def __init__(self, model_path: str | Path = "model.onnx", tokenizer_path: str | Path = "tokenizer.json", threshold: float = DEFAULT_THRESHOLD, ngram_threshold: float = DEFAULT_NGRAM_THRESHOLD, ngram_model_path: str | Path = "ngram_model.npz", vocab_map_path: str | Path | None = None, policy_model_path: str | Path | None = None, policy_threshold: float = 0.85, benign_model_path: str | Path | None = None, benign_threshold: float = 0.99, dynamic_lengths: bool | None = None, num_threads: int | None = None, inter_op_num_threads: int | None = None, enable_mem_arena: bool | None = None, providers: list[str] | None = None): self.tok = Tokenizer.from_file(str(tokenizer_path)) sess_options = ort.SessionOptions() sess_options.log_severity_level = 3 if num_threads: sess_options.intra_op_num_threads = int(num_threads) if inter_op_num_threads: sess_options.inter_op_num_threads = int(inter_op_num_threads) if enable_mem_arena is not None: sess_options.enable_cpu_mem_arena = bool(enable_mem_arena) self.sess = ort.InferenceSession( str(model_path), sess_options=sess_options, providers=providers or ["CPUExecutionProvider"] ) self.threshold = threshold self.ngram_threshold = ngram_threshold model_seq = self.sess.get_inputs()[0].shape[1] if dynamic_lengths is None: dynamic_lengths = isinstance(model_seq, str) self.dynamic_lengths = dynamic_lengths self.length_multiple = 1 self.char_multiple = 1 self.vec = HashingVectorizer(analyzer="char", ngram_range=(2, 5), n_features=NGRAM_FEATURES, alternate_sign=False, norm=None) z = np.load(str(ngram_model_path)) self.ngram_coef = z["coef"] self.ngram_intercept = float(np.asarray(z["intercept"]).reshape(-1)[0]) self.vocab_map = None if vocab_map_path is not None: z = np.load(str(vocab_map_path)) self.vocab_map = z["old_to_new"] self.policy_threshold = float(policy_threshold) self.policy_coef = None self.policy_intercept = 0.0 policy_path = (Path(policy_model_path) if policy_model_path is not None else Path(str(ngram_model_path)).with_name("policy_model.npz")) if policy_path.exists(): z = np.load(str(policy_path)) self.policy_coef = np.asarray(z["coef"], dtype=np.float64).reshape(-1) self.policy_intercept = float(np.asarray(z["intercept"]).reshape(-1)[0]) elif policy_model_path is not None: raise FileNotFoundError(f"policy model not found: {policy_path}") self.benign_threshold = float(benign_threshold) self.benign_coef = None self.benign_intercept = 0.0 benign_path = (Path(benign_model_path) if benign_model_path is not None else Path(str(ngram_model_path)).with_name("benign_model.npz")) if benign_path.exists(): z = np.load(str(benign_path)) self.benign_coef = np.asarray(z["coef"], dtype=np.float64).reshape(-1) self.benign_intercept = float(np.asarray(z["intercept"]).reshape(-1)[0]) elif benign_model_path is not None: raise FileNotFoundError(f"benign model not found: {benign_path}") def _encode(self, text: str): ids = self.tok.encode(text).ids[:MAX_LEN] if self.vocab_map is not None: ids = self.vocab_map[np.asarray(ids, dtype=np.int64)].tolist() raw = text.encode("utf-8", errors="ignore")[:CHAR_LEN] if self.dynamic_lengths: n = max(2, len(ids)) ids = ids + [0] * (n - len(ids)) mask = [1] * len(ids) + [0] * (n - len(ids)) chars = list(raw) if raw else [32] cmask = [1] * len(chars) return (np.array([ids], dtype=np.int64), np.array([mask], dtype=np.int64), np.array([chars], dtype=np.int64), np.array([cmask], dtype=np.int64)) n = len(ids) ids = ids + [0] * (MAX_LEN - n) mask = [1] * n + [0] * (MAX_LEN - n) chars = list(raw) + [0] * (CHAR_LEN - len(raw)) cmask = [1 if b else 0 for b in chars] return (np.array([ids], dtype=np.int64), np.array([mask], dtype=np.int64), np.array([chars], dtype=np.int64), np.array([cmask], dtype=np.int64)) def main_score(self, text: str) -> float: ids, mask, chars, cmask = self._encode(text) logit = self.sess.run(None, { "input_ids": ids, "attention_mask": mask, "char_ids": chars, "char_mask": cmask, })[0].flatten()[0] return float(1.0 / (1.0 + np.exp(-logit))) def ngram_score(self, text: str) -> float: x = self.vec.transform([text]) logit = float(np.asarray(x @ self.ngram_coef).reshape(-1)[0] + self.ngram_intercept) return float(1.0 / (1.0 + np.exp(-logit))) def policy_score(self, text: str) -> float: if self.policy_coef is None: return 0.0 x = self.vec.transform([text]) logit = float(np.asarray(x @ self.policy_coef).reshape(-1)[0] + self.policy_intercept) return float(1.0 / (1.0 + np.exp(-logit))) def benign_score(self, text: str) -> float: if self.benign_coef is None: return 0.0 x = self.vec.transform([text]) logit = float(np.asarray(x @ self.benign_coef).reshape(-1)[0] + self.benign_intercept) return float(1.0 / (1.0 + np.exp(-logit))) def score(self, text: str) -> float: values = [self.main_score(text), self.ngram_score(text)] if self.policy_coef is not None: values.append(self.policy_score(text)) return max(values) def check(self, text: str) -> Result: main = self.main_score(text) ngram = self.ngram_score(text) policy = self.policy_score(text) if self.policy_coef is not None else None benign = self.benign_score(text) if self.benign_coef is not None else None blocked = main >= self.threshold or ngram >= self.ngram_threshold if policy is not None and policy >= self.policy_threshold: blocked = True reason = None if blocked: if main >= self.threshold: reason = "main" elif ngram >= self.ngram_threshold: reason = "ngram" else: reason = "policy" benign_override = (benign is not None and benign >= self.benign_threshold and (policy is None or policy < 0.99)) if benign_override: blocked = False reason = "benign" values = [main, ngram] + ([policy] if policy is not None else []) return Result(blocked=blocked, score=max(values), reason=reason)