| """PyTorch loading code for send-gate v44. |
| |
| The model is a binary harm classifier: mmBERT-base backbone, a byte level |
| character CNN branch, CLS pooling, and a small MLP gate head. |
| |
| Usage: |
| import torch |
| from model import StudentGate |
| from tokenizers import Tokenizer |
| |
| ckpt = torch.load("model.pt", map_location="cpu", weights_only=False) |
| model = StudentGate("jhu-clsp/mmBERT-base") |
| model.load_state_dict(ckpt["model"]) |
| model.eval() |
| |
| tok = Tokenizer.from_file("tokenizer.json") |
| text = "kys, grapje" |
| ids = tok.encode(text).ids[:128] |
| ids = ids + [0] * (128 - len(ids)) |
| mask = [1] * len(tok.encode(text).ids[:128]) + [0] * (128 - len(tok.encode(text).ids[:128])) |
| raw = text.encode("utf-8", errors="ignore")[:256] |
| chars = list(raw) + [0] * (256 - len(raw)) |
| cmask = [1 if b else 0 for b in chars] |
| |
| import torch |
| with torch.no_grad(): |
| logit = model( |
| input_ids=torch.tensor([ids]), |
| attention_mask=torch.tensor([mask]), |
| char_ids=torch.tensor([chars]), |
| char_mask=torch.tensor([cmask]), |
| ) |
| p = torch.sigmoid(logit).item() |
| print(p, "BLOCK" if p >= 0.76 else "allow") |
| """ |
| from __future__ import annotations |
|
|
| import torch |
| import torch.nn as nn |
| from transformers import AutoModel |
|
|
| PAD = 0 |
|
|
|
|
| class CharCNN(nn.Module): |
| def __init__(self, char_dim: int = 32, filters: int = 64, |
| kernels: tuple[int, ...] = (1, 2, 3, 4, 5), dropout: float = 0.1): |
| super().__init__() |
| self.emb = nn.Embedding(256, char_dim, padding_idx=PAD) |
| self.convs = nn.ModuleList([ |
| nn.Conv1d(char_dim, filters, kernel_size=k, padding=k // 2) |
| for k in kernels |
| ]) |
| self.dropout = nn.Dropout(dropout) |
| self.out_dim = filters * len(kernels) |
|
|
| def forward(self, char_ids: torch.Tensor, char_mask: torch.Tensor | None = None): |
| x = self.emb(char_ids).transpose(1, 2) |
| if char_mask is not None: |
| x = x * char_mask.unsqueeze(1) |
| pooled = [] |
| for conv in self.convs: |
| y = torch.relu(conv(x)) |
| if y.size(2) != x.size(2): |
| y = y[:, :, :x.size(2)] |
| if char_mask is not None: |
| y = y.masked_fill(char_mask.unsqueeze(1) == 0, float("-inf")) |
| pooled.append(y.max(dim=2).values) |
| return self.dropout(torch.cat(pooled, dim=-1)) |
|
|
|
|
| class StudentGate(nn.Module): |
| def __init__(self, backbone_dir: str = "jhu-clsp/mmBERT-base", |
| dropout: float = 0.1, head_hidden: int = 256, |
| char_dim: int = 32, char_filters: int = 64, |
| char_kernels: tuple[int, ...] = (1, 2, 3, 4, 5)): |
| super().__init__() |
| self.backbone = AutoModel.from_pretrained(backbone_dir) |
| self.backbone.to(torch.float32) |
| h = self.backbone.config.hidden_size |
| self.dropout = nn.Dropout(dropout) |
| self.char_cnn = CharCNN(char_dim=char_dim, filters=char_filters, |
| kernels=char_kernels, dropout=dropout) |
| self.char_proj = nn.Linear(self.char_cnn.out_dim, h) |
| self.gate = nn.Sequential( |
| nn.Linear(h, head_hidden), nn.GELU(), nn.Dropout(dropout), |
| nn.Linear(head_hidden, 1), |
| ) |
|
|
| def forward(self, input_ids, attention_mask, char_ids, char_mask): |
| out = self.backbone(input_ids=input_ids, attention_mask=attention_mask) |
| pooled = out.last_hidden_state[:, 0] |
| char_feat = self.char_cnn(char_ids, char_mask) |
| pooled = pooled + self.char_proj(char_feat) |
| return self.gate(self.dropout(pooled)).squeeze(-1) |
|
|
|
|
| if __name__ == "__main__": |
| import sys |
| from tokenizers import Tokenizer |
|
|
| ckpt_path = sys.argv[1] if len(sys.argv) > 1 else "model.pt" |
| text = sys.argv[2] if len(sys.argv) > 2 else "kys, grapje" |
| ckpt = torch.load(ckpt_path, map_location="cpu", weights_only=False) |
| model = StudentGate() |
| model.load_state_dict(ckpt["model"]) |
| model.eval() |
| tok = Tokenizer.from_file("tokenizer.json") |
| ids = tok.encode(text).ids[:128] |
| n = len(ids) |
| ids = ids + [0] * (128 - n) |
| mask = [1] * n + [0] * (128 - n) |
| raw = text.encode("utf-8", errors="ignore")[:256] |
| chars = list(raw) + [0] * (256 - len(raw)) |
| cmask = [1 if b else 0 for b in chars] |
| with torch.no_grad(): |
| logit = model( |
| input_ids=torch.tensor([ids]), |
| attention_mask=torch.tensor([mask]), |
| char_ids=torch.tensor([chars]), |
| char_mask=torch.tensor([cmask]), |
| ) |
| p = torch.sigmoid(logit).item() |
| print(f"{text!r} -> p={p:.4f} -> {'BLOCK' if p >= 0.76 else 'allow'}") |
|
|
