wiola 0.1.0__py3-none-any.whl

wiola/__init__.py

@@ -0,0 +1,47 @@
+# coding=utf-8
+"""Wiola: a small language model with SRPE, GCLA, ATM, DSFF and WiolaRMSNorm.
+
+Importing this package registers Wiola with the HuggingFace Auto* classes so
+that ``AutoModelForCausalLM.from_pretrained("oscowlai/wiola-360m")`` works once
+weights are published.
+"""
+
+from .configuration_wiola import WiolaConfig
+from .modeling_wiola import (
+    WiolaDecoderLayer,
+    WiolaForCausalLM,
+    WiolaModel,
+    WiolaPreTrainedModel,
+)
+
+__version__ = "0.1.0"
+
+__all__ = [
+    "WiolaConfig",
+    "WiolaModel",
+    "WiolaForCausalLM",
+    "WiolaPreTrainedModel",
+    "WiolaDecoderLayer",
+]
+
+
+def _register_auto_classes():
+    try:
+        from transformers import AutoConfig, AutoModel, AutoModelForCausalLM
+    except Exception:  # transformers not installed
+        return
+    try:
+        AutoConfig.register("wiola", WiolaConfig)
+        AutoModel.register(WiolaConfig, WiolaModel)
+        AutoModelForCausalLM.register(WiolaConfig, WiolaForCausalLM)
+    except Exception:
+        # Already registered (e.g. re-import) — safe to ignore.
+        pass
+
+
+# Make the config/model loadable from the Hub with trust_remote_code too.
+WiolaConfig.register_for_auto_class("AutoConfig")
+WiolaModel.register_for_auto_class("AutoModel")
+WiolaForCausalLM.register_for_auto_class("AutoModelForCausalLM")
+
+_register_auto_classes()

wiola/components/__init__.py

@@ -0,0 +1,21 @@
+# coding=utf-8
+"""Wiola architectural components."""
+
+from .atm import merge_ratio, merge_tokens, unmerge_tokens
+from .dsff import DualStreamFeedForward
+from .gcla import GatedCrossLayerAttention, repeat_kv
+from .normalization import WiolaRMSNorm
+from .srpe import SpiralRotaryEmbedding, apply_srpe, rotate_half
+
+__all__ = [
+    "WiolaRMSNorm",
+    "SpiralRotaryEmbedding",
+    "apply_srpe",
+    "rotate_half",
+    "GatedCrossLayerAttention",
+    "repeat_kv",
+    "merge_tokens",
+    "unmerge_tokens",
+    "merge_ratio",
+    "DualStreamFeedForward",
+]

wiola/components/atm.py

@@ -0,0 +1,99 @@
+# coding=utf-8
+"""Adaptive Token Merging (ATM).
+
+Implements Section VIII of the Wiola paper. Adjacent tokens whose cosine
+similarity exceeds ``tau`` are greedily merged (averaged) in a left-to-right,
+non-overlapping scan. A merge map records the source positions so the original
+sequence length can be restored exactly after attention.
+
+ATM is **training-only**: it is disabled at inference to keep the KV-cache
+length consistent with the un-merged sequence.
+
+Because sequences in a batch may merge by different amounts, merged sequences
+are right-padded to the batch maximum and a boolean ``keep_mask`` marks the real
+(non-padded) merged positions. Unmerge ignores padded slots.
+"""
+
+from typing import List, Tuple
+
+import torch
+
+
+def _greedy_merge_one(sim: torch.Tensor, threshold: float) -> List[Tuple[int, ...]]:
+    """Greedy non-overlapping merge decisions for a single sequence.
+
+    Args:
+        sim: cosine similarities between adjacent tokens, shape [T-1].
+        threshold: tau.
+    Returns:
+        A list of groups; each group is a tuple of 1 or 2 source indices.
+    """
+    seq_len = sim.shape[0] + 1
+    groups: List[Tuple[int, ...]] = []
+    i = 0
+    while i < seq_len:
+        if i < seq_len - 1 and sim[i].item() > threshold:
+            groups.append((i, i + 1))
+            i += 2
+        else:
+            groups.append((i,))
+            i += 1
+    return groups
+
+
+def merge_tokens(hidden_states: torch.Tensor, threshold: float):
+    """Merge adjacent redundant tokens.
+
+    Args:
+        hidden_states: [B, T, d].
+        threshold: tau (cosine similarity merge threshold).
+    Returns:
+        merged: [B, T_prime_max, d] (right padded with zeros).
+        keep_mask: [B, T_prime_max] bool, True for real merged tokens.
+        merge_maps: list (len B) of lists of source-index tuples.
+    """
+    bsz, seq_len, dim = hidden_states.shape
+    if seq_len < 2:
+        keep_mask = torch.ones(bsz, seq_len, dtype=torch.bool, device=hidden_states.device)
+        merge_maps = [[(t,) for t in range(seq_len)] for _ in range(bsz)]
+        return hidden_states, keep_mask, merge_maps
+
+    normed = torch.nn.functional.normalize(hidden_states, dim=-1, eps=1e-8)
+    # rho_t = <x_hat_t, x_hat_{t+1}>  -> [B, T-1]
+    sim = (normed[:, :-1] * normed[:, 1:]).sum(-1)
+
+    merge_maps = [_greedy_merge_one(sim[b], threshold) for b in range(bsz)]
+    new_len = max(len(g) for g in merge_maps)
+
+    merged = hidden_states.new_zeros(bsz, new_len, dim)
+    keep_mask = torch.zeros(bsz, new_len, dtype=torch.bool, device=hidden_states.device)
+    for b, groups in enumerate(merge_maps):
+        for k, grp in enumerate(groups):
+            if len(grp) == 2:
+                merged[b, k] = 0.5 * (hidden_states[b, grp[0]] + hidden_states[b, grp[1]])
+            else:
+                merged[b, k] = hidden_states[b, grp[0]]
+            keep_mask[b, k] = True
+    return merged, keep_mask, merge_maps
+
+
+def unmerge_tokens(merged: torch.Tensor,
+                   merge_maps: List[List[Tuple[int, ...]]],
+                   original_len: int) -> torch.Tensor:
+    """Restore original sequence length by broadcasting each merged token back
+    to its source positions (Eq. 31)."""
+    bsz, _, dim = merged.shape
+    out = merged.new_zeros(bsz, original_len, dim)
+    for b, groups in enumerate(merge_maps):
+        for k, grp in enumerate(groups):
+            for src in grp:
+                out[b, src] = merged[b, k]
+    return out
+
+
+def merge_ratio(merge_maps: List[List[Tuple[int, ...]]], original_len: int) -> float:
+    """Average merge ratio mu = 1 - T'/T across the batch."""
+    if original_len == 0:
+        return 0.0
+    ratios = [1.0 - len(g) / original_len for g in merge_maps]
+    return float(sum(ratios) / len(ratios))

wiola/components/dsff.py

@@ -0,0 +1,38 @@
+# coding=utf-8
+"""Dual-Stream Feed-Forward (DSFF).
+
+Implements Section IX of the Wiola paper. Two parallel *dense* streams of
+different widths and activations are fused by a learned per-dimension gate:
+
+    Stream A (narrow, SwiGLU):  a = D_A( SiLU(G_A x) * (U_A x) )
+    Stream B (wide,   GELU):    b = D_B( GELU(U_B x) )
+    gate:                       alpha = sigmoid(W_f [a; b])  in (0,1)^d
+    output:                     alpha * a + (1 - alpha) * b
+
+Setting ``W_f = 0`` yields ``alpha = 0.5``, reducing DSFF to a simple ensemble
+average; DSFF therefore strictly generalises a two-stream ensemble.
+"""
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+class DualStreamFeedForward(nn.Module):
+    def __init__(self, hidden_size: int, narrow_size: int, wide_size: int):
+        super().__init__()
+        # Stream A: narrow SwiGLU.
+        self.gate_a = nn.Linear(hidden_size, narrow_size, bias=False)   # G_A
+        self.up_a = nn.Linear(hidden_size, narrow_size, bias=False)     # U_A
+        self.down_a = nn.Linear(narrow_size, hidden_size, bias=False)   # D_A
+        # Stream B: wide GELU.
+        self.up_b = nn.Linear(hidden_size, wide_size, bias=False)       # U_B
+        self.down_b = nn.Linear(wide_size, hidden_size, bias=False)     # D_B
+        # Per-dimension fusion gate from concatenated stream outputs.
+        self.fusion = nn.Linear(2 * hidden_size, hidden_size, bias=False)  # W_f
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        a = self.down_a(F.silu(self.gate_a(x)) * self.up_a(x))
+        b = self.down_b(F.gelu(self.up_b(x)))
+        alpha = torch.sigmoid(self.fusion(torch.cat((a, b), dim=-1)))
+        return alpha * a + (1.0 - alpha) * b

wiola/components/gcla.py

@@ -0,0 +1,149 @@
+# coding=utf-8
+"""Gated Cross-Layer Attention (GCLA).
+
+Implements Section VII of the Wiola paper. Each decoder layer performs:
+
+  * GQA self-attention with SRPE on the local sequence, and
+  * cross-attention to compressed summaries of up to ``Lambda`` preceding
+    layers (supplied by the model as ``context_summaries``),
+
+blended by a scalar gate ``beta = sigmoid(phi)`` and modulated by a sigmoid
+output gate ``G``:
+
+    O   = (1 - beta) * O_self + beta * O_ctx
+    A   = (G * concat(O)) W_O
+
+The context tensor is provided *per query position* as a causal cumulative mean
+of prior-layer outputs (see :class:`WiolaModel`), so a cached incremental decode
+reproduces a full forward pass exactly.
+"""
+
+import math
+from typing import Optional, Tuple
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from .srpe import SpiralRotaryEmbedding, apply_srpe
+
+
+def repeat_kv(x: torch.Tensor, n_rep: int) -> torch.Tensor:
+    """Expand [B, H_kv, T, d] -> [B, H_kv*n_rep, T, d] (GQA)."""
+    if n_rep == 1:
+        return x
+    b, kv, t, d = x.shape
+    x = x[:, :, None, :, :].expand(b, kv, n_rep, t, d)
+    return x.reshape(b, kv * n_rep, t, d)
+
+
+class GatedCrossLayerAttention(nn.Module):
+    def __init__(self, config, layer_idx: int):
+        super().__init__()
+        self.layer_idx = layer_idx
+        self.hidden_size = config.hidden_size
+        self.num_heads = config.num_attention_heads
+        self.num_kv_heads = config.num_key_value_heads
+        self.head_dim = config.head_dim
+        self.n_rep = self.num_heads // self.num_kv_heads
+        self.lookback = config.gcla_lookback
+
+        q_dim = self.num_heads * self.head_dim
+        kv_dim = self.num_kv_heads * self.head_dim
+
+        self.q_proj = nn.Linear(self.hidden_size, q_dim, bias=False)
+        self.k_proj = nn.Linear(self.hidden_size, kv_dim, bias=False)
+        self.v_proj = nn.Linear(self.hidden_size, kv_dim, bias=False)
+        self.o_proj = nn.Linear(q_dim, self.hidden_size, bias=False)
+
+        # Cross-layer context projections.
+        self.k_ctx_proj = nn.Linear(self.hidden_size, kv_dim, bias=False)
+        self.v_ctx_proj = nn.Linear(self.hidden_size, kv_dim, bias=False)
+
+        # Sigmoid output gate.
+        self.gate_proj = nn.Linear(self.hidden_size, q_dim, bias=False)
+
+        # Scalar blend gate beta = sigmoid(phi).
+        self.blend_logit = nn.Parameter(torch.tensor(float(config.gcla_gate_init)))
+
+        self.srpe = SpiralRotaryEmbedding(
+            head_dim=self.head_dim,
+            max_position_embeddings=config.max_position_embeddings,
+            theta=config.srpe_theta,
+            spiral_divisor=config.srpe_spiral_divisor,
+            radial_amplitude=config.srpe_radial_amplitude,
+            radial_frequency=config.srpe_radial_frequency,
+        )
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,                       # [B, T, d]
+        position_ids: torch.Tensor,                        # [B, T]
+        attention_mask: Optional[torch.Tensor],            # [B, 1, T, T_k] additive
+        context_summaries: Optional[torch.Tensor] = None,  # [B, T, Lambda, d]
+        past_key_value: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
+        use_cache: bool = False,
+    ):
+        bsz, q_len, _ = hidden_states.shape
+
+        q = self.q_proj(hidden_states)
+        k = self.k_proj(hidden_states)
+        v = self.v_proj(hidden_states)
+
+        q = q.view(bsz, q_len, self.num_heads, self.head_dim).transpose(1, 2)
+        k = k.view(bsz, q_len, self.num_kv_heads, self.head_dim).transpose(1, 2)
+        v = v.view(bsz, q_len, self.num_kv_heads, self.head_dim).transpose(1, 2)
+
+        # SRPE rotation on q/k (per head).
+        cos, sin = self.srpe(hidden_states, position_ids)
+        q, k = apply_srpe(q, k, cos, sin)
+
+        # Concatenate cached KV (incremental decoding).
+        # Concatenate cached KV (incremental decoding).
+        if (
+            past_key_value is not None
+            and past_key_value[0] is not None      # ← guard against placeholder None
+            and past_key_value[1] is not None
+        ):
+            past_k, past_v = past_key_value
+            k = torch.cat([past_k, k], dim=2)
+            v = torch.cat([past_v, v], dim=2)
+        present = (k, v) if use_cache else None
+
+        # GQA expansion.
+        k_rep = repeat_kv(k, self.n_rep)
+        v_rep = repeat_kv(v, self.n_rep)
+
+        scale = 1.0 / math.sqrt(self.head_dim)
+        scores = torch.matmul(q, k_rep.transpose(-1, -2)) * scale  # [B,H,T,T_k]
+        if attention_mask is not None:
+            scores = scores + attention_mask
+        attn = F.softmax(scores, dim=-1, dtype=torch.float32).to(q.dtype)
+        o_self = torch.matmul(attn, v_rep)                         # [B,H,T,dh]
+
+        # Cross-layer context attention.
+        if context_summaries is not None and context_summaries.shape[2] > 0:
+            lam = context_summaries.shape[2]
+            k_ctx = self.k_ctx_proj(context_summaries)  # [B,T,Lam,kv_dim]
+            v_ctx = self.v_ctx_proj(context_summaries)
+            k_ctx = k_ctx.view(bsz, q_len, lam, self.num_kv_heads, self.head_dim)
+            v_ctx = v_ctx.view(bsz, q_len, lam, self.num_kv_heads, self.head_dim)
+            # Expand kv heads to full head count.
+            k_ctx = k_ctx.repeat_interleave(self.n_rep, dim=3)  # [B,T,Lam,H,dh]
+            v_ctx = v_ctx.repeat_interleave(self.n_rep, dim=3)
+            # scores[b,h,t,l] = q[b,h,t,:] . k_ctx[b,t,l,h,:]
+            ctx_scores = torch.einsum("bhtd,btlhd->bhtl", q, k_ctx) * scale
+            ctx_attn = F.softmax(ctx_scores, dim=-1, dtype=torch.float32).to(q.dtype)
+            o_ctx = torch.einsum("bhtl,btlhd->bhtd", ctx_attn, v_ctx)
+            beta = torch.sigmoid(self.blend_logit)
+            o = (1.0 - beta) * o_self + beta * o_ctx
+        else:
+            o = o_self
+
+        # Merge heads.
+        o = o.transpose(1, 2).contiguous().view(bsz, q_len, self.num_heads * self.head_dim)
+
+        # Sigmoid output gate.
+        g = torch.sigmoid(self.gate_proj(hidden_states))
+        o = self.o_proj(g * o)
+        return o, present

wiola/components/normalization.py

@@ -0,0 +1,33 @@
+# coding=utf-8
+"""WiolaRMSNorm: RMS normalisation with a learned per-dimension input offset.
+
+Implements Eq. (3) of the Wiola paper:
+
+    WiolaRMSNorm(x) = gamma * (x + delta) / sqrt(mean((x + delta)^2) + eps)
+
+Setting ``delta = 0`` recovers standard RMSNorm exactly, so this strictly
+generalises RMSNorm. The offset shifts the *input before* normalisation,
+changing the normalisation target itself rather than adding a post-norm bias.
+"""
+
+import torch
+import torch.nn as nn
+
+
+class WiolaRMSNorm(nn.Module):
+    def __init__(self, hidden_size: int, eps: float = 1e-6):
+        super().__init__()
+        self.weight = nn.Parameter(torch.ones(hidden_size))   # gamma
+        self.offset = nn.Parameter(torch.zeros(hidden_size))  # delta
+        self.variance_epsilon = eps
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        input_dtype = hidden_states.dtype
+        # Compute the normalisation in fp32 for numerical stability.
+        z = hidden_states.to(torch.float32) + self.offset.to(torch.float32)
+        variance = z.pow(2).mean(-1, keepdim=True)
+        z = z * torch.rsqrt(variance + self.variance_epsilon)
+        return (self.weight.to(torch.float32) * z).to(input_dtype)
+
+    def extra_repr(self) -> str:
+        return f"{tuple(self.weight.shape)}, eps={self.variance_epsilon}"

wiola/components/srpe.py

@@ -0,0 +1,117 @@
+# coding=utf-8
+"""Spiral Rotary Positional Encoding (SRPE).
+
+Implements Section VI of the Wiola paper. For position ``p`` and dimension-pair
+index ``j in [d_h/2]``:
+
+    omega_j   = theta_0 ** (-2j / d_h)
+    Theta_j(p)= p * omega_j * (1 + 1/k_s)
+    r_j(p)    = 1 + a_s * sin(p * f_s * omega_j)
+    c_j(p)    = r_j(p) * cos(Theta_j(p))
+    s_j(p)    = r_j(p) * sin(Theta_j(p))
+
+and the rotation is applied to q/k with the standard rotate-half trick, which
+is algebraically identical to Eqs. (12)-(13):
+
+    out_lo = q_lo * c - q_hi * s
+    out_hi = q_hi * c + q_lo * s
+
+Unlike RoPE, ``Theta`` carries a second winding angle (via ``1 + 1/k_s``) and
+the magnitude is modulated by a radial term ``r_j(p)``, embedding positions on a
+3D helical manifold with no learned parameters.
+"""
+
+import torch
+import torch.nn as nn
+
+
+def rotate_half(x: torch.Tensor) -> torch.Tensor:
+    """Rotate the two halves of the last dimension: [x1, x2] -> [-x2, x1]."""
+    half = x.shape[-1] // 2
+    x1 = x[..., :half]
+    x2 = x[..., half:]
+    return torch.cat((-x2, x1), dim=-1)
+
+
+class SpiralRotaryEmbedding(nn.Module):
+    def __init__(
+        self,
+        head_dim: int,
+        max_position_embeddings: int = 2048,
+        theta: float = 10000.0,
+        spiral_divisor: int = 8,
+        radial_amplitude: float = 0.1,
+        radial_frequency: float = 0.01,
+    ):
+        super().__init__()
+        if head_dim % 2 != 0:
+            raise ValueError(f"head_dim must be even, got {head_dim}.")
+        self.head_dim = head_dim
+        self.max_position_embeddings = max_position_embeddings
+        self.theta = theta
+        self.k_s = spiral_divisor
+        self.a_s = radial_amplitude
+        self.f_s = radial_frequency
+
+        # omega_j = theta ** (-2j / d_h), j in [0, d_h/2).
+        j = torch.arange(0, head_dim, 2, dtype=torch.float32)
+        inv_freq = theta ** (-(j / head_dim))          # shape [d_h/2]
+        self.register_buffer("inv_freq", inv_freq, persistent=False)
+
+        self._cached_len = 0
+        self._build_cache(max_position_embeddings)
+
+    def _build_cache(self, seq_len: int, device=None, dtype=torch.float32):
+        positions = torch.arange(seq_len, dtype=torch.float32,
+                                 device=device if device is not None else self.inv_freq.device)
+        inv_freq = self.inv_freq.to(positions.device)
+        # Combined winding angle Theta_j(p) = p * omega_j * (1 + 1/k_s).
+        combined = 1.0 + 1.0 / self.k_s
+        angle = torch.outer(positions, inv_freq) * combined        # [T, d_h/2]
+        # Radial modulation r_j(p) = 1 + a_s * sin(p * f_s * omega_j).
+        radial = 1.0 + self.a_s * torch.sin(
+            torch.outer(positions, inv_freq) * self.f_s
+        )
+        cos = radial * torch.cos(angle)
+        sin = radial * torch.sin(angle)
+        # Duplicate to full head_dim so it lines up with rotate_half.
+        cos = torch.cat((cos, cos), dim=-1)                        # [T, d_h]
+        sin = torch.cat((sin, sin), dim=-1)
+        self.register_buffer("cos_cached", cos.to(dtype), persistent=False)
+        self.register_buffer("sin_cached", sin.to(dtype), persistent=False)
+        self._cached_len = seq_len
+
+    @torch.no_grad()
+    def _maybe_extend(self, seq_len: int, device, dtype):
+        if seq_len > self._cached_len or self.cos_cached.device != device:
+            self._build_cache(max(seq_len, self._cached_len), device=device, dtype=dtype)
+
+    def forward(self, x: torch.Tensor, position_ids: torch.Tensor):
+        """Return (cos, sin) gathered for ``position_ids``.
+
+        Args:
+            x: any tensor only used to read dtype/device.
+            position_ids: LongTensor of shape [batch, seq_len].
+        Returns:
+            cos, sin each of shape [batch, seq_len, head_dim].
+        """
+        max_pos = int(position_ids.max().item()) + 1
+        self._maybe_extend(max_pos, x.device, torch.float32)
+        cos = self.cos_cached[position_ids]   # [B, S, d_h]
+        sin = self.sin_cached[position_ids]
+        return cos.to(x.dtype), sin.to(x.dtype)
+
+
+def apply_srpe(q: torch.Tensor, k: torch.Tensor,
+               cos: torch.Tensor, sin: torch.Tensor):
+    """Apply SRPE rotation to query and key tensors.
+
+    Args:
+        q, k: shape [batch, num_heads, seq_len, head_dim].
+        cos, sin: shape [batch, seq_len, head_dim].
+    """
+    cos = cos.unsqueeze(1)  # [B, 1, S, d_h] -> broadcast over heads
+    sin = sin.unsqueeze(1)
+    q_rot = (q * cos) + (rotate_half(q) * sin)
+    k_rot = (k * cos) + (rotate_half(k) * sin)
+    return q_rot, k_rot

wiola/configuration_wiola.py

@@ -0,0 +1,146 @@
+# coding=utf-8
+# Copyright 2025 The Wiola / OSCOWL-AI authors.
+# Licensed under the Apache License, Version 2.0 (the "License").
+"""Wiola model configuration."""
+
+from transformers.configuration_utils import PretrainedConfig
+
+
+class WiolaConfig(PretrainedConfig):
+    r"""
+    Configuration class for a :class:`WiolaForCausalLM` model.
+
+    This stores every hyper-parameter described in the Wiola paper. Defaults
+    correspond to the **wiola-360m** variant. The four published sizes are
+    available as YAML files under ``configs/`` and as named presets via
+    :meth:`WiolaConfig.from_preset`.
+
+    Args:
+        vocab_size (int): Vocabulary size of the BPE tokenizer.
+        hidden_size (int): Model hidden dimension ``d``.
+        num_hidden_layers (int): Number of decoder layers ``L``.
+        num_attention_heads (int): Number of query heads ``H``.
+        num_key_value_heads (int): Number of key/value heads ``H_kv`` (GQA).
+        max_position_embeddings (int): Maximum context length ``T``.
+        dsff_narrow_size (int): DSFF Stream A width ``d_A``.
+        dsff_wide_size (int): DSFF Stream B width ``d_B``.
+        srpe_theta (float): SRPE base theta ``theta_0``.
+        srpe_spiral_divisor (int): SRPE spiral divisor ``k_s``.
+        srpe_radial_amplitude (float): SRPE radial amplitude ``a_s``.
+        srpe_radial_frequency (float): SRPE radial frequency ``f_s``.
+        atm_threshold (float): ATM cosine-similarity merge threshold ``tau``.
+        atm_enabled (bool): Master switch for Adaptive Token Merging in training.
+        gcla_lookback (int): GCLA lookback depth ``Lambda``.
+        gcla_gate_init (float): Logit ``phi`` used to initialise the blend gate
+            ``beta = sigmoid(phi)``.
+        rms_norm_eps (float): Epsilon for WiolaRMSNorm.
+        initializer_range (float): Stddev for truncated-normal init.
+        tie_word_embeddings (bool): Tie input embedding and LM head.
+    """
+
+    model_type = "wiola"
+    keys_to_ignore_at_inference = ["past_key_values"]
+
+    def __init__(
+        self,
+        vocab_size: int = 32000,
+        hidden_size: int = 1024,
+        num_hidden_layers: int = 16,
+        num_attention_heads: int = 16,
+        num_key_value_heads: int = 4,
+        max_position_embeddings: int = 2048,
+        dsff_narrow_size: int = 1024,
+        dsff_wide_size: int = 4096,
+        srpe_theta: float = 10000.0,
+        srpe_spiral_divisor: int = 8,
+        srpe_radial_amplitude: float = 0.1,
+        srpe_radial_frequency: float = 0.01,
+        atm_threshold: float = 0.92,
+        atm_enabled: bool = True,
+        gcla_lookback: int = 2,
+        gcla_gate_init: float = -3.0,
+        rms_norm_eps: float = 1e-6,
+        initializer_range: float = 0.02,
+        use_cache: bool = True,
+        pad_token_id: int = 0,
+        bos_token_id: int = 1,
+        eos_token_id: int = 2,
+        tie_word_embeddings: bool = True,
+        **kwargs,
+    ):
+        self.vocab_size = vocab_size
+        self.hidden_size = hidden_size
+        self.num_hidden_layers = num_hidden_layers
+        self.num_attention_heads = num_attention_heads
+        # GQA: default num_key_value_heads to num_attention_heads (MHA) when unset.
+        if num_key_value_heads is None:
+            num_key_value_heads = num_attention_heads
+        self.num_key_value_heads = num_key_value_heads
+        self.max_position_embeddings = max_position_embeddings
+
+        self.dsff_narrow_size = dsff_narrow_size
+        self.dsff_wide_size = dsff_wide_size
+
+        self.srpe_theta = srpe_theta
+        self.srpe_spiral_divisor = srpe_spiral_divisor
+        self.srpe_radial_amplitude = srpe_radial_amplitude
+        self.srpe_radial_frequency = srpe_radial_frequency
+
+        self.atm_threshold = atm_threshold
+        self.atm_enabled = atm_enabled
+
+        self.gcla_lookback = gcla_lookback
+        self.gcla_gate_init = gcla_gate_init
+
+        self.rms_norm_eps = rms_norm_eps
+        self.initializer_range = initializer_range
+        self.use_cache = use_cache
+
+        if hidden_size % num_attention_heads != 0:
+            raise ValueError(
+                f"hidden_size ({hidden_size}) must be divisible by "
+                f"num_attention_heads ({num_attention_heads})."
+            )
+        if num_attention_heads % num_key_value_heads != 0:
+            raise ValueError(
+                f"num_attention_heads ({num_attention_heads}) must be divisible "
+                f"by num_key_value_heads ({num_key_value_heads})."
+            )
+
+        super().__init__(
+            pad_token_id=pad_token_id,
+            bos_token_id=bos_token_id,
+            eos_token_id=eos_token_id,
+            tie_word_embeddings=tie_word_embeddings,
+            **kwargs,
+        )
+
+    @property
+    def head_dim(self) -> int:
+        return self.hidden_size // self.num_attention_heads
+
+    # Convenience presets matching the paper's model family. -----------------
+    _PRESETS = {
+        "wiola-120m": dict(hidden_size=768, num_hidden_layers=12,
+                           num_attention_heads=12, num_key_value_heads=4,
+                           dsff_narrow_size=768, dsff_wide_size=3072),
+        "wiola-360m": dict(hidden_size=1024, num_hidden_layers=16,
+                           num_attention_heads=16, num_key_value_heads=4,
+                           dsff_narrow_size=1024, dsff_wide_size=4096),
+        "wiola-700m": dict(hidden_size=1536, num_hidden_layers=24,
+                           num_attention_heads=16, num_key_value_heads=8,
+                           dsff_narrow_size=1536, dsff_wide_size=6144),
+        "wiola-1.5b": dict(hidden_size=2048, num_hidden_layers=28,
+                           num_attention_heads=16, num_key_value_heads=8,
+                           dsff_narrow_size=2048, dsff_wide_size=8192),
+    }
+
+    @classmethod
+    def from_preset(cls, name: str, **overrides) -> "WiolaConfig":
+        if name not in cls._PRESETS:
+            raise KeyError(
+                f"Unknown preset '{name}'. Choose from {list(cls._PRESETS)}."
+            )
+        params = dict(cls._PRESETS[name])
+        params.update(overrides)
+        return cls(**params)

wiola/data/__init__.py

@@ -0,0 +1,22 @@
+# coding=utf-8
+"""Wiola data pipeline."""
+
+from .dataset import (
+    SPECIAL_TOKENS,
+    PackedDatasetConfig,
+    PackedTextDataset,
+    TextSource,
+    build_tokenizer,
+    causal_collate,
+    load_or_build_tokenizer,
+)
+
+__all__ = [
+    "TextSource",
+    "PackedTextDataset",
+    "PackedDatasetConfig",
+    "causal_collate",
+    "build_tokenizer",
+    "load_or_build_tokenizer",
+    "SPECIAL_TOKENS",
+]