"""Compact oblivious-tree ensemble (CatBoost's symmetric trees) evaluated with whole-batch numpy operations. In an oblivious tree every node of one depth level asks the same question, so a depth-d tree is just d (feature, border) tests and a table of 2^d leaves: the d test results, read as bits, are the leaf index. All distinct tests of the ensemble are evaluated once per batch into a boolean matrix [n, S]; the leaf index of every (row, tree) pair is then built with d gathers and the leaves are summed. No python loop over rows or trees, no CatBoost at inference: the .npz file (tests + leaves, a few hundred KB) is the whole model. """ from __future__ import annotations import json import numpy as np class Oblivious: def __init__(self, feat, border, tree_tests, leaves, bias, names=None, classes=None, meta=None): self.feat = np.asarray(feat, np.int32) # [S] feature of test s (last test: always false) self.border = np.asarray(border, np.float32) # [S] test s is x[feat] > border self.tree_tests = np.asarray(tree_tests, np.int32) # [T, D] tests of each tree, bit d = depth d self.leaves = np.asarray(leaves, np.float32) # [T, 2^D, C] self.bias = np.asarray(bias, np.float32) # [C] self.names, self.classes, self.meta = list(names or []), list(classes or []), dict(meta or {}) T, L, C = self.leaves.shape D = self.tree_tests.shape[1] self._flat = np.ascontiguousarray(self.leaves.transpose(2, 0, 1).reshape(C, T * L)) # [C, T*L] self._off = (np.arange(T, dtype=np.int32) * L)[:, None] # tests sorted by feature: each feature's tests are one broadcast compare into a slice of the test matrix order = np.lexsort((self.border, self.feat)) inv = np.empty_like(order) inv[order] = np.arange(len(order)) fs, bs = self.feat[order], self.border[order] self._bcol, self._tt = bs[:, None], inv[self.tree_tests] self._groups = [(int(f), int(np.searchsorted(fs, f, "left")), int(np.searchsorted(fs, f, "right"))) for f in np.unique(fs)] self._w = (1 << np.arange(D)).astype(np.uint8) @property def n_trees(self) -> int: return self.leaves.shape[0] def raw(self, x: np.ndarray, chunk: int = 8192) -> np.ndarray: """x: float32 [n, F] -> raw scores [n, C] (logits for Logloss / MultiClass). Feature-major ([F, m]) so every gather copies whole contiguous rows; leaf indices built as uint8.""" x = np.asarray(x, np.float32) out = np.empty((len(x), self.leaves.shape[2]), np.float32) for s in range(0, len(x), chunk): xt = np.ascontiguousarray(x[s:s + chunk].T) # [F, m] b = np.empty((len(self._bcol), xt.shape[1]), bool) # [S, m] test results for f, lo, hi in self._groups: np.greater(xt[f], self._bcol[lo:hi], out=b[lo:hi]) idx = np.einsum("tdm,d->tm", b.view(np.uint8)[self._tt], self._w) # [T, m] leaf index, bit d = depth d flat = idx.astype(np.int32) flat += self._off for c in range(out.shape[1]): out[s:s + chunk, c] = self._flat[c][flat].sum(0) return out + self.bias def proba_attack(self, x: np.ndarray) -> np.ndarray: """P(attack) for a binary detector (one logit) or 1 - P(benign) for a multi-class model.""" r = self.raw(x) if r.shape[1] == 1: return 1.0 / (1.0 + np.exp(-r[:, 0])) z = np.exp(r - r.max(1, keepdims=True)) p = z / z.sum(1, keepdims=True) return 1.0 - p[:, self.classes.index("benign")] if "benign" in self.classes else p.max(1) # --- io --- def save(self, path: str, leaf_dtype=np.float32) -> None: np.savez_compressed(path, feat=self.feat.astype(np.int16), border=self.border, tree_tests=self.tree_tests.astype(np.int16), leaves=self.leaves.astype(leaf_dtype), bias=self.bias, info=np.frombuffer(json.dumps({"names": self.names, "classes": self.classes, "meta": self.meta}).encode(), np.uint8)) @classmethod def load(cls, path: str) -> "Oblivious": z = np.load(path) info = json.loads(bytes(z["info"]).decode()) return cls(z["feat"], z["border"], z["tree_tests"], z["leaves"].astype(np.float32), z["bias"], info["names"], info["classes"], info["meta"]) @classmethod def from_catboost(cls, model_json: dict, names=None, classes=None, meta=None) -> "Oblivious": """From CatBoost's model.save_model(path, format='json') (float features only, symmetric trees).""" ff = model_json["features_info"]["float_features"] flat = {f["feature_index"]: f["flat_feature_index"] for f in ff} trees = model_json["oblivious_trees"] for t in trees: t["splits"] = t.get("splits") or [] # a depth-0 tree (one leaf) is written with splits: null D = max(1, max(len(t["splits"]) for t in trees)) C = len(trees[0]["leaf_values"]) // (1 << len(trees[0]["splits"])) tests: dict = {} tt = np.full((len(trees), D), -1, np.int64) leaves = np.zeros((len(trees), 1 << D, C), np.float32) for i, t in enumerate(trees): for d, sp in enumerate(t["splits"]): if sp.get("split_type", "FloatFeature") != "FloatFeature": raise ValueError(f"unsupported split {sp}") key = (flat[sp["float_feature_index"]], float(sp["border"])) tt[i, d] = tests.setdefault(key, len(tests)) n = 1 << len(t["splits"]) leaves[i, :n] = np.asarray(t["leaf_values"], np.float32).reshape(n, C) # leaf-major, class inner never = len(tests) # padding test for shallower trees: x[0] > +inf is always false tt[tt < 0] = never feat = [k[0] for k in tests] + [0] border = [k[1] for k in tests] + [np.inf] scale, bias = model_json.get("scale_and_bias", [1.0, [0.0] * C]) bias = np.broadcast_to(np.asarray(bias, np.float32), (C,)) return cls(feat, border, tt, leaves * np.float32(scale), bias, names, classes, meta)