| """Run termite-classifier with PyTorch. It reads plain data files and never runs a command. |
| Usage: python termite_classifier.py 'ls -la' [FOLDER] [SHELL] [OS] [SESSION]""" |
| import json, pathlib, struct, sys |
| import torch, torch.nn as nn, torch.nn.functional as F |
|
|
| HERE = pathlib.Path(__file__).resolve().parent |
| QUESTIONS = ("only_reads", "no_secrets", "harmful") |
| TYPES = {"F32": torch.float32} |
| CLS, FIRST = 1, 3 |
|
|
|
|
| def pieces(data): |
| """A piece is a run of ASCII letters, digits or bytes over 127, or one other byte.""" |
| start = 0 |
| for i, b in enumerate(data): |
| if not (48 <= b <= 57 or 65 <= b <= 90 or 97 <= b <= 122 or b > 127): |
| if i > start: |
| yield data[start:i] |
| yield data[i:i + 1] |
| start = i + 1 |
| if start < len(data): |
| yield data[start:] |
|
|
|
|
| class Tokeniser: |
| def __init__(self, merges): |
| self.rank = {tuple(m): i for i, m in enumerate(merges)} |
|
|
| def piece(self, data): |
| w = [FIRST + b for b in data] |
| while len(w) > 1: |
| found = min((self.rank.get(p, 1 << 30), i) for i, p in enumerate(zip(w, w[1:]))) |
| if found[0] == 1 << 30: |
| break |
| w[found[1]:found[1] + 2] = [FIRST + 256 + found[0]] |
| return w |
|
|
| def encode(self, text): |
| ids = [CLS] |
| for p in pieces(text.encode("utf8", "surrogateescape")): |
| ids += self.piece(p) |
| return ids |
|
|
|
|
| class Block(nn.Module): |
| def __init__(self, width, heads): |
| super().__init__() |
| self.heads = heads |
| self.n1, self.n2 = nn.LayerNorm(width), nn.LayerNorm(width) |
| self.qkv, self.out = nn.Linear(width, 3 * width), nn.Linear(width, width) |
| self.up, self.down = nn.Linear(width, 4 * width), nn.Linear(4 * width, width) |
|
|
| def forward(self, x, keep): |
| b, t, w = x.shape |
| q, k, v = self.qkv(self.n1(x)).view(b, t, 3, self.heads, w // self.heads).permute(2, 0, 3, 1, 4) |
| a = F.scaled_dot_product_attention(q, k, v, attn_mask=keep[:, None, None, :]) |
| x = x + self.out(a.transpose(1, 2).reshape(b, t, w)) |
| return x + self.down(F.gelu(self.up(self.n2(x)))) |
|
|
|
|
| class Encoder(nn.Module): |
| def __init__(self, vocab, longest, layers, width, heads): |
| super().__init__() |
| self.tokens, self.places = nn.Embedding(vocab, width), nn.Embedding(longest, width) |
| self.blocks = nn.ModuleList(Block(width, heads) for _ in range(layers)) |
| self.last = nn.LayerNorm(width) |
| self.heads = nn.Linear(width, 3) |
|
|
| def forward(self, ids): |
| keep = ids != 0 |
| x = self.tokens(ids) + self.places(torch.arange(ids.shape[1])) |
| for block in self.blocks: |
| x = block(x, keep) |
| return self.heads(self.last(x)[:, 0]) |
|
|
|
|
| def read_weights(path): |
| with open(path, "rb") as f: |
| size = struct.unpack("<Q", f.read(8))[0] |
| header = json.loads(f.read(size)) |
| body = bytearray(f.read()) |
| return {name: torch.frombuffer(body, dtype=TYPES[t["dtype"]], offset=t["data_offsets"][0], |
| count=(t["data_offsets"][1] - t["data_offsets"][0]) // 4).view(t["shape"]) |
| for name, t in header.items() if name != "__metadata__"} |
|
|
|
|
| class Classifier: |
| def __init__(self, folder=HERE): |
| folder = pathlib.Path(folder) |
| self.config = json.load(open(folder / "config.json")) |
| c = self.config |
| self.tok = Tokeniser(json.load(open(folder / "merges.json"))) |
| self.net = Encoder(c["vocab"], c["longest"], c["layers"], c["width"], c["heads"]) |
| self.net.load_state_dict(read_weights(folder / c["weights"])) |
| self.net.eval() |
|
|
| @torch.no_grad() |
| def answer(self, text): |
| """Chances for each question. A command too long to read gets the strict answer.""" |
| ids = self.tok.encode(text) |
| if len(ids) > self.config["longest"]: |
| return {"only_reads": 0.0, "no_secrets": 0.0, "harmful": 1.0} |
| z = self.net(torch.tensor([ids]))[0] |
| return {q: torch.sigmoid(z[i] / self.config["temperatures"][i]).item() for i, q in enumerate(QUESTIONS)} |
|
|
| def safe(self, a): |
| limit = self.config["limits"] |
| return a["only_reads"] >= limit["only_reads"] and a["no_secrets"] >= limit["no_secrets"] and a["harmful"] <= limit["harmful"] |
|
|
|
|
| if __name__ == "__main__": |
| command, *rest = sys.argv[1:] |
| folder, shell, system, session = (rest + ["/home/dev/src/project", "bash", "linux", "local"][len(rest):])[:4] |
| model = Classifier() |
| answer = model.answer(f"{shell} {system} {session}\n{folder}\n{command}") |
| print(json.dumps({**answer, "safe": model.safe(answer)}, indent=1)) |
|
|
