| """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) |
|
|
