Download runtime/python/dibaguard/http/model.py from Dibachain/DibaGuard: direct link, hf CLI and curl.
- Browser
- Download file 8.19 kB
-
https://huggingface.co/Dibachain/DibaGuard/resolve/main/runtime/python/dibaguard/http/model.py
- Command line
-
hf download hf://Dibachain/DibaGuard/runtime/python/dibaguard/http/model.py
-
curl -L -o model.py https://huggingface.co/Dibachain/DibaGuard/resolve/main/runtime/python/dibaguard/http/model.py
8.19 kB
| """Tier-1 HTTP model at inference: hashed byte n-grams -> int8 embedding bag -> small MLP, numpy only. | |
| File (.npz): emb int8 [B, E] + scale per column, W1 [E+S, H] / b1, W2 [H, C] / b2 (float32), info json | |
| (buckets, orders, side-feature names, classes, temperature, thresholds). The same hashing (dibaguard.hashing) | |
| runs at training time inside torch's EmbeddingBag, so bucket ids agree bit for bit. Side features (encoding depth, | |
| charset profile, lengths) ride along the bag. One request costs one pass of ngram_buckets + two small matmuls. | |
| """ | |
| from __future__ import annotations | |
| import json | |
| import numpy as np | |
| from dibaguard.hashing import bag_mean, ngram_buckets | |
| from dibaguard.http.canon import charset_profile | |
| SIDE = ["deepest", "len", "n_params", "persian", "non_ascii", "special", "digits"] | |
| def side_features(texts: list[str], deepest) -> np.ndarray: | |
| out = np.zeros((len(texts), len(SIDE)), np.float32) | |
| for i, (t, d) in enumerate(zip(texts, deepest)): | |
| pr = charset_profile(t) | |
| out[i] = (float(d), min(len(t), 8192) / 1024.0, t.count(" | "), pr.get("persian", 0.0), pr.get("non_ascii", 0.0), | |
| pr.get("special", 0.0), pr.get("digits", 0.0)) | |
| return out | |
| class HttpModel: | |
| def __init__(self, emb_q: np.ndarray, scale: np.ndarray, W1, b1, W2, b2, info: dict): | |
| self.emb_q, self.scale = np.asarray(emb_q, np.int8), np.asarray(scale, np.float32) | |
| self.W1, self.b1, self.W2, self.b2 = (np.asarray(a, np.float32) for a in (W1, b1, W2, b2)) | |
| self.info = info | |
| self.buckets, self.orders = int(info["buckets"]), tuple(info["orders"]) | |
| self.max_bytes = int(info.get("max_bytes", 1024)) | |
| def embed(self, texts: list[str]) -> np.ndarray: | |
| idx, doc = ngram_buckets(texts, self.buckets, self.orders, self.max_bytes) | |
| return bag_mean(self.emb_q, idx, doc, len(texts), self.scale) | |
| def logits(self, texts: list[str], deepest=None) -> np.ndarray: | |
| v = self.embed(texts) | |
| if self.W1.shape[0] > v.shape[1]: | |
| side = side_features(texts, deepest if deepest is not None else [0] * len(texts)) | |
| v = np.hstack([v, side]) | |
| h = np.maximum(v @ self.W1 + self.b1, 0.0) | |
| return h @ self.W2 + self.b2 | |
| def proba_attack(self, texts: list[str], deepest=None) -> np.ndarray: | |
| z = self.logits(texts, deepest)[:, 0] / float(self.info.get("temperature", 1.0)) | |
| return 1.0 / (1.0 + np.exp(-z)) | |
| def save(self, path: str) -> None: | |
| np.savez_compressed(path, emb=self.emb_q, scale=self.scale, W1=self.W1, b1=self.b1, W2=self.W2, b2=self.b2, | |
| info=np.frombuffer(json.dumps(self.info).encode(), np.uint8)) | |
| def load(cls, path: str) -> "HttpModel": | |
| z = np.load(path) | |
| return cls(z["emb"], z["scale"], z["W1"], z["b1"], z["W2"], z["b2"], json.loads(bytes(z["info"]).decode())) | |
| def quantise(emb: np.ndarray) -> tuple[np.ndarray, np.ndarray]: | |
| """float [B, E] -> (int8, per-column scale) so that int8 * scale ~ emb.""" | |
| scale = np.maximum(np.abs(emb).max(0), 1e-8) / 127.0 | |
| return np.clip(np.round(emb / scale), -127, 127).astype(np.int8), scale.astype(np.float32) | |
| def novelty_scores(texts, seen: np.ndarray, buckets: int, orders, max_bytes: int = 1024, leave_one_out: bool = False) -> np.ndarray: | |
| idx, doc = ngram_buckets(texts, buckets, orders, max_bytes) | |
| out = np.zeros(len(texts)) | |
| if idx.size == 0: | |
| return out | |
| cnt = seen[idx].astype(np.float64) | |
| if leave_one_out: # subtract this request's own contribution to the vocabulary | |
| key = doc.astype(np.int64) * buckets + idx | |
| _, inv, own = np.unique(key, return_inverse=True, return_counts=True) | |
| cnt = np.maximum(cnt - own[inv], 0.0) | |
| w = 1.0 / (1.0 + cnt) # unseen -> 1, seen once -> 0.5, common -> ~0 | |
| tot = np.bincount(doc, minlength=len(texts)).astype(np.float64) | |
| return np.bincount(doc, weights=w, minlength=len(texts)) / np.maximum(tot, 1.0) | |
| class HttpModel2(HttpModel): | |
| """Site-conditioned tier-1 (http_v2): the same bag + MLP, plus a *site profile* taken from a sample of the site's | |
| own benign requests - the mean bag vector (a request's cosine distance to it is an input feature) and the mean | |
| hidden vector, which modulates the hidden layer (FiLM). A site without a profile uses the one shipped from | |
| training. Weights: W1 [E+S+1, H], Wc [H, 2H] / bc, W2 [H, C].""" | |
| def __init__(self, emb_q, scale, W1, b1, Wc, bc, W2, b2, info: dict, site: dict | None = None): | |
| super().__init__(emb_q, scale, W1, b1, W2, b2, info) | |
| self.Wc, self.bc = np.asarray(Wc, np.float32), np.asarray(bc, np.float32) | |
| self.site = site or {} | |
| def _hidden(self, texts, deepest, s_mean: np.ndarray) -> np.ndarray: | |
| v = self.embed(texts) | |
| side = side_features(texts, deepest if deepest is not None else [0] * len(texts)) | |
| num = (v * s_mean[None, :]).sum(1) | |
| den = np.linalg.norm(v, axis=1) * np.linalg.norm(s_mean) + 1e-6 | |
| dist = (1.0 - num / den)[:, None].astype(np.float32) | |
| return np.maximum(np.hstack([v, side, dist]) @ self.W1 + self.b1, 0.0) | |
| def site_profile(self, texts: list[str], deepest=None) -> dict: | |
| """Benign requests of a site (a few hundred suffice) -> {'s_mean': [E], 'ctx': [H], 'seen': bucket counts [B]}. | |
| 'seen' is the site's own n-gram vocabulary: what this application's normal requests are made of.""" | |
| v = self.embed(texts) | |
| s_mean = v.mean(0).astype(np.float32) | |
| idx, _ = ngram_buckets(texts, self.buckets, self.orders, self.max_bytes) | |
| seen = np.bincount(idx, minlength=self.buckets).astype(np.int32) if idx.size else np.zeros(self.buckets, np.int32) | |
| return {"s_mean": s_mean, "ctx": self._hidden(texts, deepest, s_mean).mean(0).astype(np.float32), "seen": seen, "n_ref": len(texts)} | |
| def novelty(self, texts: list[str], site: dict | None = None, leave_one_out: bool = False) -> np.ndarray: | |
| """Site-only tier: how much of a request is made of n-grams this site's normal traffic never contains, weighted | |
| by rarity (an n-gram seen in one reference request counts almost as new) -> [n] in [0, 1]. No labels. | |
| leave_one_out: the texts are the reference itself - each request's own n-grams are taken out of the counts, | |
| so reference and live requests are scored the same way (exchangeable, as the threshold needs).""" | |
| site = site or self.site | |
| seen = site.get("seen") | |
| if seen is None: | |
| return np.zeros(len(texts)) | |
| return novelty_scores(texts, seen, self.buckets, self.orders, self.max_bytes, leave_one_out) | |
| def logits(self, texts: list[str], deepest=None, site: dict | None = None) -> np.ndarray: | |
| site = site or self.site | |
| h = self._hidden(texts, deepest, site["s_mean"]) | |
| g = np.tanh(site["ctx"] @ self.Wc + self.bc) | |
| H = h.shape[1] | |
| h = h * (1.0 + g[:H]) + g[H:] | |
| return h @ self.W2 + self.b2 | |
| def proba_attack(self, texts: list[str], deepest=None, site: dict | None = None) -> np.ndarray: | |
| z = self.logits(texts, deepest, site)[:, 0] / float(self.info.get("temperature", 1.0)) | |
| return 1.0 / (1.0 + np.exp(-z)) | |
| def save(self, path: str) -> None: | |
| np.savez_compressed(path, emb=self.emb_q, scale=self.scale, W1=self.W1, b1=self.b1, Wc=self.Wc, bc=self.bc, W2=self.W2, | |
| b2=self.b2, site_s_mean=self.site["s_mean"], site_ctx=self.site["ctx"], | |
| site_seen=self.site.get("seen", np.zeros(self.buckets, np.int32)), | |
| info=np.frombuffer(json.dumps(self.info).encode(), np.uint8)) | |
| def load(cls, path: str) -> "HttpModel2": | |
| z = np.load(path) | |
| return cls(z["emb"], z["scale"], z["W1"], z["b1"], z["Wc"], z["bc"], z["W2"], z["b2"], | |
| json.loads(bytes(z["info"]).decode()), {"s_mean": z["site_s_mean"], "ctx": z["site_ctx"], | |
| "seen": z["site_seen"] if "site_seen" in z.files else None}) | |