File size: 6,035 Bytes
c6f3061 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 | 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)
# Training projects only supervised positions, reducing vocabulary memory.
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
# Training and inference use the same meaning of t: actual masked fraction.
noise = selected.sum(-1).float() / valid.sum(-1)
return clean.masked_fill(selected, MASK), selected, noise
|