"""Architecture dimensions of the fixed inference export."""
from dataclasses import dataclass


@dataclass(frozen=True)
class ModelConfig:
    vocab_size: int = 2048
    width: int = 384
    heads: int = 6
    mlp_width: int = 2096
    loops: int = 6
    rank: int = 128
    controller_width: int = 16
    max_seq_len: int = 8192
    norm_eps: float = 1e-06
    routing_eps: float = 1.1920928955078125e-07

    def validate(self):
        if self != ModelConfig():
            raise ValueError("This runtime supports the released Purrence architecture only")


"""Six-loop Attention Residuals inference with per-head XSA and gates."""
import torch
from torch import nn
import torch.nn.functional as F


def rms(x, eps):
    work = x if x.dtype == torch.float64 else x.float()
    return (work * torch.rsqrt(work.square().mean(-1, keepdim=True) + eps)).to(x.dtype)


def depth_read(sources, query, eps=2 ** -23):
    values = torch.stack(sources)
    work = values if values.dtype == torch.float64 else values.float()
    inverse_rms = torch.rsqrt(work.square().mean(-1) + eps)
    scores = (work[..., -query.numel():] * query.to(work.dtype)).sum(-1) * inverse_rms
    weights = scores.softmax(dim=0)
    return (weights[..., None] * work).sum(0).to(values.dtype)


def document_layout(ids, boundary_id):
    flat = ids.reshape(-1)
    starts = flat.eq(boundary_id)
    starts[0] = True
    if ids.ndim == 2:
        starts[::ids.shape[1]] = True
    offsets = starts.nonzero().flatten()
    cu = torch.cat((offsets, offsets.new_tensor([flat.numel()]))).to(torch.int32)
    positions = torch.arange(flat.numel(), device=ids.device) - torch.repeat_interleave(offsets, cu.diff().long())
    return cu, positions


def half_rope(x, positions, frequencies, factors=None):
    rotary_dim = x.shape[-1] // 2
    if factors is None:
        angle = positions.float()[:, None] * frequencies[None]
        cosine, sine = angle.cos().to(x.dtype)[:, None], angle.sin().to(x.dtype)[:, None]
    else:
        cosine, sine = (factor.to(x.dtype) for factor in factors)
    even, odd = x[..., :rotary_dim:2], x[..., 1:rotary_dim:2]
    rotated = torch.stack((even * cosine + odd * sine, odd * cosine - even * sine), dim=-1)
    return torch.cat((rotated.flatten(-2), x[..., rotary_dim:]), dim=-1)


def offset_keys(keys, positions):
    rotary_dim = keys.shape[-1] // 2
    previous = torch.cat((keys[:1, :, rotary_dim:], keys[:-1, :, rotary_dim:]))
    stationary = torch.where(positions[:, None, None] > 0, previous, keys[..., rotary_dim:])
    return torch.cat((keys[..., :rotary_dim], stationary), dim=-1)


def attention(q, k, v, cu):
    result = []
    boundaries = cu.tolist()
    for start, end in zip(boundaries, boundaries[1:]):
        query, key, value = (item[start:end].transpose(0, 1)[None] for item in (q, k, v))
        out = F.scaled_dot_product_attention(query, key, value, is_causal=True)
        result.append(out[0].transpose(0, 1))
    return torch.cat(result)


class SharedBlock(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.config = config
        self.q = nn.Linear(config.width, config.width, bias=False)
        self.k = nn.Linear(config.width, config.width, bias=False)
        self.v = nn.Linear(config.width, config.width, bias=False)
        self.o = nn.Linear(config.width, config.width, bias=False)
        self.up = nn.Linear(config.width, config.mlp_width, bias=False)
        self.down = nn.Linear(config.mlp_width, config.width, bias=False)
        self.xsa = nn.Linear(config.controller_width, config.heads, bias=False)
        self.gate = nn.Linear(config.controller_width, config.heads, bias=False)
        head_dim = config.width // config.heads
        self.register_buffer("frequencies", (1 / 1024) ** torch.linspace(0, 1, head_dim // 4), persistent=False)

    def rotary_factors(self, positions, dtype):
        angle = positions.float()[:, None] * self.frequencies[None]
        return angle.cos().to(dtype)[:, None], angle.sin().to(dtype)[:, None]

    def attend(self, x, cu, positions, rope_factors):
        config = self.config
        q, k, v = [F.linear(x, getattr(self, name).weight.to(x.dtype)).view(
            -1, config.heads, config.width // config.heads) for name in ('q', 'k', 'v')]
        q = half_rope(rms(q, config.routing_eps), positions, self.frequencies, rope_factors)
        k = half_rope(rms(k, config.routing_eps), positions, self.frequencies, rope_factors)
        k = offset_keys(k, positions)
        y = attention(q, k, v, cu)
        local = x[..., :config.controller_width]
        alpha, gate = self.xsa(local), self.gate(local)
        dot = (y * v).sum(-1, keepdim=True)
        denominator = v.square().sum(-1, keepdim=True).clamp_min(1e-8)
        y = y - alpha.tanh().to(y.dtype)[..., None] * (dot / denominator) * v
        y = y * gate.sigmoid().to(y.dtype)[..., None]
        return F.linear(y.contiguous().flatten(-2), self.o.weight.to(y.dtype))

    def mlp(self, x):
        return self.down(self.up(x).relu().square())


class RecurrentLM(nn.Module):
    def __init__(self, config=None):
        super().__init__()
        self.config = config or ModelConfig()
        self.config.validate()
        self.embedding = nn.Embedding(self.config.vocab_size, self.config.width)
        self.block = SharedBlock(self.config)
        self.queries = nn.Parameter(torch.zeros(2 * self.config.loops, self.config.rank))

    def read(self, sources, query):
        x = sources[0] if query is None else depth_read(sources, query, self.config.routing_eps)
        return rms(x, self.config.norm_eps)

    def forward(self, ids, cu_seqlens=None, positions=None):
        if ids.ndim not in (1, 2) or ids.numel() == 0:
            raise ValueError("Expected nonempty one- or two-dimensional token IDs")
        if cu_seqlens is None:
            cu_seqlens, positions = document_layout(ids, self.config.vocab_size - 1)
        if positions is None:
            raise ValueError("Explicit packed boundaries require explicit document positions")
        if int(cu_seqlens.diff().max()) > self.config.max_seq_len:
            raise ValueError("Document segment exceeds max_seq_len")
        embedding = self.embedding(ids.reshape(-1))
        rope_factors = self.block.rotary_factors(positions, embedding.dtype)
        sources = [embedding]
        for loop in range(self.config.loops):
            x = self.read(sources, None if loop == 0 else self.queries[2 * loop - 1])
            a = self.block.attend(x, cu_seqlens, positions, rope_factors)
            x = self.read(sources + [a], self.queries[2 * loop])
            m = self.block.mlp(x)
            sources.append(a + m)
        x = self.read(sources, self.queries[-1])
        return F.linear(x, self.embedding.weight).view(*ids.shape, self.config.vocab_size)

    def parameter_count(self):
        return sum(parameter.numel() for parameter in self.parameters())


#!/usr/bin/env python3
"""Small runtime for the reviewed 2048-ID tokenizer and structural boundaries."""
import hashlib
import json
from pathlib import Path


class Sub3MTokenizer:
    vocab_size = 2048
    boundary_id = 2047

    def __init__(self, artifact_dir):
        import tokenizers
        directory = Path(artifact_dir)
        metadata = json.loads((directory / "artifact.json").read_text(encoding="utf-8"))
        if tokenizers.__version__ != metadata["tokenizers_version"]:
            raise RuntimeError("Use tokenizers==" + metadata["tokenizers_version"])
        path = directory / "ordinary-tokenizer.json"
        if hashlib.sha256(path.read_bytes()).hexdigest() != metadata["tokenizer_sha256"]:
            raise ValueError("Tokenizer file does not match its artifact hash")
        self.backend = tokenizers.Tokenizer.from_file(str(path))
        if self.backend.get_vocab_size() != 2047 or metadata["vocab_size"] != 2048:
            raise ValueError("Expected 2047 ordinary IDs and one structural boundary")
        if json.loads(path.read_text(encoding="utf-8")).get("added_tokens"):
            raise ValueError("Ordinary encoder must not recognize structural token spellings")
        self.metadata = metadata

    def encode(self, text):
        """Encode literal text; even '<|endoftext|>' remains ordinary content."""
        if not isinstance(text, str):
            raise TypeError("encode expects a Unicode string")
        ids = self.backend.encode(text, add_special_tokens=False).ids
        if any(token_id >= self.boundary_id for token_id in ids):
            raise ValueError("Ordinary encoding emitted a reserved ID")
        return ids

    def decode(self, ids, skip_boundaries=False):
        """Decode generated IDs; optionally omit structural boundary markers.

        Arbitrary generated byte sequences can be invalid UTF-8; the ByteLevel
        decoder replaces invalid sequences. Every encoded valid string roundtrips.
        """
        chunks = []
        ordinary = []
        for token_id in ids:
            if type(token_id) is not int or token_id not in range(self.vocab_size):
                raise ValueError("Token ID must be an integer in [0, 2048)")
            if token_id == self.boundary_id:
                chunks.append(self.backend.decode(ordinary, skip_special_tokens=False))
                ordinary = []
                if not skip_boundaries:
                    chunks.append(self.metadata["boundary_special"]["spelling"])
            else:
                ordinary.append(token_id)
        chunks.append(self.backend.decode(ordinary, skip_special_tokens=False))
        return "".join(chunks)