| """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) |
| self.border = np.asarray(border, np.float32) |
| self.tree_tests = np.asarray(tree_tests, np.int32) |
| self.leaves = np.asarray(leaves, np.float32) |
| self.bias = np.asarray(bias, np.float32) |
| 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)) |
| self._off = (np.arange(T, dtype=np.int32) * L)[:, None] |
| |
| 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) |
| b = np.empty((len(self._bcol), xt.shape[1]), bool) |
| 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) |
| 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) |
|
|
| |
| 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 [] |
| 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) |
| never = len(tests) |
| 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) |
|
|
