Text Classification
ONNX
English
multilingual
pico-type
pico_type
byte-level
content-classification
multi-head
classifier
edge-ai
code-detection
language-detection
clipboard
tiny
open-source
onnx-runtime
content-moderation
risk-detection
secrets-detection
prompt-injection-detection
pii-detection
pico-type-arth
Eval Results (legacy)
| """pico-type: a tiny byte-level multi-head content classifier.""" | |
| from __future__ import annotations | |
| from dataclasses import dataclass, field | |
| import torch | |
| import torch.nn.functional as F | |
| from torch import nn | |
| TIERS: dict[str, int] = { | |
| "tiny": 16, | |
| "small": 64, | |
| "base": 192, | |
| "pro": 576, | |
| } | |
| class PicoTypeConfig: | |
| max_bytes: int = 1024 | |
| embed_dim: int = 96 | |
| trunk_dim: int = 192 | |
| num_heads: int = 4 | |
| num_attn_layers: int = 2 | |
| rope_theta: float = 500_000.0 | |
| conv_kernels: tuple[int, ...] = (3, 5, 7) | |
| tiers: dict[str, int] = field(default_factory=lambda: dict(TIERS)) | |
| num_coarse: int = 12 | |
| num_modality: int = 8 | |
| num_subtype: int = 24 | |
| num_code_lang: int = 62 | |
| num_text_lang: int = 30 | |
| num_file_mime: int = 90 | |
| num_risk: int = 6 | |
| dropout: float = 0.1 | |
| undetected_threshold: float = 0.4 | |
| risk_threshold: float = 0.5 | |
| pad_byte: int = 0 | |
| def tier_dim(self, tier: str) -> int: | |
| if tier not in self.tiers: | |
| raise KeyError(f"unknown tier '{tier}', available: {list(self.tiers)}") | |
| return self.tiers[tier] | |
| class ByteEmbed(nn.Module): | |
| def __init__(self, num_bytes: int = 256, dim: int = 96): | |
| super().__init__() | |
| self.embed = nn.Embedding(num_bytes, dim) | |
| nn.init.normal_(self.embed.weight, std=0.02) | |
| def forward(self, x: torch.Tensor) -> torch.Tensor: | |
| return self.embed(x) | |
| class ConvBlock(nn.Module): | |
| def __init__(self, in_dim: int, out_dim: int, kernel_size: int, dropout: float = 0.1): | |
| super().__init__() | |
| padding = kernel_size // 2 | |
| self.conv = nn.Conv1d(in_dim, out_dim, kernel_size, padding=padding, bias=False) | |
| self.norm = nn.LayerNorm(out_dim) | |
| self.act = nn.GELU() | |
| self.drop = nn.Dropout(dropout) | |
| if in_dim != out_dim: | |
| self.proj = nn.Conv1d(in_dim, out_dim, kernel_size=1, bias=False) | |
| else: | |
| self.proj = nn.Identity() | |
| def forward(self, x: torch.Tensor) -> torch.Tensor: | |
| residual = self.proj(x) | |
| h = self.conv(x) | |
| h = h.transpose(1, 2) | |
| h = self.norm(h) | |
| h = self.act(h) | |
| h = self.drop(h) | |
| h = h.transpose(1, 2) | |
| return h + residual | |
| def _apply_rope(x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor) -> torch.Tensor: | |
| d = x.size(-1) | |
| x1, x2 = x[..., : d // 2], x[..., d // 2 :] | |
| rot = torch.cat([-x2, x1], dim=-1) | |
| return x * cos + rot * sin | |
| class RotaryPosEmb(nn.Module): | |
| def __init__(self, dim: int, theta: float = 500_000.0, max_seq: int = 4096): | |
| super().__init__() | |
| if dim % 2 != 0: | |
| raise ValueError("rope dim must be even") | |
| inv_freq = 1.0 / (theta ** (torch.arange(0, dim, 2).float() / dim)) | |
| self.register_buffer("inv_freq", inv_freq, persistent=False) | |
| self._max_seq = max_seq | |
| self._cached_seq = 0 | |
| self._cached_dim = dim | |
| self._build_cache(max_seq, dim) | |
| def _build_cache(self, max_seq: int, dim: int) -> None: | |
| t = torch.arange(max_seq, dtype=self.inv_freq.dtype, device=self.inv_freq.device) | |
| freqs = torch.einsum("i,j->ij", t, self.inv_freq) | |
| emb = torch.cat([freqs, freqs], dim=-1) | |
| self.register_buffer("cos_cached", emb.cos()[None, None, :, :], persistent=False) | |
| self.register_buffer("sin_cached", emb.sin()[None, None, :, :], persistent=False) | |
| self._cached_seq = max_seq | |
| self._cached_dim = dim | |
| def forward(self, q: torch.Tensor, k: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: | |
| seq = q.size(-2) | |
| if seq > self._cached_seq: | |
| self._build_cache(max(seq, self._cached_seq * 2), self._cached_dim) | |
| cos = self.cos_cached[:, :, :seq, :] | |
| sin = self.sin_cached[:, :, :seq, :] | |
| return _apply_rope(q, cos, sin), _apply_rope(k, cos, sin) | |
| class AttnBlock(nn.Module): | |
| def __init__( | |
| self, | |
| dim: int, | |
| num_heads: int, | |
| dropout: float = 0.1, | |
| rope_theta: float = 500_000.0, | |
| ): | |
| super().__init__() | |
| if dim % num_heads != 0: | |
| raise ValueError(f"dim {dim} not divisible by num_heads {num_heads}") | |
| self.num_heads = num_heads | |
| self.head_dim = dim // num_heads | |
| self.qkv = nn.Linear(dim, 3 * dim, bias=False) | |
| self.out_proj = nn.Linear(dim, dim) | |
| self.norm1 = nn.LayerNorm(dim) | |
| self.norm2 = nn.LayerNorm(dim) | |
| self.mlp = nn.Sequential( | |
| nn.Linear(dim, 4 * dim), | |
| nn.GELU(), | |
| nn.Dropout(dropout), | |
| nn.Linear(4 * dim, dim), | |
| ) | |
| self.attn_drop = dropout | |
| self.resid_drop = nn.Dropout(dropout) | |
| self.rope = RotaryPosEmb(self.head_dim, rope_theta) | |
| def forward(self, x: torch.Tensor, mask: torch.Tensor | None = None) -> torch.Tensor: | |
| B, L, D = x.shape | |
| h = self.norm1(x) | |
| qkv = self.qkv(h).reshape(B, L, 3, self.num_heads, self.head_dim).permute(2, 0, 3, 1, 4) | |
| q, k, v = qkv[0], qkv[1], qkv[2] | |
| q, k = self.rope(q, k) | |
| attn_mask_float = None | |
| if mask is not None: | |
| attn_mask_float = mask[:, None, None, :].to(dtype=q.dtype) | |
| attn_mask_float = torch.where(attn_mask_float.to(torch.bool), 0.0, float("-inf")) | |
| out = F.scaled_dot_product_attention( | |
| q, k, v, | |
| attn_mask=attn_mask_float, | |
| dropout_p=self.attn_drop if self.training else 0.0, | |
| is_causal=False, | |
| ) | |
| out = out.transpose(1, 2).reshape(B, L, D) | |
| x = x + self.resid_drop(self.out_proj(out)) | |
| x = x + self.mlp(self.norm2(x)) | |
| return x | |
| class Pool(nn.Module): | |
| def forward( | |
| self, | |
| x: torch.Tensor, | |
| mask: torch.Tensor | None = None, | |
| ) -> torch.Tensor: | |
| if mask is not None: | |
| m = mask.unsqueeze(-1).to(x.dtype) | |
| denom = m.sum(dim=1).clamp(min=1.0) | |
| mean = (x * m).sum(dim=1) / denom | |
| x_for_max = x.masked_fill(m == 0, float("-inf")) | |
| max_ = x_for_max.max(dim=1).values | |
| x_for_std = x * m | |
| mean_for_std = x_for_std.sum(dim=1, keepdim=True) / denom.unsqueeze(1) | |
| sq = ((x_for_std - mean_for_std * m) ** 2) * m | |
| std = sq.sum(dim=1) / denom | |
| std = std.sqrt().clamp(min=0.0) | |
| else: | |
| mean = x.mean(dim=1) | |
| max_ = x.max(dim=1).values | |
| std = x.std(dim=1, unbiased=False) | |
| return torch.cat([mean, max_, std], dim=-1) | |
| class MatryoshkaHead(nn.Module): | |
| def __init__( | |
| self, | |
| in_dim: int, | |
| tier_dims: dict[str, int], | |
| out_dim: int, | |
| ): | |
| super().__init__() | |
| if in_dim != max(tier_dims.values()): | |
| raise ValueError( | |
| f"trunk dim {in_dim} must equal max tier dim {max(tier_dims.values())}" | |
| ) | |
| self.tier_dims = dict(tier_dims) | |
| self.linears = nn.ModuleDict({ | |
| name: nn.Linear(d, out_dim) for name, d in tier_dims.items() | |
| }) | |
| for linear in self.linears.values(): | |
| nn.init.zeros_(linear.bias) | |
| nn.init.normal_(linear.weight, std=0.02) | |
| def forward(self, x: torch.Tensor, tier: str) -> torch.Tensor: | |
| if tier not in self.linears: | |
| raise KeyError(f"unknown tier '{tier}', available: {list(self.linears)}") | |
| d = self.tier_dims[tier] | |
| sliced = x[..., :d] | |
| return self.linears[tier](sliced) | |
| class PicoType(nn.Module): | |
| def __init__(self, config: PicoTypeConfig | None = None): | |
| super().__init__() | |
| self.config = config or PicoTypeConfig() | |
| cfg = self.config | |
| self.embed = ByteEmbed(256, cfg.embed_dim) | |
| in_dim = cfg.embed_dim | |
| self.conv_blocks = nn.ModuleList() | |
| for k in cfg.conv_kernels: | |
| self.conv_blocks.append( | |
| ConvBlock(in_dim, cfg.trunk_dim, k, cfg.dropout) | |
| ) | |
| in_dim = cfg.trunk_dim | |
| self.attn_blocks = nn.ModuleList([ | |
| AttnBlock(cfg.trunk_dim, cfg.num_heads, cfg.dropout, cfg.rope_theta) | |
| for _ in range(cfg.num_attn_layers) | |
| ]) | |
| self.pool = Pool() | |
| self.trunk_dim = 3 * cfg.trunk_dim | |
| self.heads = nn.ModuleDict({ | |
| "coarse": MatryoshkaHead(self.trunk_dim, cfg.tiers, cfg.num_coarse), | |
| "modality": MatryoshkaHead(self.trunk_dim, cfg.tiers, cfg.num_modality), | |
| "subtype": MatryoshkaHead(self.trunk_dim, cfg.tiers, cfg.num_subtype), | |
| "code_lang": MatryoshkaHead(self.trunk_dim, cfg.tiers, cfg.num_code_lang), | |
| "text_lang": MatryoshkaHead(self.trunk_dim, cfg.tiers, cfg.num_text_lang), | |
| "file_mime": MatryoshkaHead(self.trunk_dim, cfg.tiers, cfg.num_file_mime), | |
| "risk": MatryoshkaHead(self.trunk_dim, cfg.tiers, cfg.num_risk), | |
| }) | |
| def forward( | |
| self, | |
| bytes_input: torch.Tensor, | |
| mask: torch.Tensor | None = None, | |
| tier: str = "base", | |
| ) -> dict[str, torch.Tensor]: | |
| if bytes_input.dtype != torch.long: | |
| bytes_input = bytes_input.long() | |
| if bytes_input.size(1) > self.config.max_bytes: | |
| bytes_input = bytes_input[:, : self.config.max_bytes] | |
| if mask is not None: | |
| mask = mask[:, : self.config.max_bytes] | |
| if mask is None: | |
| mask = torch.ones_like(bytes_input, dtype=torch.bool) | |
| else: | |
| mask = mask.to(device=bytes_input.device, dtype=torch.bool) | |
| empty_rows = mask.sum(dim=1) == 0 | |
| if empty_rows.any(): | |
| mask = mask.clone() | |
| mask[empty_rows, 0] = True | |
| x = self.embed(bytes_input) | |
| x = x.transpose(1, 2) | |
| for block in self.conv_blocks: | |
| x = block(x) | |
| x = x.transpose(1, 2) | |
| for block in self.attn_blocks: | |
| x = block(x, mask) | |
| pooled = self.pool(x, mask) | |
| return {name: head(pooled, tier) for name, head in self.heads.items()} | |
| def parameter_count(self, tier: str | None = None) -> int: | |
| if tier is None: | |
| return sum(p.numel() for p in self.parameters()) | |
| kept = 0 | |
| for name, p in self.named_parameters(): | |
| if ".linears." in name: | |
| if f".linears.{tier}." in name: | |
| kept += p.numel() | |
| else: | |
| kept += p.numel() | |
| return kept | |
| def tier_sizes(self) -> dict[str, int]: | |
| return {tier: self.parameter_count(tier) for tier in self.config.tiers} | |
| def encode_bytes(data: bytes, max_len: int = 1024, pad: int = 0) -> tuple[torch.Tensor, torch.Tensor]: | |
| raw = list(data[:max_len]) | |
| seq = raw + [pad] * (max_len - len(raw)) | |
| mask = [1] * len(raw) + [0] * (max_len - len(raw)) | |
| return torch.tensor([seq], dtype=torch.long), torch.tensor([mask], dtype=torch.long) | |
| def smoke_test() -> dict[str, int]: | |
| cfg = PicoTypeConfig(max_bytes=128) | |
| model = PicoType(cfg).eval() | |
| sample = ( | |
| b"def hello():\n print('hi')\n\nif __name__ == '__main__':\n hello()\n" | |
| ) | |
| x, mask = encode_bytes(sample, max_len=cfg.max_bytes, pad=cfg.pad_byte) | |
| with torch.no_grad(): | |
| logits = model(x, mask, tier="base") | |
| sizes = {tier: model.parameter_count(tier) for tier in cfg.tiers} | |
| return { | |
| "logits_keys": len(logits), | |
| "tiers": len(cfg.tiers), | |
| **sizes, | |
| } | |
| if __name__ == "__main__": | |
| info = smoke_test() | |
| for k, v in info.items(): | |
| print(f"{k}: {v}") | |