| """Tiny CED model architecture required to load model.safetensors.""" |
|
|
| import torch |
| import torch.nn.functional as F |
| from torch import nn |
|
|
|
|
| class RMSNorm(nn.Module): |
| def __init__(self, dim): |
| super().__init__() |
| self.weight = nn.Parameter(torch.ones(dim)) |
|
|
| def forward(self, x): |
| z = x.float() |
| return ( |
| z |
| * torch.rsqrt(z.square().mean(-1, keepdim=True) + 1e-5) |
| * self.weight |
| ).to(x.dtype) |
|
|
|
|
| def rope(x, pos): |
| d = x.shape[-1] |
| freq = 10000 ** (-torch.arange(0, d, 2, device=x.device) / d) |
| angle = pos[:, None, :, None].float() * freq |
| a, b = x.float()[..., ::2], x.float()[..., 1::2] |
| c, s = angle.cos(), angle.sin() |
| return torch.stack((a * c - b * s, a * s + b * c), dim=-1).flatten(-2).to(x.dtype) |
|
|
|
|
| class Attention(nn.Module): |
| def __init__(self, dim, heads, window): |
| super().__init__() |
| self.heads = heads |
| self.window = window |
| for name in ["wq", "wkg", "wvg", "wkl", "wvl", "wo"]: |
| setattr(self, name, nn.Linear(dim, dim, bias=False)) |
|
|
| def forward(self, u, source, pos, valid): |
| batch, length, dim = u.shape |
|
|
| def split(z): |
| return z.view(batch, length, self.heads, dim // self.heads).transpose(1, 2) |
|
|
| q = rope(split(self.wq(u)), pos) |
| kg, vg = rope(split(self.wkg(source)), pos), split(self.wvg(source)) |
| kl, vl = rope(split(self.wkl(u)), pos), split(self.wvl(u)) |
| distance = pos[:, :, None] - pos[:, None, :] |
| global_mask = (distance >= 0) & valid[:, :, None] & valid[:, None, :] |
| local_mask = global_mask & (distance < self.window) |
| mask = torch.cat((global_mask, local_mask), -1)[:, None] |
| k, v = torch.cat((kg, kl), 2), torch.cat((vg, vl), 2) |
| out = F.scaled_dot_product_attention( |
| q, |
| k, |
| v, |
| attn_mask=mask, |
| dropout_p=0, |
| is_causal=False, |
| ) |
| return self.wo(out.transpose(1, 2).contiguous().view(batch, length, dim)) |
|
|
|
|
| class Block(nn.Module): |
| def __init__(self, dim, heads, ff, window): |
| super().__init__() |
| self.attn_norm, self.ffn_norm = RMSNorm(dim), RMSNorm(dim) |
| self.attn = Attention(dim, heads, window) |
| self.gate, self.up = (nn.Linear(dim, ff, bias=False) for _ in range(2)) |
| self.down = nn.Linear(ff, dim, bias=False) |
|
|
| def forward(self, h, e, pos, valid): |
| u = self.attn_norm(h) |
| h = h + self.attn(u, u if e is None else e, pos, valid) |
| z = self.ffn_norm(h) |
| return h + self.down(F.silu(self.gate(z)) * self.up(z)) |
|
|
|
|
| class CED(nn.Module): |
| def __init__(self, vocab=8192, dim=384, heads=6, ff=1024, layers=4, window=64): |
| super().__init__() |
| assert dim % heads == 0 and (dim // heads) % 2 == 0 |
| self.embedding = nn.Embedding(vocab, dim) |
| self.encoder = nn.ModuleList( |
| [Block(dim, heads, ff, window) for _ in range(layers)] |
| ) |
| self.decoder = nn.ModuleList( |
| [Block(dim, heads, ff, window) for _ in range(layers)] |
| ) |
| self.encoder_norm, self.final_norm = RMSNorm(dim), RMSNorm(dim) |
| self.head = nn.Linear(dim, vocab, bias=False) |
| self.head.weight = self.embedding.weight |
| for parameter in self.parameters(): |
| if parameter.ndim > 1: |
| nn.init.normal_(parameter, std=0.02) |
|
|
| def forward(self, ids, positions=None, valid=None): |
| if positions is None: |
| positions = torch.arange(ids.shape[1], device=ids.device)[None].expand_as(ids) |
| if valid is None: |
| valid = torch.ones_like(ids, dtype=torch.bool) |
| h = self.embedding(ids) |
| for layer in self.encoder: |
| h = layer(h, None, positions, valid) |
| e = self.encoder_norm(h) |
| h = e |
| for layer in self.decoder: |
| h = layer(h, e, positions, valid) |
| return self.head(self.final_norm(h)) |
|
|
