| """Glyph model definition — standalone, no dependencies beyond PyTorch.""" |
|
|
| import torch |
| import torch.nn as nn |
| import torch.nn.functional as F |
|
|
|
|
| class TrigramHashEmbedding(nn.Module): |
| def __init__(self, n_buckets=8192, d_embed=64, prime=31): |
| super().__init__() |
| self.n_buckets, self.prime = n_buckets, prime |
| self.embed = nn.Embedding(n_buckets, d_embed) |
| def forward(self, x): |
| xp = F.pad(x.long(), (2, 0), value=0) |
| h = (xp[:, :-2] * self.prime * self.prime + xp[:, 1:-1] * self.prime + xp[:, 2:]) % self.n_buckets |
| return self.embed(h) |
|
|
|
|
| class RoPE(nn.Module): |
| def __init__(self, head_dim, max_len=1024, theta=10000.0): |
| super().__init__() |
| inv_freq = 1.0 / (theta ** (torch.arange(0, head_dim, 2).float() / head_dim)) |
| self.register_buffer("inv_freq", inv_freq, persistent=False) |
| self._build(max_len) |
| def _build(self, seq_len): |
| t = torch.arange(seq_len, device=self.inv_freq.device).float() |
| freqs = torch.outer(t, self.inv_freq); emb = torch.cat([freqs, freqs], dim=-1) |
| self.register_buffer("cos", emb.cos()[None, None], persistent=False) |
| self.register_buffer("sin", emb.sin()[None, None], persistent=False); self._max = seq_len |
| @staticmethod |
| def _rotate(x): |
| x1, x2 = x.chunk(2, dim=-1); return torch.cat([-x2, x1], dim=-1) |
| def forward(self, q, k): |
| T = q.size(2) |
| if T > self._max: self._build(T) |
| c, s = self.cos[:,:,:T], self.sin[:,:,:T] |
| return q*c + self._rotate(q)*s, k*c + self._rotate(k)*s |
|
|
|
|
| class Attention(nn.Module): |
| def __init__(self, d_model, n_heads): |
| super().__init__() |
| self.n_heads, self.head_dim = n_heads, d_model // n_heads |
| self.qkv = nn.Linear(d_model, 3 * d_model); self.out = nn.Linear(d_model, d_model) |
| self.rope = RoPE(self.head_dim); self.norm = nn.LayerNorm(d_model) |
| def forward(self, x): |
| res = x; x = self.norm(x); B, T, C = x.shape |
| qkv = self.qkv(x).view(B, T, 3, self.n_heads, self.head_dim) |
| q, k, v = qkv.permute(2, 0, 3, 1, 4); q, k = self.rope(q, k) |
| out = F.scaled_dot_product_attention(q, k, v, is_causal=False) |
| return res + self.out(out.transpose(1, 2).contiguous().view(B, T, C)) |
|
|
|
|
| class ConvBlock(nn.Module): |
| def __init__(self, d_in, d_out, kernel=3): |
| super().__init__() |
| self.conv = nn.Conv1d(d_in, d_out, kernel, padding=kernel // 2) |
| self.bn = nn.BatchNorm1d(d_out) |
| self.residual = nn.Conv1d(d_in, d_out, 1) if d_in != d_out else nn.Identity() |
| def forward(self, x): return F.gelu(self.bn(self.conv(x))) + self.residual(x) |
|
|
|
|
| class MultiTaskLID(nn.Module): |
| """ |
| Glyph: Multi-task byte-level text classifier. |
| ~4M shared parameters + per-task classification heads. |
| """ |
| def __init__(self, task_configs, max_len=512, d_byte=64, d_tri=64, |
| n_buckets=8192, d_model=384, n_conv=4, n_attn=2, n_heads=6, dropout=0.0): |
| super().__init__() |
| self.max_len = max_len; self.d_model = d_model |
| self.byte_embed = nn.Embedding(256, d_byte) |
| self.tri_embed = TrigramHashEmbedding(n_buckets, d_tri) |
| self.proj = nn.Linear(d_byte + d_tri, d_model) |
| self.convs = nn.ModuleList([ConvBlock(d_model, d_model, 3) for _ in range(n_conv)]) |
| self.attns = nn.ModuleList([Attention(d_model, n_heads) for _ in range(n_attn)]) |
| self.drop = nn.Dropout(dropout); self.norm = nn.LayerNorm(d_model) |
| self.heads = nn.ModuleDict({t: nn.Linear(d_model, n) for t, n in task_configs.items()}) |
|
|
| def forward(self, x, task): |
| h = self.proj(torch.cat([self.byte_embed(x), self.tri_embed(x)], dim=-1)) |
| h = h.transpose(1, 2) |
| for conv in self.convs: h = conv(h) |
| h = h.transpose(1, 2) |
| for attn in self.attns: h = attn(h) |
| return {"logits": self.heads[task](self.drop(self.norm(h).mean(dim=1)))} |
|
|
