infer.py

#!/usr/bin/env python3
"""Standalone inference for Blink-1. One file, no install beyond torch.

    pip install torch
    python infer.py                       # base, default prompts
    python infer.py --instruct            # instruct (ChatML) variant
    python infer.py --prompt "the " -n 120

Blink is a *byte* model: input and output are raw UTF-8 bytes (vocab 257 =
256 bytes + EOS=256). There is no tokenizer to download. It is 1,087
parameters, so coherent words are the exception, not the rule. That's the
point.

The model code below is the exact Blink architecture (a looped shared
transformer block with a per-loop LoRA and per-iteration embedding). It is
copied here on purpose so this file runs on its own — you do not need the
training repo to load a 13 KB model.
"""
from __future__ import annotations

import argparse
import math
import os
from dataclasses import dataclass, fields

import torch
from torch import Tensor, nn
from torch.nn import functional as F

EOS_TOKEN_ID = 256
# Blink is deployed at 8 loops through the shared block (the trained
# shared_loops=2 in the config is a floor; inference uses the full loop budget,
# exactly like the training repo's CLI). At 2 loops the model collapses to
# emitting spaces; at 8 it actually moves.
BLINK_LOOPS = 8


# --------------------------------------------------------------------------
# config (only the fields the checkpoint carries; extras are ignored on load)
# --------------------------------------------------------------------------
@dataclass
class ModelConfig:
    vocab_size: int = 257
    dim: int = 3
    n_heads: int = 1
    prelude_layers: int = 0
    coda_layers: int = 0
    shared_loops: int = 2
    lora_rank: int = 2
    ffn_hidden: int = 6
    max_seq_len: int = 4096
    max_context_len: int = 8192
    rope_base: float = 10_000_000.0
    index_dim: int = 2
    index_top_k: int = 16
    local_window: int = 16
    sparse_chunk_threshold: int = 4096
    sparse_chunk_queries: int = 1024
    ssa_block_size: int = 32
    ssa_top_k_blocks: int = 2
    attention_window: int = 256
    fp4_weights: bool = False
    thinking_enabled: bool = False
    thinking_steps_min: int = 0
    thinking_steps_max: int = 0
    thinking_loss_weight: float = 0.0


# --------------------------------------------------------------------------
# model (verbatim Blink architecture)
# --------------------------------------------------------------------------
def build_rope_cache(config: ModelConfig, length: int) -> tuple[Tensor, Tensor]:
    head_dim = config.dim // config.n_heads
    positions = torch.arange(length, dtype=torch.float32)
    inv_freq = 1.0 / (config.rope_base ** (torch.arange(0, head_dim, 2, dtype=torch.float32) / head_dim))
    angles = torch.outer(positions, inv_freq)
    return torch.cos(angles), torch.sin(angles)


def apply_rope(x: Tensor, cos: Tensor, sin: Tensor) -> Tensor:
    x_even = x[..., 0::2]
    x_odd = x[..., 1::2]
    rotated_even = x_even * cos - x_odd * sin
    rotated_odd = x_even * sin + x_odd * cos
    return torch.stack((rotated_even, rotated_odd), dim=-1).flatten(-2)


def causal_attention(q: Tensor, k: Tensor, v: Tensor, window: int) -> Tensor:
    """Causal (optionally sliding-window) attention.

    Folds Blink's two training-time backends — padded SDPA for short
    sequences and flex sliding-window attention once seq_len exceeds the
    window — into one masked softmax. Both use scale = 1/sqrt(head_dim) on
    the post-RoPE query dim, so the result is identical to either backend.
    """
    head_dim_qk = q.shape[-1]
    head_dim_v = v.shape[-1]
    scale = 1.0 / math.sqrt(head_dim_qk)
    scores = (q @ k.transpose(-2, -1)) * scale
    seq = q.shape[-2]
    qi = torch.arange(seq, device=q.device).view(seq, 1)
    ki = torch.arange(seq, device=q.device).view(1, seq)
    mask = ki <= qi
    if window > 0:
        mask = mask & (qi - ki < window)
    scores = scores.masked_fill(~mask, float("-inf"))
    return F.softmax(scores, dim=-1) @ v[..., :head_dim_v]


class SwiGlu(nn.Module):
    def __init__(self, dim: int, hidden: int) -> None:
        super().__init__()
        self.gate_up = nn.Linear(dim, 2 * hidden, bias=False)
        self.down = nn.Linear(hidden, dim, bias=False)

    def forward(self, x: Tensor) -> Tensor:
        gate, up = self.gate_up(x).chunk(2, dim=-1)
        return self.down(F.silu(gate) * up)


class Attention(nn.Module):
    def __init__(self, config: ModelConfig) -> None:
        super().__init__()
        self.config = config
        self.n_heads = config.n_heads
        self.head_dim = config.dim // config.n_heads
        self.qkv = nn.Linear(config.dim, 3 * config.dim, bias=False)
        self.out = nn.Linear(config.dim, config.dim, bias=False)

    def forward(self, x: Tensor, cos: Tensor, sin: Tensor, qkv_delta: Tensor | None) -> Tensor:
        batch, seq_len, dim = x.shape
        qkv = self.qkv(x)
        if qkv_delta is not None:
            qkv = qkv + qkv_delta
        q, k, v = qkv.split(dim, dim=-1)
        q = q.view(batch, seq_len, self.n_heads, self.head_dim).transpose(1, 2)
        k = k.view(batch, seq_len, self.n_heads, self.head_dim).transpose(1, 2)
        v = v.view(batch, seq_len, self.n_heads, self.head_dim).transpose(1, 2)
        q = apply_rope(q, cos, sin)
        k = apply_rope(k, cos, sin)
        window = self.config.attention_window
        use_window = 0 < window < seq_len
        attended = causal_attention(q, k, v, window if use_window else 0)
        merged = attended.transpose(1, 2).reshape(batch, seq_len, dim)
        return self.out(merged)


class Block(nn.Module):
    def __init__(self, config: ModelConfig) -> None:
        super().__init__()
        self.attn_norm = nn.RMSNorm(config.dim)
        self.attn = Attention(config)
        self.ffn_norm = nn.RMSNorm(config.dim)
        self.ffn = SwiGlu(config.dim, config.ffn_hidden)

    def forward(self, x: Tensor, cos: Tensor, sin: Tensor, qkv_delta: Tensor | None) -> Tensor:
        x = x + self.attn(self.attn_norm(x), cos, sin, qkv_delta)
        return x + self.ffn(self.ffn_norm(x))


class LoopLora(nn.Module):
    def __init__(self, config: ModelConfig, max_loops: int) -> None:
        super().__init__()
        self.down = nn.ModuleList(nn.Linear(config.dim, config.lora_rank, bias=False) for _ in range(max_loops))
        self.up = nn.ModuleList(nn.Linear(config.lora_rank, 3 * config.dim, bias=False) for _ in range(max_loops))

    def forward(self, x: Tensor, loop_index: int) -> Tensor:
        clamped = min(loop_index, len(self.down) - 1)
        return self.up[clamped](self.down[clamped](x))


class SparseIndexer(nn.Module):
    """Holds the SSA gate parameter. Never fires for Blink (short sequences,
    tiny dim), so it is load-only — present to match the released weights."""

    def __init__(self) -> None:
        super().__init__()
        self.gate = nn.Parameter(torch.tensor([0.1]))


class Blink(nn.Module):
    def __init__(self, config: ModelConfig, max_loops: int = 8) -> None:
        super().__init__()
        self.config = config
        self.max_loops = max_loops
        self.embed = nn.Embedding(config.vocab_size, config.dim)
        self.indexer = SparseIndexer()
        self.prelude = nn.ModuleList(Block(config) for _ in range(config.prelude_layers))
        self.shared = Block(config)
        self.loop_lora = LoopLora(config, max_loops)
        self.loop_embed = nn.Embedding(max_loops, config.dim)
        self.coda = nn.ModuleList(Block(config) for _ in range(config.coda_layers))
        self.final_norm = nn.RMSNorm(config.dim)

    @torch.no_grad()
    def forward(self, tokens: Tensor, loops: int | None = None) -> Tensor:
        loop_count = loops if loops is not None else self.config.shared_loops
        seq_len = tokens.shape[1]
        cos, sin = build_rope_cache(self.config, seq_len)
        cos, sin = cos.to(tokens.device), sin.to(tokens.device)
        x = self.embed(tokens)
        for block in self.prelude:
            x = block(x, cos, sin, None)
        for loop_index in range(loop_count):
            clamped = min(loop_index, self.max_loops - 1)
            gated = x + self.loop_embed.weight[clamped]
            delta = self.loop_lora(gated, loop_index)
            x = self.shared(gated, cos, sin, delta)
        for block in self.coda:
            x = block(x, cos, sin, None)
        x = self.final_norm(x)
        return F.linear(x, self.embed.weight)


# --------------------------------------------------------------------------
# loading + generation
# --------------------------------------------------------------------------
def load_model(checkpoint_path: str) -> tuple[Blink, ModelConfig]:
    payload = torch.load(checkpoint_path, map_location="cpu", weights_only=False)
    known = {f.name for f in fields(ModelConfig)}
    config = ModelConfig(**{k: v for k, v in payload["model_config"].items() if k in known})
    max_loops = payload["model"]["loop_embed.weight"].shape[0]
    model = Blink(config, max_loops=max_loops)
    model.load_state_dict(payload["model"])
    model.eval()
    return model, config


def sample_next_token(logits, temperature, top_k, repetition_penalty, recent, rng):
    if repetition_penalty != 1.0 and recent:
        logits = logits.clone()
        idx = torch.tensor(sorted(set(recent)), dtype=torch.long)
        sel = logits[idx]
        logits[idx] = torch.where(sel > 0, sel / repetition_penalty, sel * repetition_penalty)
    if temperature <= 0:
        return int(torch.argmax(logits))
    probs = F.softmax(logits / temperature, dim=-1)
    if top_k and top_k > 0:
        v, i = torch.topk(probs, min(top_k, probs.size(-1)))
        return int(i[torch.multinomial(v, 1, generator=rng)].item())
    return int(torch.multinomial(probs, 1, generator=rng).item())


def generate(model, prompt, max_new_tokens=80, temperature=0.5, top_k=5,
             repetition_penalty=1.1, repetition_window=128, seed=0):
    rng = torch.Generator()
    if seed is not None:
        rng.manual_seed(seed)
    tokens = list(prompt.encode("utf-8"))
    generated: list[int] = []
    for _ in range(max_new_tokens):
        x = torch.tensor([tokens[-model.config.max_context_len:]], dtype=torch.long)
        logits = model(x, loops=BLINK_LOOPS)[0, -1]
        nxt = sample_next_token(logits, temperature, top_k, repetition_penalty,
                                tokens[-repetition_window:], rng)
        if nxt == EOS_TOKEN_ID:
            break
        tokens.append(nxt)
        generated.append(nxt)
        if b"<|im_end|>" in bytes(generated[-60:]):
            break
    return bytes(generated).decode("utf-8", errors="replace").removesuffix("<|im_end|>")


CHATML = "<|im_start|>user\n{}<|im_end|>\n<|im_start|>assistant\n"
BASE_DECODE = dict(temperature=0.5, top_k=5, repetition_penalty=1.1)
INSTRUCT_DECODE = dict(temperature=0.2, top_k=10, repetition_penalty=1.5)


def main() -> None:
    ap = argparse.ArgumentParser(description="Run Blink-1 byte-level generation.")
    ap.add_argument("--instruct", action="store_true", help="use the instruct (ChatML) champion")
    ap.add_argument("--prompt", help="single prompt; with --instruct it is wrapped in ChatML")
    ap.add_argument("-n", "--max-new-tokens", type=int, default=80)
    ap.add_argument("--seed", type=int, default=0)
    args = ap.parse_args()

    here = os.path.dirname(os.path.abspath(__file__))
    ckpt = "blink-1-instruct.pt" if args.instruct else "blink-1-base.pt"
    model, config = load_model(os.path.join(here, ckpt))
    nparams = sum(p.numel() for p in model.parameters())
    print(f"loaded {ckpt}: {nparams} params, dim={config.dim}, shared_loops={config.shared_loops}\n")

    decode = INSTRUCT_DECODE if args.instruct else BASE_DECODE
    if args.prompt is not None:
        prompts = [CHATML.format(args.prompt) if args.instruct else args.prompt]
    elif args.instruct:
        prompts = [CHATML.format("hello"), CHATML.format("what is your name?")]
    else:
        prompts = ["the ", "I think ", "once upon "]

    for p in prompts:
        out = generate(model, p, max_new_tokens=args.max_new_tokens, seed=args.seed, **decode)
        print(f"prompt:  {p!r}")
        print(f"output:  {out!r}\n")


if __name__ == "__main__":
    main()