Glint-Research/Router

1"""glint-router-1m: standalone inference. torch + tokenizers + safetensors, nothing else.2 3one forward pass over the prompt gives you every field at once: domain,4complexity, code, math, reasoning, long_output, route, and a 64-d projection5used for user-defined categories. no decoding loop, no output parsing.6"""7 8from __future__ import annotations9 10import json11import math12from dataclasses import dataclass13from pathlib import Path14 15import torch16from torch import Tensor, nn17from torch.nn import functional18 19HERE = Path(__file__).parent20PAD_ID = 021BOS_ID = 122BINARY_FIELDS = ("code", "math", "reasoning", "long_output")23COMPLEXITY_LEVELS = 524 25DOMAINS = (26    "programming", "web_dev", "databases", "devops_sysadmin", "security",27    "data_science_ml", "math", "science", "engineering", "reasoning_puzzle",28    "creative_writing", "professional_writing", "editing_grammar", "translation",29    "factual_qa", "education", "business", "finance", "legal", "health",30    "travel", "food", "entertainment", "lifestyle", "other",31)32 33 34@dataclass(frozen=True)35class RouterConfig:36    vocab_size: int = 409637    dim: int = 12838    n_heads: int = 839    layers: int = 340    ffn_hidden: int = 20841    max_len: int = 25642    rope_base: float = 10_000.043    proj_dim: int = 6444    n_domains: int = len(DOMAINS)45 46 47def build_rope_cache(config: RouterConfig) -> tuple[Tensor, Tensor]:48    head_dim = config.dim // config.n_heads49    positions = torch.arange(config.max_len, dtype=torch.float32)50    inv_freq = 1.0 / (config.rope_base ** (torch.arange(0, head_dim, 2).float() / head_dim))51    angles = torch.outer(positions, inv_freq)52    return torch.cos(angles), torch.sin(angles)53 54 55def apply_rope(x: Tensor, cos: Tensor, sin: Tensor) -> Tensor:56    x_even, x_odd = x[..., 0::2], x[..., 1::2]57    rotated_even = x_even * cos - x_odd * sin58    rotated_odd = x_even * sin + x_odd * cos59    return torch.stack((rotated_even, rotated_odd), dim=-1).flatten(-2)60 61 62class SwiGlu(nn.Module):63    def __init__(self, dim: int, hidden: int) -> None:64        super().__init__()65        self.gate_up = nn.Linear(dim, 2 * hidden, bias=False)66        self.down = nn.Linear(hidden, dim, bias=False)67 68    def forward(self, x: Tensor) -> Tensor:69        gate, up = self.gate_up(x).chunk(2, dim=-1)70        return self.down(functional.silu(gate) * up)71 72 73class RouterAttention(nn.Module):74    def __init__(self, config: RouterConfig) -> None:75        super().__init__()76        self.n_heads = config.n_heads77        self.head_dim = config.dim // config.n_heads78        self.qkv = nn.Linear(config.dim, 3 * config.dim, bias=False)79        self.out = nn.Linear(config.dim, config.dim, bias=False)80 81    def forward(self, x: Tensor, cos: Tensor, sin: Tensor, attn_mask: Tensor) -> Tensor:82        batch, seq_len, dim = x.shape83        q, k, v = self.qkv(x).split(dim, dim=-1)84        shape = (batch, seq_len, self.n_heads, self.head_dim)85        q = apply_rope(q.view(shape).transpose(1, 2), cos, sin)86        k = apply_rope(k.view(shape).transpose(1, 2), cos, sin)87        v = v.view(shape).transpose(1, 2)88        attended = functional.scaled_dot_product_attention(q, k, v, attn_mask=attn_mask)89        return self.out(attended.transpose(1, 2).reshape(batch, seq_len, dim))90 91 92class RouterBlock(nn.Module):93    def __init__(self, config: RouterConfig) -> None:94        super().__init__()95        self.attn_norm = nn.RMSNorm(config.dim)96        self.attn = RouterAttention(config)97        self.ffn_norm = nn.RMSNorm(config.dim)98        self.ffn = SwiGlu(config.dim, config.ffn_hidden)99 100    def forward(self, x: Tensor, cos: Tensor, sin: Tensor, attn_mask: Tensor) -> Tensor:101        x = x + self.attn(self.attn_norm(x), cos, sin, attn_mask)102        return x + self.ffn(self.ffn_norm(x))103 104 105class GlintRouter(nn.Module):106    def __init__(self, config: RouterConfig | None = None) -> None:107        super().__init__()108        config = config or RouterConfig()109        self.config = config110        self.embed = nn.Embedding(config.vocab_size, config.dim, padding_idx=PAD_ID)111        self.blocks = nn.ModuleList(RouterBlock(config) for _ in range(config.layers))112        self.final_norm = nn.RMSNorm(config.dim)113        cos, sin = build_rope_cache(config)114        self.register_buffer("rope_cos", cos, persistent=False)115        self.register_buffer("rope_sin", sin, persistent=False)116 117        pooled = 2 * config.dim118        self.domain_head = nn.Linear(pooled, config.n_domains)119        self.complexity_dir = nn.Linear(pooled, 1, bias=False)120        self.complexity_bias = nn.Parameter(torch.zeros(COMPLEXITY_LEVELS - 1))121        self.binary_head = nn.Linear(pooled, len(BINARY_FIELDS))122        self.route_head = nn.Linear(pooled, 1)123        self.proj_head = nn.Linear(pooled, config.proj_dim)124        self.register_buffer("temperature", torch.ones(3), persistent=True)125 126    def encode(self, tokens: Tensor) -> Tensor:127        valid = tokens != PAD_ID128        seq_len = tokens.shape[1]129        cos = self.rope_cos[:seq_len].to(self.embed.weight.dtype)130        sin = self.rope_sin[:seq_len].to(self.embed.weight.dtype)131        attn_mask = torch.zeros(tokens.shape, dtype=self.embed.weight.dtype,132                                device=tokens.device)133        attn_mask = attn_mask.masked_fill(~valid, float("-inf"))[:, None, None, :]134        x = self.embed(tokens)135        for block in self.blocks:136            x = block(x, cos, sin, attn_mask)137        x = self.final_norm(x)138        mask = valid.unsqueeze(-1)139        mean = (x * mask).sum(dim=1) / mask.sum(dim=1).clamp(min=1)140        maximum = torch.nan_to_num(x.masked_fill(~mask, float("-inf")).max(dim=1).values,141                                   neginf=0.0)142        return torch.cat((mean, maximum), dim=-1)143 144    def forward(self, tokens: Tensor) -> dict[str, Tensor]:145        pooled = self.encode(tokens)146        return {147            "domain": self.domain_head(pooled),148            "complexity": self.complexity_dir(pooled) + self.complexity_bias,149            "binary": self.binary_head(pooled),150            "route": self.route_head(pooled).squeeze(-1),151            "proj": functional.normalize(self.proj_head(pooled), dim=-1),152        }153 154    def calibrated(self, tokens: Tensor) -> dict[str, Tensor]:155        """probabilities with the fitted temperatures applied. the routing156        arithmetic in policy.py runs on these numbers, so they have to mean157        something. temperatures were fitted on held-out data after training."""158        out = self.forward(tokens)159        route_t, complexity_t, binary_t = self.temperature.unbind()160        return {161            "domain": out["domain"].softmax(dim=-1),162            "complexity": (out["complexity"] / complexity_t).sigmoid(),163            "binary": (out["binary"] / binary_t).sigmoid(),164            "route": (out["route"] / route_t).sigmoid(),165            "proj": out["proj"],166        }167 168 169def complexity_from_cumulative(probabilities: Tensor) -> Tensor:170    """coral decode. level = 1 + how many thresholds the prompt clears."""171    return 1 + (probabilities > 0.5).sum(dim=-1)172 173 174def encode_batch(tokenizer, texts: list[str], max_len: int) -> Tensor:175    """bos-prefixed, right-padded. long prompts lose their tail, because the176    instruction verb lives at the front and the pasted context lives at the back."""177    out = torch.full((len(texts), max_len), PAD_ID, dtype=torch.long)178    for row, text in enumerate(texts):179        ids = [BOS_ID] + tokenizer.encode(text).ids[: max_len - 1]180        out[row, : len(ids)] = torch.tensor(ids, dtype=torch.long)181    return out182 183 184def load_router(directory: Path = HERE, device: str = "cpu"):185    """returns (model, tokenizer). reads config.json + model.safetensors + tokenizer.json."""186    from safetensors.torch import load_file187    from tokenizers import Tokenizer188 189    directory = Path(directory)190    config = RouterConfig(**json.loads((directory / "config.json").read_text())["model"])191    model = GlintRouter(config).to(device)192    model.load_state_dict(load_file(directory / "model.safetensors", device=device))193    model.eval()194    return model, Tokenizer.from_file(str(directory / "tokenizer.json"))195 196 197def count_parameters(model: nn.Module) -> int:198    return sum(p.numel() for p in model.parameters())199