| from dataclasses import dataclass, asdict |
| import json |
| import math |
| from pathlib import Path |
|
|
| import torch |
| from torch import nn |
| from torch.nn import functional as F |
| from diffuthink.model import Block, RMSNorm |
|
|
| PAD, BOS, EOS, MASK, UNK = range(5) |
| SPECIAL_TOKENS = ["[PAD]", "[BOS]", "[EOS]", "[MASK]", "[UNK]"] |
|
|
|
|
| @dataclass |
| class Config: |
| vocab_size: int = 8192 |
| width: int = 320 |
| heads: int = 8 |
| layers: int = 6 |
| max_length: int = 192 |
| architecture: str = "diffuthink_v2" |
|
|
| def __post_init__(self): |
| if min(self.width, self.heads, self.layers, self.max_length) < 1: |
| raise ValueError("Dimensions must be positive") |
| if self.width % self.heads or (self.width // self.heads) % 2: |
| raise ValueError("RoPE head dimension must be even") |
| if self.vocab_size <= len(SPECIAL_TOKENS): |
| raise ValueError("Vocabulary too small") |
|
|
|
|
| class Denoiser(nn.Module): |
| """Original small Transformer, initialized randomly, with tied embeddings. |
| |
| Learned absolute positions supplement RoPE so all-masked sequences have |
| position-specific information. No external model weights are used. |
| """ |
| def __init__(self, config): |
| super().__init__() |
| self.config = config |
| self.embedding = nn.Embedding(config.vocab_size, config.width) |
| self.position = nn.Embedding(config.max_length, config.width) |
| self.time_mlp = nn.Sequential(nn.Linear(config.width, config.width), nn.SiLU(), nn.Linear(config.width, config.width)) |
| self.blocks = nn.ModuleList(Block(config) for _ in range(config.layers)) |
| self.norm = RMSNorm(config.width) |
| self.apply(self._init) |
| for block in self.blocks: |
| nn.init.normal_(block.out.weight, std=0.02 / math.sqrt(2 * config.layers)) |
| nn.init.normal_(block.down.weight, std=0.02 / math.sqrt(2 * config.layers)) |
|
|
| @staticmethod |
| def _init(module): |
| if isinstance(module, (nn.Linear, nn.Embedding)): |
| nn.init.normal_(module.weight, std=0.02) |
| if isinstance(module, nn.Linear) and module.bias is not None: |
| nn.init.zeros_(module.bias) |
|
|
| def forward(self, tokens, noise, selected=None, causal=False, right_padded=False): |
| if tokens.ndim != 2 or tokens.shape[1] > self.config.max_length: |
| raise ValueError("Sequence exceeds configured context") |
| half = self.config.width // 2 |
| freq = torch.exp(-math.log(10000) * torch.arange(half, device=tokens.device) / max(half - 1, 1)) |
| angle = noise[:, None] * 1000 * freq |
| t = self.time_mlp(torch.cat((angle.sin(), angle.cos()), -1)) |
| x = self.embedding(tokens) + self.position(torch.arange(tokens.shape[1], device=tokens.device)) |
| for block in self.blocks: |
| x = block(x, t, None if causal and right_padded else tokens.ne(PAD), causal=causal) |
| x = self.norm(x) |
| |
| if selected is not None: |
| x = x[selected] |
| return F.linear(x, self.embedding.weight) |
|
|
| def save_pretrained(self, directory): |
| from safetensors.torch import save_file |
| directory = Path(directory) |
| directory.mkdir(parents=True, exist_ok=True) |
| (directory / "config.json").write_text(json.dumps(asdict(self.config), indent=2), encoding="utf-8") |
| temporary = directory / "model.safetensors.tmp" |
| save_file({k: v.detach().cpu().contiguous() for k, v in self.state_dict().items()}, str(temporary)) |
| temporary.replace(directory / "model.safetensors") |
|
|
| def extend_context(self, length): |
| """Keep learned positions exactly; initialize new positions near their mean.""" |
| if length < self.config.max_length: |
| raise ValueError("Cannot shrink context during extension") |
| if length == self.config.max_length: |
| return |
| old = self.position.weight.detach() |
| position = nn.Embedding(length, self.config.width, device=old.device, dtype=old.dtype) |
| with torch.no_grad(): |
| nn.init.normal_(position.weight, std=0.02) |
| position.weight.add_(old.mean(0)) |
| position.weight[:len(old)].copy_(old) |
| self.position = position |
| self.config.max_length = length |
|
|
| @classmethod |
| def from_pretrained(cls, directory, device="cpu"): |
| from safetensors.torch import load_file |
| directory = Path(directory) |
| model = cls(Config(**json.loads((directory / "config.json").read_text(encoding="utf-8")))) |
| model.load_state_dict(load_file(str(directory / "model.safetensors"))) |
| return model.to(device).eval() |
|
|
|
|
| def corrupt(clean, rng, mode="mixed", rate=None): |
| """Mix random holes and continuation masks. BOS/PAD are never targets.""" |
| valid = clean.ne(PAD) & clean.ne(BOS) |
| b, n = clean.shape |
| if not valid.any(-1).all(): |
| raise ValueError("All examples need a target") |
| noise = 0.05 + 0.95 * torch.rand(b, device=clean.device, generator=rng) if rate is None else torch.full((b,), rate, device=clean.device) |
| selected = (torch.rand(clean.shape, device=clean.device, generator=rng) < noise[:, None]) & valid |
| if mode == "mixed": |
| suffix_rows = torch.rand(b, device=clean.device, generator=rng) < 0.5 |
| lengths = valid.sum(-1) |
| starts = 1 + (torch.rand(b, device=clean.device, generator=rng) * lengths).long() |
| suffix = (torch.arange(n, device=clean.device)[None] >= starts[:, None]) & valid |
| selected = torch.where(suffix_rows[:, None], suffix, selected) |
| elif mode != "random": |
| raise ValueError("Unknown corruption mode") |
| empty = ~selected.any(-1) |
| fallback = torch.rand(clean.shape, device=clean.device, generator=rng).masked_fill(~valid, -1).argmax(-1) |
| selected[empty, fallback[empty]] = True |
| |
| noise = selected.sum(-1).float() / valid.sum(-1) |
| return clean.masked_fill(selected, MASK), selected, noise |
|
|
