oruk-loop-153M / modeling_oruk_loop.py

oruk-loop-153M: weights, inference code, model card, licenses, results, media

58b0065 verified 1 day ago

22 kB

"""oruk-loop-153M: a double-looped Qwen3-ASR student.
The Qwen3-ASR architecture is kept, but each of its two transformer stacks is replaced by ONE shared block that is
applied repeatedly:
* audio encoder: one block applied 24 times, made depth-aware by a rank-64 adapter on every linear layer plus
  per-depth biases and LayerNorm gains;
* text decoder: one block applied 28 times, with rank-16 adapters and per-depth RMSNorm gains;
* token table: the kept English vocabulary (36,974 rows); 368 frequent rows are stored exactly and the rest are
  factorised at rank 729. The same table is used for input embeddings and the (tied) output layer.
Token ids stay in the original Qwen3 id space, so the stock Qwen3-ASR processor (tokenizer, chat template,
feature extractor) works unchanged.
Usage:
    from transformers import AutoModel, AutoProcessor
    model = AutoModel.from_pretrained("oruk/oruk-loop-153M", trust_remote_code=True)
    processor = AutoProcessor.from_pretrained("oruk/oruk-loop-153M")
    print(model.transcribe("speech.wav", processor))
"""
from __future__ import annotations

import copy
import math
import os
from pathlib import Path

import numpy as np
import torch
import torch.nn.functional as F
from torch import nn
from transformers import AutoConfig, GenerationConfig, Qwen3ASRForConditionalGeneration

SAMPLE_RATE = 16000
PAD_ID = 151645
EOS_IDS = (151643, 151645)


# ---------------------------------------------------------------------------------------------- looped modules
class DepthCursor:
    """Which virtual depth the shared block is computing (set by each call site before the block runs)."""

    def __init__(self, n: int):
        self.n = n
        self.depth = 0


class LoopLinear(nn.Module):
    """y = W x + b[d] + B[d] (A[d] x): one shared weight, per-depth bias and per-depth rank-r adapter."""

    def __init__(self, cursor: DepthCursor, n: int, in_features: int, out_features: int, rank: int, bias: bool):
        super().__init__()
        object.__setattr__(self, "cursor", cursor)
        self.in_features, self.out_features, self.rank = in_features, out_features, rank
        self.weight = nn.Parameter(torch.zeros(out_features, in_features))
        self.bias = nn.Parameter(torch.zeros(n, out_features)) if bias else None
        if rank > 0:
            self.lora_A = nn.Parameter(torch.zeros(n, rank, in_features))
            self.lora_B = nn.Parameter(torch.zeros(n, out_features, rank))
        object.__setattr__(self, "_merged", None)
        object.__setattr__(self, "_merged_bias", None)

    @torch.no_grad()
    def merge(self, dtype=torch.bfloat16):
        """Materialise W + B[d] A[d] for every depth (faster decoding; ~3.5 GB in bf16 for the whole model)."""
        n = self.lora_A.shape[0] if self.rank > 0 else self.cursor.n
        W = torch.empty((n, self.out_features, self.in_features), dtype=dtype, device=self.weight.device)
        for i in range(n):
            w = self.weight.float()
            if self.rank > 0:
                w = w + self.lora_B[i].float() @ self.lora_A[i].float()
            W[i] = w.to(dtype)
        object.__setattr__(self, "_merged", W)
        object.__setattr__(self, "_merged_bias", None if self.bias is None else self.bias.detach().to(dtype))

    def unmerge(self):
        object.__setattr__(self, "_merged", None)
        object.__setattr__(self, "_merged_bias", None)

    def forward(self, x):
        d = self.cursor.depth
        if self._merged is not None:
            b = None if self._merged_bias is None else self._merged_bias[d]
            return F.linear(x.to(self._merged.dtype), self._merged[d], b)
        y = F.linear(x, self.weight, None if self.bias is None else self.bias[d])
        if self.rank > 0:
            y = y + F.linear(F.linear(x, self.lora_A[d]), self.lora_B[d])
        return y


class LoopLayerNorm(nn.Module):
    def __init__(self, cursor, n, shape, eps):
        super().__init__()
        object.__setattr__(self, "cursor", cursor)
        self.shape, self.eps = tuple(shape), eps
        self.weight = nn.Parameter(torch.ones(n, *self.shape))
        self.bias = nn.Parameter(torch.zeros(n, *self.shape))

    def forward(self, x):
        d = self.cursor.depth
        w, b = self.weight[d], self.bias[d]
        if not torch.is_autocast_enabled(x.device.type):  # under autocast layer_norm runs in fp32 with fp32 gains
            w, b = w.to(x.dtype), b.to(x.dtype)
        return F.layer_norm(x, self.shape, w, b, self.eps)


class LoopRMSNorm(nn.Module):
    def __init__(self, cursor, n, dim, eps):
        super().__init__()
        object.__setattr__(self, "cursor", cursor)
        self.weight = nn.Parameter(torch.ones(n, dim))
        self.variance_epsilon = eps

    def forward(self, x):
        d = self.cursor.depth
        dtype = x.dtype
        h = x.to(torch.float32)
        h = h * torch.rsqrt(h.pow(2).mean(-1, keepdim=True) + self.variance_epsilon)
        return (self.weight[d] * h.to(dtype)).to(dtype)


def build_loop_block(template: nn.Module, cursor: DepthCursor, n: int, rank: int) -> nn.Module:
    """Turn one transformer layer into the shared looped block (same module paths, depth-aware leaves)."""
    block = copy.deepcopy(template)
    for name, mod in list(block.named_modules()):
        if not name:
            continue
        if isinstance(mod, nn.Linear):
            new = LoopLinear(cursor, n, mod.in_features, mod.out_features, rank, mod.bias is not None)
        elif isinstance(mod, nn.LayerNorm):
            new = LoopLayerNorm(cursor, n, mod.normalized_shape, mod.eps)
        elif type(mod).__name__.endswith("RMSNorm"):
            new = LoopRMSNorm(cursor, n, mod.weight.shape[0], mod.variance_epsilon)
        else:
            continue
        parent_name, _, child = name.rpartition(".")
        setattr(block.get_submodule(parent_name) if parent_name else block, child, new)
    return block


class LoopCall(nn.Module):
    """Parameter-free call site for virtual depth `index` of a shared block (one per original layer)."""

    def __init__(self, block, cursor, index, kv_slot=None):
        super().__init__()
        object.__setattr__(self, "_block", block)
        object.__setattr__(self, "_cursor", cursor)
        self.index, self.kv_slot = index, kv_slot

    def forward(self, hidden_states, *args, **kwargs):
        self._cursor.depth = self.index
        if self.kv_slot is not None:  # each virtual decoder depth keeps its own KV-cache slot
            self._block.self_attn.layer_idx = self.kv_slot
        return self._block(hidden_states, *args, **kwargs)


# ---------------------------------------------------------------------------------------------- token table
class TokenTable(nn.Module):
    """Tied table over the kept vocabulary: exact frequent rows + rank-r factorised rare rows (U @ P)."""

    def __init__(self, n_full, dim, n_kept, n_exact, rank):
        super().__init__()
        self.n_full, self.embedding_dim, self.rank = n_full, dim, rank
        self.register_buffer("keep", torch.zeros(n_kept, dtype=torch.long))
        self.register_buffer("lut", torch.zeros(n_full, dtype=torch.long))
        self.register_buffer("exact_idx", torch.zeros(n_exact, dtype=torch.long))
        self.register_buffer("rare_idx", torch.zeros(n_kept - n_exact, dtype=torch.long))
        self.exact_rows = nn.Parameter(torch.zeros(n_exact, dim, dtype=torch.bfloat16), requires_grad=False)
        self.U = nn.Parameter(torch.zeros(n_kept - n_exact, rank))
        self.P = nn.Parameter(torch.zeros(rank, dim))
        object.__setattr__(self, "_merged", None)

    def weight(self, dtype=None):
        dtype = dtype or self.U.dtype
        if self._merged is not None:
            return self._merged if self._merged.dtype == dtype else self._merged.to(dtype)
        W = torch.empty(self.keep.shape[0], self.embedding_dim, dtype=dtype, device=self.U.device)
        W = W.index_copy(0, self.exact_idx, self.exact_rows.to(dtype))
        return W.index_copy(0, self.rare_idx, (self.U.float() @ self.P.float()).to(dtype))

    @torch.no_grad()
    def merge(self, dtype=torch.bfloat16):
        object.__setattr__(self, "_merged", None)
        object.__setattr__(self, "_merged", self.weight(dtype).detach())

    def forward(self, ids):
        # read the merged table in its own dtype (bf16 on GPU), as evaluated; never re-cast it per call
        return F.embedding(self.lut[ids], self._merged if self._merged is not None else self.weight())


class TableEmbedding(nn.Module):
    def __init__(self, table):
        super().__init__()
        object.__setattr__(self, "table", table)
        self.embedding_dim = table.embedding_dim

    @property
    def num_embeddings(self):
        return self.table.n_full

    @property
    def weight(self):
        return self.table.weight()

    def forward(self, ids):
        return self.table(ids)


class TableHead(nn.Module):
    """Output layer tied to the token table; full-width logits with -inf at ids outside the kept vocabulary."""

    def __init__(self, table):
        super().__init__()
        object.__setattr__(self, "table", table)
        self.in_features, self.out_features = table.embedding_dim, table.n_full

    @property
    def weight(self):
        return self.table.weight()

    def forward(self, h):
        kept = F.linear(h, self.table.weight(h.dtype))
        out = torch.full(h.shape[:-1] + (self.table.n_full,), float("-inf"), dtype=kept.dtype, device=kept.device)
        return out.index_copy_(-1, self.table.keep, kept)


# ---------------------------------------------------------------------------------------------- audio helpers
def load_audio(source, sampling_rate=None):
    """path / (array, sr) / 1-D array -> float32 mono at 16 kHz."""
    if isinstance(source, (str, os.PathLike)):
        import soundfile as sf
        wav, sr = sf.read(str(source), dtype="float32", always_2d=False)
    elif isinstance(source, tuple):
        wav, sr = source
    else:
        wav, sr = source, sampling_rate or SAMPLE_RATE
    wav = np.asarray(wav, dtype=np.float32)
    if wav.ndim > 1:
        wav = wav.mean(axis=1 if wav.shape[0] > wav.shape[1] else 0)
    if sr != SAMPLE_RATE:
        try:
            import soxr
            wav = soxr.resample(wav, sr, SAMPLE_RATE, quality="HQ").astype(np.float32)
        except ImportError:
            from math import gcd
            from scipy.signal import resample_poly
            g = gcd(int(sr), SAMPLE_RATE)
            wav = resample_poly(wav, SAMPLE_RATE // g, int(sr) // g).astype(np.float32)
    # the evaluation pipeline stores 16-bit PCM; quantise the same way so results reproduce exactly
    return np.clip(np.rint(wav * 32767.0), -32768, 32767).astype(np.float32) / 32768.0


def split_points(w, sr=SAMPLE_RATE, max_s=30.0, search_s=5.0, frame=0.1):
    """Chunk boundaries for audio longer than max_s: each cut sits at the quietest 100 ms frame within the last
    search_s seconds before the limit, so words are rarely split."""
    n, cuts, start = len(w), [], 0
    L, S, Fr = int(max_s * sr), int(search_s * sr), int(frame * sr)
    while n - start > L:
        lo, hi = start + L - S, start + L
        seg = w[lo:hi].astype(np.float32)
        k = len(seg) // Fr
        e = (seg[:k * Fr].reshape(k, Fr) ** 2).mean(1)
        cut = lo + int(np.argmin(e)) * Fr + Fr // 2 if k else hi
        cuts.append(cut)
        start = cut
    return [0] + cuts + [n]


# ---------------------------------------------------------------------------------------------- model
class OrukLoopForConditionalGeneration(Qwen3ASRForConditionalGeneration):
    """Qwen3-ASR with both stacks looped. `from_pretrained` loads the 153M stored parameters; see `transcribe`."""

    _tied_weights_keys = {}

    def __init__(self, config):
        loop = dict(getattr(config, "oruk_loop", {}))
        tc, ac = config.text_config, config.audio_config
        n_dec, n_enc, vocab = tc.num_hidden_layers, ac.encoder_layers, tc.vocab_size
        types = list(tc.layer_types)
        # build the stock skeleton with ONE template layer per stack and a tiny vocabulary (cheap), then restore
        tc.num_hidden_layers, tc.layer_types, tc.vocab_size, ac.encoder_layers = 1, types[:1], 8, 1
        try:
            super().__init__(config)
        finally:
            tc.num_hidden_layers, tc.layer_types, tc.vocab_size, ac.encoder_layers = n_dec, types, vocab, n_enc
        at, lm = self.model.audio_tower, self.model.language_model
        for c in {id(at.config): at.config}.values():  # sub-modules may hold copies of the sub-configs
            c.encoder_layers = n_enc
        lm.config.num_hidden_layers, lm.config.layer_types, lm.config.vocab_size = n_dec, types, vocab

        self.enc_cursor, self.dec_cursor = DepthCursor(n_enc), DepthCursor(n_dec)
        at.loop = build_loop_block(at.layers[0], self.enc_cursor, n_enc, int(loop.get("encoder_rank", 64)))
        at.layers = nn.ModuleList(LoopCall(at.loop, self.enc_cursor, i) for i in range(n_enc))
        lm.loop = build_loop_block(lm.layers[0], self.dec_cursor, n_dec, int(loop.get("decoder_rank", 16)))
        lm.layers = nn.ModuleList(LoopCall(lm.loop, self.dec_cursor, i, kv_slot=i) for i in range(n_dec))
        lm.token_table = TokenTable(vocab, tc.hidden_size, int(loop.get("table_rows", 36974)),
                                    int(loop.get("table_exact", 368)), int(loop.get("table_rank", 729)))
        lm.embed_tokens = TableEmbedding(lm.token_table)
        self.lm_head = TableHead(lm.token_table)
        self._source = None

    # the table is tied by construction; never let transformers re-tie or resize it
    def tie_weights(self, *args, **kwargs):
        return None

    def tie_embeddings_and_encoder_decoder(self, *args, **kwargs):
        return None

    def num_stored_parameters(self):
        return sum(p.numel() for p in self.parameters())

    def merge_loops(self, dtype=torch.bfloat16):
        for m in self.modules():
            if isinstance(m, LoopLinear):
                m.merge(dtype)

    def unmerge_loops(self):
        for m in self.modules():
            if isinstance(m, LoopLinear):
                m.unmerge()

    @classmethod
    def from_pretrained(cls, pretrained_model_name_or_path, *model_args, config=None, dtype=None, torch_dtype=None,
                        device=None, device_map=None, merge=None, revision=None, cache_dir=None, token=None,
                        local_files_only=False, **kwargs):
        """Load the release weights.
        device: "cuda", "cpu", ... (default: cuda if available). dtype: compute dtype (default bf16 on GPU, fp32 on CPU).
        merge: materialise per-depth weights for fast decoding (default: True on GPU, False on CPU).
        """
        ignored = [k for k in kwargs if not k.startswith("_") and k not in ("trust_remote_code", "low_cpu_mem_usage",
                                                                            "attn_implementation", "use_safetensors")]
        if ignored:
            import warnings
            warnings.warn(f"oruk-loop from_pretrained ignores these arguments: {ignored}")
        src = str(pretrained_model_name_or_path)
        if os.path.isdir(src):
            path = Path(src)
        else:
            from huggingface_hub import snapshot_download
            path = Path(snapshot_download(src, revision=revision, cache_dir=cache_dir, token=token,
                                          local_files_only=local_files_only,
                                          allow_patterns=["*.json", "*.safetensors", "*.py", "*.jinja", "*.txt"]))
        if config is None:
            config = AutoConfig.from_pretrained(str(path))
        if device is None:
            device = device_map if isinstance(device_map, str) and device_map != "auto" else None
        device = torch.device(device or ("cuda" if torch.cuda.is_available() else "cpu"))
        dtype = dtype or torch_dtype
        if isinstance(dtype, str) and dtype != "auto":
            dtype = getattr(torch, dtype)
        if dtype is None or dtype == "auto":
            dtype = torch.bfloat16 if device.type == "cuda" else torch.float32
        if merge is None:
            merge = device.type == "cuda"

        model = cls(config)
        from safetensors.torch import load_file
        state = load_file(str(path / "model.safetensors"))
        missing, unexpected = model.load_state_dict(state, strict=False)
        if missing or unexpected:
            raise RuntimeError(f"oruk-loop weights do not match the model: missing={missing[:8]} "
                               f"unexpected={unexpected[:8]}")
        del state
        try:
            model.generation_config = GenerationConfig.from_pretrained(str(path))
        except Exception:
            pass
        model.to(device)
        model.model.language_model.token_table.merge(dtype)
        if merge:
            model.merge_loops(dtype)
        if dtype == torch.bfloat16 and device.type == "cuda":
            # evaluated precision: modules kept from Qwen3-ASR in bf16, looped weights merged to bf16, per-depth norms
            # and biases in fp32, everything run under bf16 autocast
            at, lm = model.model.audio_tower, model.model.language_model
            for mod in (at.conv2d1, at.conv2d2, at.conv2d3, at.conv_out, at.ln_post, model.model.multi_modal_projector,
                        lm.norm):
                mod.to(torch.bfloat16)
            model._autocast = torch.bfloat16
        elif dtype != torch.float32:
            model.to(dtype)
        model.eval()
        model._source = str(path)
        return model

    _autocast = None

    def generate(self, *args, **kwargs):
        if self._autocast is not None and not torch.is_autocast_enabled(self.device.type):
            with torch.autocast(self.device.type, dtype=self._autocast):
                return super().generate(*args, **kwargs)
        return super().generate(*args, **kwargs)

    # ------------------------------------------------------------------------------------------ transcription
    @torch.no_grad()
    def transcribe(self, audio, processor=None, *, sampling_rate=None, batch_size=16, num_beams=1, language="English",
                   length_penalty=1.0, return_chunks=False):
        """Transcribe one recording or a list of recordings (paths, 1-D arrays at `sampling_rate`, or (array, sr)).
        Audio longer than 30 s is cut at quiet points into <=30 s chunks; chunk texts are joined with spaces.
        Returns a string (or a list of strings for a list input).
        """
        if processor is None:
            from transformers import AutoProcessor
            processor = AutoProcessor.from_pretrained(self._source or "oruk/oruk-loop-153M")
        single = not isinstance(audio, list)
        clips = [load_audio(a, sampling_rate) for a in ([audio] if single else audio)]
        items, owners = [], []
        for w in clips:
            b = split_points(w)
            ids = []
            for a, z in zip(b[:-1], b[1:]):
                piece = w[a:z]
                if len(piece) < 1600:  # <0.1 s tail
                    continue
                enc = processor.apply_transcription_request(audio=piece, language=language, return_tensors="pt",
                                                            processor_kwargs={"audio_kwargs": {"sampling_rate": SAMPLE_RATE}})
                ids.append(len(items))
                items.append({"input_ids": enc["input_ids"][0], "input_features": enc["input_features"][0].float(),
                              "input_features_mask": enc["input_features_mask"][0], "dur": len(piece) / SAMPLE_RATE})
            owners.append(ids)
        texts = [""] * len(items)
        order = sorted(range(len(items)), key=lambda i: items[i]["dur"])
        dev, dt = self.device, next(p for p in self.model.audio_tower.conv2d1.parameters()).dtype
        for k in range(0, len(order), batch_size):
            idx = order[k:k + batch_size]
            batch = self._collate([items[i] for i in idx], dev, dt)
            max_new = int(min(448, math.ceil(10 * max(items[i]["dur"] for i in idx)) + 24))
            out = self.generate(**batch, do_sample=False, num_beams=num_beams, length_penalty=length_penalty,
                                max_new_tokens=max_new, use_cache=True)
            cont = out[:, batch["input_ids"].shape[1]:]
            for j, i in enumerate(idx):
                toks = cont[j].tolist()
                for e_pos, t in enumerate(toks):
                    if t in EOS_IDS:
                        toks = toks[:e_pos + 1]
                        break
                texts[i] = processor.decode(toks, return_format="transcription_only").strip()
        results = [[texts[i] for i in own] for own in owners]
        joined = [" ".join(t for t in r if t) for r in results]
        if return_chunks:
            return (results[0] if single else results), (joined[0] if single else joined)
        return joined[0] if single else joined

    @staticmethod
    def _collate(items, device, dtype):
        B = len(items)
        P = max(len(it["input_ids"]) for it in items)
        Fm = max(it["input_features"].shape[-1] for it in items)
        Fm = (Fm + 99) // 100 * 100
        ids = torch.full((B, P), PAD_ID, dtype=torch.long)
        am = torch.zeros((B, P), dtype=torch.long)
        feats = torch.zeros((B, items[0]["input_features"].shape[0], Fm), dtype=torch.float32)
        fm = torch.zeros((B, Fm), dtype=torch.long)
        for b, it in enumerate(items):
            n = len(it["input_ids"])
            ids[b, P - n:] = it["input_ids"]
            am[b, P - n:] = 1
            f = it["input_features"].shape[-1]
            feats[b, :, :f] = it["input_features"]
            fm[b, :f] = it["input_features_mask"]
        return {"input_ids": ids.to(device), "attention_mask": am.to(device),
                "input_features": feats.to(device, dtype=dtype), "input_features_mask": fm.to(device)}

	"""oruk-loop-153M: a double-looped Qwen3-ASR student.

	The Qwen3-ASR architecture is kept, but each of its two transformer stacks is replaced by ONE shared block that is
	applied repeatedly:

	* audio encoder: one block applied 24 times, made depth-aware by a rank-64 adapter on every linear layer plus
	per-depth biases and LayerNorm gains;
	* text decoder: one block applied 28 times, with rank-16 adapters and per-depth RMSNorm gains;
	* token table: the kept English vocabulary (36,974 rows); 368 frequent rows are stored exactly and the rest are
	factorised at rank 729. The same table is used for input embeddings and the (tied) output layer.

	Token ids stay in the original Qwen3 id space, so the stock Qwen3-ASR processor (tokenizer, chat template,
	feature extractor) works unchanged.

	Usage:
	from transformers import AutoModel, AutoProcessor
	model = AutoModel.from_pretrained("oruk/oruk-loop-153M", trust_remote_code=True)
	processor = AutoProcessor.from_pretrained("oruk/oruk-loop-153M")
	print(model.transcribe("speech.wav", processor))
	"""
	from __future__ import annotations

	import copy
	import math
	import os
	from pathlib import Path

	import numpy as np
	import torch
	import torch.nn.functional as F
	from torch import nn
	from transformers import AutoConfig, GenerationConfig, Qwen3ASRForConditionalGeneration

	SAMPLE_RATE = 16000
	PAD_ID = 151645
	EOS_IDS = (151643, 151645)


	# ---------------------------------------------------------------------------------------------- looped modules
	class DepthCursor:
	"""Which virtual depth the shared block is computing (set by each call site before the block runs)."""

	def __init__(self, n: int):
	self.n = n
	self.depth = 0


	class LoopLinear(nn.Module):
	"""y = W x + b[d] + B[d] (A[d] x): one shared weight, per-depth bias and per-depth rank-r adapter."""

	def __init__(self, cursor: DepthCursor, n: int, in_features: int, out_features: int, rank: int, bias: bool):
	super().__init__()
	object.__setattr__(self, "cursor", cursor)
	self.in_features, self.out_features, self.rank = in_features, out_features, rank
	self.weight = nn.Parameter(torch.zeros(out_features, in_features))
	self.bias = nn.Parameter(torch.zeros(n, out_features)) if bias else None
	if rank > 0:
	self.lora_A = nn.Parameter(torch.zeros(n, rank, in_features))
	self.lora_B = nn.Parameter(torch.zeros(n, out_features, rank))
	object.__setattr__(self, "_merged", None)
	object.__setattr__(self, "_merged_bias", None)

	@torch.no_grad()
	def merge(self, dtype=torch.bfloat16):
	"""Materialise W + B[d] A[d] for every depth (faster decoding; ~3.5 GB in bf16 for the whole model)."""
	n = self.lora_A.shape[0] if self.rank > 0 else self.cursor.n
	W = torch.empty((n, self.out_features, self.in_features), dtype=dtype, device=self.weight.device)
	for i in range(n):
	w = self.weight.float()
	if self.rank > 0:
	w = w + self.lora_B[i].float() @ self.lora_A[i].float()
	W[i] = w.to(dtype)
	object.__setattr__(self, "_merged", W)
	object.__setattr__(self, "_merged_bias", None if self.bias is None else self.bias.detach().to(dtype))

	def unmerge(self):
	object.__setattr__(self, "_merged", None)
	object.__setattr__(self, "_merged_bias", None)

	def forward(self, x):
	d = self.cursor.depth
	if self._merged is not None:
	b = None if self._merged_bias is None else self._merged_bias[d]
	return F.linear(x.to(self._merged.dtype), self._merged[d], b)
	y = F.linear(x, self.weight, None if self.bias is None else self.bias[d])
	if self.rank > 0:
	y = y + F.linear(F.linear(x, self.lora_A[d]), self.lora_B[d])
	return y


	class LoopLayerNorm(nn.Module):
	def __init__(self, cursor, n, shape, eps):
	super().__init__()
	object.__setattr__(self, "cursor", cursor)
	self.shape, self.eps = tuple(shape), eps
	self.weight = nn.Parameter(torch.ones(n, *self.shape))
	self.bias = nn.Parameter(torch.zeros(n, *self.shape))

	def forward(self, x):
	d = self.cursor.depth
	w, b = self.weight[d], self.bias[d]
	if not torch.is_autocast_enabled(x.device.type): # under autocast layer_norm runs in fp32 with fp32 gains
	w, b = w.to(x.dtype), b.to(x.dtype)
	return F.layer_norm(x, self.shape, w, b, self.eps)


	class LoopRMSNorm(nn.Module):
	def __init__(self, cursor, n, dim, eps):
	super().__init__()
	object.__setattr__(self, "cursor", cursor)
	self.weight = nn.Parameter(torch.ones(n, dim))
	self.variance_epsilon = eps

	def forward(self, x):
	d = self.cursor.depth
	dtype = x.dtype
	h = x.to(torch.float32)
	h = h * torch.rsqrt(h.pow(2).mean(-1, keepdim=True) + self.variance_epsilon)
	return (self.weight[d] * h.to(dtype)).to(dtype)


	def build_loop_block(template: nn.Module, cursor: DepthCursor, n: int, rank: int) -> nn.Module:
	"""Turn one transformer layer into the shared looped block (same module paths, depth-aware leaves)."""
	block = copy.deepcopy(template)
	for name, mod in list(block.named_modules()):
	if not name:
	continue
	if isinstance(mod, nn.Linear):
	new = LoopLinear(cursor, n, mod.in_features, mod.out_features, rank, mod.bias is not None)
	elif isinstance(mod, nn.LayerNorm):
	new = LoopLayerNorm(cursor, n, mod.normalized_shape, mod.eps)
	elif type(mod).__name__.endswith("RMSNorm"):
	new = LoopRMSNorm(cursor, n, mod.weight.shape[0], mod.variance_epsilon)
	else:
	continue
	parent_name, _, child = name.rpartition(".")
	setattr(block.get_submodule(parent_name) if parent_name else block, child, new)
	return block


	class LoopCall(nn.Module):
	"""Parameter-free call site for virtual depth `index` of a shared block (one per original layer)."""

	def __init__(self, block, cursor, index, kv_slot=None):
	super().__init__()
	object.__setattr__(self, "_block", block)
	object.__setattr__(self, "_cursor", cursor)
	self.index, self.kv_slot = index, kv_slot

	def forward(self, hidden_states, args, *kwargs):
	self._cursor.depth = self.index
	if self.kv_slot is not None: # each virtual decoder depth keeps its own KV-cache slot
	self._block.self_attn.layer_idx = self.kv_slot
	return self._block(hidden_states, args, *kwargs)


	# ---------------------------------------------------------------------------------------------- token table
	class TokenTable(nn.Module):
	"""Tied table over the kept vocabulary: exact frequent rows + rank-r factorised rare rows (U @ P)."""

	def __init__(self, n_full, dim, n_kept, n_exact, rank):
	super().__init__()
	self.n_full, self.embedding_dim, self.rank = n_full, dim, rank
	self.register_buffer("keep", torch.zeros(n_kept, dtype=torch.long))
	self.register_buffer("lut", torch.zeros(n_full, dtype=torch.long))
	self.register_buffer("exact_idx", torch.zeros(n_exact, dtype=torch.long))
	self.register_buffer("rare_idx", torch.zeros(n_kept - n_exact, dtype=torch.long))
	self.exact_rows = nn.Parameter(torch.zeros(n_exact, dim, dtype=torch.bfloat16), requires_grad=False)
	self.U = nn.Parameter(torch.zeros(n_kept - n_exact, rank))
	self.P = nn.Parameter(torch.zeros(rank, dim))
	object.__setattr__(self, "_merged", None)

	def weight(self, dtype=None):
	dtype = dtype or self.U.dtype
	if self._merged is not None:
	return self._merged if self._merged.dtype == dtype else self._merged.to(dtype)
	W = torch.empty(self.keep.shape[0], self.embedding_dim, dtype=dtype, device=self.U.device)
	W = W.index_copy(0, self.exact_idx, self.exact_rows.to(dtype))
	return W.index_copy(0, self.rare_idx, (self.U.float() @ self.P.float()).to(dtype))

	@torch.no_grad()
	def merge(self, dtype=torch.bfloat16):
	object.__setattr__(self, "_merged", None)
	object.__setattr__(self, "_merged", self.weight(dtype).detach())

	def forward(self, ids):
	# read the merged table in its own dtype (bf16 on GPU), as evaluated; never re-cast it per call
	return F.embedding(self.lut[ids], self._merged if self._merged is not None else self.weight())


	class TableEmbedding(nn.Module):
	def __init__(self, table):
	super().__init__()
	object.__setattr__(self, "table", table)
	self.embedding_dim = table.embedding_dim

	@property
	def num_embeddings(self):
	return self.table.n_full

	@property
	def weight(self):
	return self.table.weight()

	def forward(self, ids):
	return self.table(ids)


	class TableHead(nn.Module):
	"""Output layer tied to the token table; full-width logits with -inf at ids outside the kept vocabulary."""

	def __init__(self, table):
	super().__init__()
	object.__setattr__(self, "table", table)
	self.in_features, self.out_features = table.embedding_dim, table.n_full

	@property
	def weight(self):
	return self.table.weight()

	def forward(self, h):
	kept = F.linear(h, self.table.weight(h.dtype))
	out = torch.full(h.shape[:-1] + (self.table.n_full,), float("-inf"), dtype=kept.dtype, device=kept.device)
	return out.index_copy_(-1, self.table.keep, kept)


	# ---------------------------------------------------------------------------------------------- audio helpers
	def load_audio(source, sampling_rate=None):
	"""path / (array, sr) / 1-D array -> float32 mono at 16 kHz."""
	if isinstance(source, (str, os.PathLike)):
	import soundfile as sf
	wav, sr = sf.read(str(source), dtype="float32", always_2d=False)
	elif isinstance(source, tuple):
	wav, sr = source
	else:
	wav, sr = source, sampling_rate or SAMPLE_RATE
	wav = np.asarray(wav, dtype=np.float32)
	if wav.ndim > 1:
	wav = wav.mean(axis=1 if wav.shape[0] > wav.shape[1] else 0)
	if sr != SAMPLE_RATE:
	try:
	import soxr
	wav = soxr.resample(wav, sr, SAMPLE_RATE, quality="HQ").astype(np.float32)
	except ImportError:
	from math import gcd
	from scipy.signal import resample_poly
	g = gcd(int(sr), SAMPLE_RATE)
	wav = resample_poly(wav, SAMPLE_RATE // g, int(sr) // g).astype(np.float32)
	# the evaluation pipeline stores 16-bit PCM; quantise the same way so results reproduce exactly
	return np.clip(np.rint(wav * 32767.0), -32768, 32767).astype(np.float32) / 32768.0


	def split_points(w, sr=SAMPLE_RATE, max_s=30.0, search_s=5.0, frame=0.1):
	"""Chunk boundaries for audio longer than max_s: each cut sits at the quietest 100 ms frame within the last
	search_s seconds before the limit, so words are rarely split."""
	n, cuts, start = len(w), [], 0
	L, S, Fr = int(max_s * sr), int(search_s * sr), int(frame * sr)
	while n - start > L:
	lo, hi = start + L - S, start + L
	seg = w[lo:hi].astype(np.float32)
	k = len(seg) // Fr
	e = (seg[:k * Fr].reshape(k, Fr) ** 2).mean(1)
	cut = lo + int(np.argmin(e)) * Fr + Fr // 2 if k else hi
	cuts.append(cut)
	start = cut
	return [0] + cuts + [n]


	# ---------------------------------------------------------------------------------------------- model
	class OrukLoopForConditionalGeneration(Qwen3ASRForConditionalGeneration):
	"""Qwen3-ASR with both stacks looped. `from_pretrained` loads the 153M stored parameters; see `transcribe`."""

	_tied_weights_keys = {}

	def __init__(self, config):
	loop = dict(getattr(config, "oruk_loop", {}))
	tc, ac = config.text_config, config.audio_config
	n_dec, n_enc, vocab = tc.num_hidden_layers, ac.encoder_layers, tc.vocab_size
	types = list(tc.layer_types)
	# build the stock skeleton with ONE template layer per stack and a tiny vocabulary (cheap), then restore
	tc.num_hidden_layers, tc.layer_types, tc.vocab_size, ac.encoder_layers = 1, types[:1], 8, 1
	try:
	super().__init__(config)
	finally:
	tc.num_hidden_layers, tc.layer_types, tc.vocab_size, ac.encoder_layers = n_dec, types, vocab, n_enc
	at, lm = self.model.audio_tower, self.model.language_model
	for c in {id(at.config): at.config}.values(): # sub-modules may hold copies of the sub-configs
	c.encoder_layers = n_enc
	lm.config.num_hidden_layers, lm.config.layer_types, lm.config.vocab_size = n_dec, types, vocab

	self.enc_cursor, self.dec_cursor = DepthCursor(n_enc), DepthCursor(n_dec)
	at.loop = build_loop_block(at.layers[0], self.enc_cursor, n_enc, int(loop.get("encoder_rank", 64)))
	at.layers = nn.ModuleList(LoopCall(at.loop, self.enc_cursor, i) for i in range(n_enc))
	lm.loop = build_loop_block(lm.layers[0], self.dec_cursor, n_dec, int(loop.get("decoder_rank", 16)))
	lm.layers = nn.ModuleList(LoopCall(lm.loop, self.dec_cursor, i, kv_slot=i) for i in range(n_dec))
	lm.token_table = TokenTable(vocab, tc.hidden_size, int(loop.get("table_rows", 36974)),
	int(loop.get("table_exact", 368)), int(loop.get("table_rank", 729)))
	lm.embed_tokens = TableEmbedding(lm.token_table)
	self.lm_head = TableHead(lm.token_table)
	self._source = None

	# the table is tied by construction; never let transformers re-tie or resize it
	def tie_weights(self, args, *kwargs):
	return None

	def tie_embeddings_and_encoder_decoder(self, args, *kwargs):
	return None

	def num_stored_parameters(self):
	return sum(p.numel() for p in self.parameters())

	def merge_loops(self, dtype=torch.bfloat16):
	for m in self.modules():
	if isinstance(m, LoopLinear):
	m.merge(dtype)

	def unmerge_loops(self):
	for m in self.modules():
	if isinstance(m, LoopLinear):
	m.unmerge()

	@classmethod
	def from_pretrained(cls, pretrained_model_name_or_path, *model_args, config=None, dtype=None, torch_dtype=None,
	device=None, device_map=None, merge=None, revision=None, cache_dir=None, token=None,
	local_files_only=False, **kwargs):
	"""Load the release weights.

	device: "cuda", "cpu", ... (default: cuda if available). dtype: compute dtype (default bf16 on GPU, fp32 on CPU).
	merge: materialise per-depth weights for fast decoding (default: True on GPU, False on CPU).
	"""
	ignored = [k for k in kwargs if not k.startswith("_") and k not in ("trust_remote_code", "low_cpu_mem_usage",
	"attn_implementation", "use_safetensors")]
	if ignored:
	import warnings
	warnings.warn(f"oruk-loop from_pretrained ignores these arguments: {ignored}")
	src = str(pretrained_model_name_or_path)
	if os.path.isdir(src):
	path = Path(src)
	else:
	from huggingface_hub import snapshot_download
	path = Path(snapshot_download(src, revision=revision, cache_dir=cache_dir, token=token,
	local_files_only=local_files_only,
	allow_patterns=[".json", ".safetensors", ".py", ".jinja", "*.txt"]))
	if config is None:
	config = AutoConfig.from_pretrained(str(path))
	if device is None:
	device = device_map if isinstance(device_map, str) and device_map != "auto" else None
	device = torch.device(device or ("cuda" if torch.cuda.is_available() else "cpu"))
	dtype = dtype or torch_dtype
	if isinstance(dtype, str) and dtype != "auto":
	dtype = getattr(torch, dtype)
	if dtype is None or dtype == "auto":
	dtype = torch.bfloat16 if device.type == "cuda" else torch.float32
	if merge is None:
	merge = device.type == "cuda"

	model = cls(config)
	from safetensors.torch import load_file
	state = load_file(str(path / "model.safetensors"))
	missing, unexpected = model.load_state_dict(state, strict=False)
	if missing or unexpected:
	raise RuntimeError(f"oruk-loop weights do not match the model: missing={missing[:8]} "
	f"unexpected={unexpected[:8]}")
	del state
	try:
	model.generation_config = GenerationConfig.from_pretrained(str(path))
	except Exception:
	pass
	model.to(device)
	model.model.language_model.token_table.merge(dtype)
	if merge:
	model.merge_loops(dtype)
	if dtype == torch.bfloat16 and device.type == "cuda":
	# evaluated precision: modules kept from Qwen3-ASR in bf16, looped weights merged to bf16, per-depth norms
	# and biases in fp32, everything run under bf16 autocast
	at, lm = model.model.audio_tower, model.model.language_model
	for mod in (at.conv2d1, at.conv2d2, at.conv2d3, at.conv_out, at.ln_post, model.model.multi_modal_projector,
	lm.norm):
	mod.to(torch.bfloat16)
	model._autocast = torch.bfloat16
	elif dtype != torch.float32:
	model.to(dtype)
	model.eval()
	model._source = str(path)
	return model

	_autocast = None

	def generate(self, args, *kwargs):
	if self._autocast is not None and not torch.is_autocast_enabled(self.device.type):
	with torch.autocast(self.device.type, dtype=self._autocast):
	return super().generate(args, *kwargs)
	return super().generate(args, *kwargs)

	# ------------------------------------------------------------------------------------------ transcription
	@torch.no_grad()
	def transcribe(self, audio, processor=None, *, sampling_rate=None, batch_size=16, num_beams=1, language="English",
	length_penalty=1.0, return_chunks=False):
	"""Transcribe one recording or a list of recordings (paths, 1-D arrays at `sampling_rate`, or (array, sr)).

	Audio longer than 30 s is cut at quiet points into <=30 s chunks; chunk texts are joined with spaces.
	Returns a string (or a list of strings for a list input).
	"""
	if processor is None:
	from transformers import AutoProcessor
	processor = AutoProcessor.from_pretrained(self._source or "oruk/oruk-loop-153M")
	single = not isinstance(audio, list)
	clips = [load_audio(a, sampling_rate) for a in ([audio] if single else audio)]
	items, owners = [], []
	for w in clips:
	b = split_points(w)
	ids = []
	for a, z in zip(b[:-1], b[1:]):
	piece = w[a:z]
	if len(piece) < 1600: # <0.1 s tail
	continue
	enc = processor.apply_transcription_request(audio=piece, language=language, return_tensors="pt",
	processor_kwargs={"audio_kwargs": {"sampling_rate": SAMPLE_RATE}})
	ids.append(len(items))
	items.append({"input_ids": enc["input_ids"][0], "input_features": enc["input_features"][0].float(),
	"input_features_mask": enc["input_features_mask"][0], "dur": len(piece) / SAMPLE_RATE})
	owners.append(ids)
	texts = [""] * len(items)
	order = sorted(range(len(items)), key=lambda i: items[i]["dur"])
	dev, dt = self.device, next(p for p in self.model.audio_tower.conv2d1.parameters()).dtype
	for k in range(0, len(order), batch_size):
	idx = order[k:k + batch_size]
	batch = self._collate([items[i] for i in idx], dev, dt)
	max_new = int(min(448, math.ceil(10 * max(items[i]["dur"] for i in idx)) + 24))
	out = self.generate(**batch, do_sample=False, num_beams=num_beams, length_penalty=length_penalty,
	max_new_tokens=max_new, use_cache=True)
	cont = out[:, batch["input_ids"].shape[1]:]
	for j, i in enumerate(idx):
	toks = cont[j].tolist()
	for e_pos, t in enumerate(toks):
	if t in EOS_IDS:
	toks = toks[:e_pos + 1]
	break
	texts[i] = processor.decode(toks, return_format="transcription_only").strip()
	results = [[texts[i] for i in own] for own in owners]
	joined = [" ".join(t for t in r if t) for r in results]
	if return_chunks:
	return (results[0] if single else results), (joined[0] if single else joined)
	return joined[0] if single else joined

	@staticmethod
	def _collate(items, device, dtype):
	B = len(items)
	P = max(len(it["input_ids"]) for it in items)
	Fm = max(it["input_features"].shape[-1] for it in items)
	Fm = (Fm + 99) // 100 * 100
	ids = torch.full((B, P), PAD_ID, dtype=torch.long)
	am = torch.zeros((B, P), dtype=torch.long)
	feats = torch.zeros((B, items[0]["input_features"].shape[0], Fm), dtype=torch.float32)
	fm = torch.zeros((B, Fm), dtype=torch.long)
	for b, it in enumerate(items):
	n = len(it["input_ids"])
	ids[b, P - n:] = it["input_ids"]
	am[b, P - n:] = 1
	f = it["input_features"].shape[-1]
	feats[b, :, :f] = it["input_features"]
	fm[b, :f] = it["input_features_mask"]
	return {"input_ids": ids.to(device), "attention_mask": am.to(device),
	"input_features": feats.to(device, dtype=dtype), "input_features_mask": fm.to(device)}