| """Mateli: streaming Finnish speech recognition (MinGRU encoder + RNN-T), PyTorch reference. |
| |
| python mateli.py speech.wav # text |
| python mateli.py speech.wav --words # one word per line with start/end times |
| |
| from mateli import Mateli, load_audio |
| asr = Mateli.from_pretrained(".") # directory with config.json and model.safetensors |
| result = asr.transcribe(load_audio("speech.wav")) |
| print(result.text) |
| |
| The encoder is chunk-causal: audio is processed in 320 ms chunks (32 log-mel frames -> 8 encoder |
| frames of 40 ms), each chunk sees only itself and the past, and the state carried between chunks |
| has a fixed size. transcribe() runs the same chunk loop a live stream would, so its output equals |
| streaming recognition of the same audio. |
| |
| Model: log-mel (80 bins, 25 ms window, 10 ms hop) -> two stride-2 Conv2d -> 16 macaron blocks |
| (feed-forward / MinGRU / causal depthwise conv / feed-forward), with a windowed self-attention |
| sublayer after blocks 4, 7, 10 and 13 (keys = the current chunk + 64 previous frames) -> additive |
| joiner with a one-layer MinGRU prediction network over 44 characters. Greedy decoding. |
| """ |
| from __future__ import annotations |
|
|
| import argparse |
| import json |
| import math |
| from dataclasses import dataclass |
| from pathlib import Path |
|
|
| import torch |
| import torch.nn as nn |
| import torch.nn.functional as F |
|
|
|
|
| |
| class MinGRU(nn.Module): |
| """minGRU (Feng et al. 2024, "Were RNNs All We Needed?") with a sigmoid highway. |
| h_t = (1 - z_t) * h_{t-1} + z_t * g(x_t), computed for a whole chunk with a log-space scan.""" |
|
|
| def __init__(self, hidden: int): |
| super().__init__() |
| self.layers = nn.ModuleList([nn.Linear(hidden, 3 * hidden, bias=False)]) |
|
|
| @staticmethod |
| def _log_g(x): |
| return torch.where(x >= 0, (F.relu(x) + 0.5).log(), -F.softplus(-x)) |
|
|
| def forward(self, x, state): |
| """x [B, T, H], state [B, H] (>= 0) -> (output [B, T, H], new state [B, H]).""" |
| hidden, gate, proj = self.layers[0](x).chunk(3, dim=-1) |
| log_coeffs = -F.softplus(gate) |
| log_values = -F.softplus(-gate) + self._log_g(hidden) |
| |
| log_coeffs = torch.cat([torch.zeros_like(log_coeffs[:, :1]), log_coeffs], dim=1) |
| log_values = torch.cat([state.clamp_min(0).log().unsqueeze(1), log_values], dim=1) |
| a_star = log_coeffs.cumsum(dim=1) |
| h = (a_star + (log_values - a_star).logcumsumexp(dim=1)).exp()[:, 1:] |
| g = proj.sigmoid() |
| return g * h + (1.0 - g) * x, h[:, -1] |
|
|
| def step(self, x, state): |
| """One time step without the scan: x [B, H], state [B, H] -> (output, new state).""" |
| hidden, gate, proj = self.layers[0](x).chunk(3, dim=-1) |
| g_hidden = torch.where(hidden >= 0, hidden + 0.5, hidden.sigmoid()) |
| h = torch.lerp(state, g_hidden, gate.sigmoid()) |
| g = proj.sigmoid() |
| return g * h + (1.0 - g) * x, h |
|
|
|
|
| |
| class FeedForward(nn.Module): |
| def __init__(self, d: int, expansion: int): |
| super().__init__() |
| self.net = nn.Sequential(nn.LayerNorm(d), nn.Linear(d, expansion * d), nn.SiLU(), |
| nn.Linear(expansion * d, d)) |
|
|
| def forward(self, x): |
| return self.net(x) |
|
|
|
|
| class MinGRUMixer(nn.Module): |
| def __init__(self, d: int): |
| super().__init__() |
| self.norm = nn.LayerNorm(d) |
| self.gru = MinGRU(d) |
|
|
| def forward(self, x, state): |
| return self.gru(self.norm(x), state) |
|
|
|
|
| class ConvModule(nn.Module): |
| """Conformer convolution module with a causal depthwise convolution.""" |
|
|
| def __init__(self, d: int, kernel: int): |
| super().__init__() |
| self.norm = nn.LayerNorm(d) |
| self.pw1 = nn.Conv1d(d, 2 * d, 1) |
| self.dw = nn.Conv1d(d, d, kernel, groups=d) |
| self.act_norm = nn.LayerNorm(d) |
| self.pw2 = nn.Conv1d(d, d, 1) |
|
|
| def forward(self, x, cache): |
| """x [B, T, D], cache [B, D, kernel-1] (the previous GLU outputs).""" |
| y = F.glu(self.pw1(self.norm(x).transpose(1, 2)), dim=1) |
| y = torch.cat([cache, y], dim=2) |
| new_cache = y[:, :, -cache.shape[2]:] |
| y = F.silu(self.act_norm(self.dw(y).transpose(1, 2))) |
| return self.pw2(y.transpose(1, 2)).transpose(1, 2), new_cache |
|
|
|
|
| class Block(nn.Module): |
| def __init__(self, d: int, kernel: int, expansion: int): |
| super().__init__() |
| self.ff1 = FeedForward(d, expansion) |
| self.mixer = MinGRUMixer(d) |
| self.conv = ConvModule(d, kernel) |
| self.ff2 = FeedForward(d, expansion) |
| self.norm = nn.LayerNorm(d) |
|
|
| def forward(self, x, state): |
| gru_h, conv_cache = state |
| x = x + 0.5 * self.ff1(x) |
| m, gru_h = self.mixer(x, gru_h) |
| x = x + m |
| c, conv_cache = self.conv(x, conv_cache) |
| x = x + c |
| x = x + 0.5 * self.ff2(x) |
| return self.norm(x), (gru_h, conv_cache) |
|
|
|
|
| class _Attn(nn.Module): |
| def __init__(self, d: int, heads: int): |
| super().__init__() |
| self.norm = nn.LayerNorm(d) |
| self.mha = nn.MultiheadAttention(d, heads, batch_first=True) |
|
|
|
|
| class WindowedAttention(nn.Module): |
| """Residual self-attention: the chunk's frames attend to the chunk and up to `left` cached |
| earlier inputs. No positional encoding (the causal convolutions provide position).""" |
|
|
| def __init__(self, d: int, heads: int, left: int): |
| super().__init__() |
| self.attn = _Attn(d, heads) |
| self.left = left |
|
|
| def forward(self, x, cache): |
| ctx = torch.cat([cache, x], dim=1) |
| q, kv = self.attn.norm(x), self.attn.norm(ctx) |
| y, _ = self.attn.mha(q, kv, kv, need_weights=False) |
| return x + y, ctx[:, -self.left:] |
|
|
|
|
| class Encoder(nn.Module): |
| def __init__(self, cfg: dict): |
| super().__init__() |
| d = cfg["d_model"] |
| self.blocks = nn.ModuleList([Block(d, cfg["conv_kernel"], cfg["ffn_expansion"]) |
| for _ in range(cfg["n_blocks"])]) |
| self.attn = nn.ModuleDict({str(p): WindowedAttention(d, cfg["attn_heads"], cfg["attn_left_context"]) |
| for p in cfg["attn_after_blocks"]}) |
|
|
| def initial_state(self, device): |
| d, k = self.blocks[0].norm.normalized_shape[0], self.blocks[0].conv.dw.kernel_size[0] |
| return {"blocks": [(torch.zeros(1, d, device=device), torch.zeros(1, d, k - 1, device=device)) |
| for _ in self.blocks], |
| "attn": {p: torch.zeros(1, 0, d, device=device) for p in self.attn}} |
|
|
| def forward(self, x, state): |
| new_blocks, new_attn = [], {} |
| for i, blk in enumerate(self.blocks): |
| x, st = blk(x, state["blocks"][i]) |
| new_blocks.append(st) |
| p = str(i + 1) |
| if p in self.attn: |
| x, new_attn[p] = self.attn[p](x, state["attn"][p]) |
| return x, {"blocks": new_blocks, "attn": new_attn} |
|
|
|
|
| |
| class LogMel(nn.Module): |
| def __init__(self, n_mels=80, n_fft=512, win=400, hop=160): |
| super().__init__() |
| self.n_fft, self.win, self.hop = n_fft, win, hop |
| self.register_buffer("fb", torch.zeros(n_mels, n_fft // 2 + 1)) |
| self.register_buffer("window", torch.hann_window(win)) |
|
|
| def forward(self, wav): |
| """[B, samples] -> [B, frames, n_mels], frames = 1 + samples // hop.""" |
| spec = torch.stft(wav, self.n_fft, self.hop, win_length=self.win, window=self.window, |
| center=True, return_complex=True) |
| mel = torch.einsum("mf,bft->bmt", self.fb, spec.real ** 2 + spec.imag ** 2) |
| return torch.log(mel + 1e-6).transpose(1, 2) |
|
|
|
|
| class Subsampler(nn.Module): |
| """Two Conv2d (3x3, stride 2) over (time, mel) -> 40 ms frames.""" |
|
|
| def __init__(self, n_mels: int, d: int): |
| super().__init__() |
| self.conv = nn.Sequential(nn.Conv2d(1, d, 3, stride=2, padding=1), nn.ReLU(), |
| nn.Conv2d(d, d, 3, stride=2, padding=1), nn.ReLU()) |
| f_out = ((n_mels - 1) // 2 + 1 - 1) // 2 + 1 |
| self.out = nn.Linear(d * f_out, d) |
|
|
| def forward(self, mel): |
| x = self.conv(mel.unsqueeze(1)) |
| b, c, t, f = x.shape |
| return self.out(x.transpose(1, 2).reshape(b, t, c * f)) |
|
|
|
|
| class FrontEnd(nn.Module): |
| def __init__(self, n_mels: int, d: int): |
| super().__init__() |
| self.logmel = LogMel(n_mels) |
| self.subsample = Subsampler(n_mels, d) |
|
|
|
|
| |
| class PredictionNetwork(nn.Module): |
| def __init__(self, vocab_size: int, hidden: int): |
| super().__init__() |
| self.embed = nn.Embedding(vocab_size, hidden) |
| self.rnn = MinGRU(hidden) |
|
|
| def step(self, token: int, state): |
| return self.rnn.step(self.embed(torch.tensor([token], device=state.device)), state) |
|
|
|
|
| class Joiner(nn.Module): |
| def __init__(self, enc_dim: int, pred_dim: int, joint_dim: int, vocab_size: int): |
| super().__init__() |
| self.enc_proj = nn.Linear(enc_dim, joint_dim) |
| self.pred_proj = nn.Linear(pred_dim, joint_dim) |
| self.out = nn.Linear(joint_dim, vocab_size) |
|
|
|
|
| @dataclass(frozen=True) |
| class Word: |
| text: str |
| start: float |
| end: float |
|
|
|
|
| @dataclass |
| class Transcript: |
| words: list[Word] |
|
|
| @property |
| def text(self) -> str: |
| return " ".join(w.text for w in self.words) |
|
|
|
|
| class Mateli(nn.Module): |
| def __init__(self, cfg: dict): |
| super().__init__() |
| self.cfg = cfg |
| d = cfg["d_model"] |
| self.frontend = FrontEnd(cfg["n_mels"], d) |
| self.encoder = Encoder(cfg) |
| self.prednet = PredictionNetwork(len(cfg["vocab"]), cfg["pred_hidden"]) |
| self.joiner = Joiner(d, cfg["pred_hidden"], cfg["joint_dim"], len(cfg["vocab"])) |
|
|
| @classmethod |
| def from_pretrained(cls, model_dir=".", device: str | None = None) -> "Mateli": |
| from safetensors.torch import load_file |
| d = Path(model_dir) |
| model = cls(json.loads((d / "config.json").read_text(encoding="utf-8"))) |
| model.load_state_dict(load_file(str(d / "model.safetensors"))) |
| device = device or ("cuda" if torch.cuda.is_available() else "cpu") |
| return model.to(device).eval() |
|
|
| @property |
| def device(self): |
| return self.joiner.out.weight.device |
|
|
| @torch.inference_mode() |
| def transcribe(self, samples) -> Transcript: |
| """float32 mono 16 kHz samples in [-1, 1] -> words with times.""" |
| cfg = self.cfg |
| wav = torch.as_tensor(samples, dtype=torch.float32, device=self.device) |
| if wav.ndim != 1: |
| raise ValueError(f"expected mono samples of shape [n], got {tuple(wav.shape)}; mix channels down first") |
| wav = wav[None] |
| if wav.shape[1] <= self.frontend.logmel.n_fft // 2: |
| return Transcript([]) |
| mel = self.frontend.logmel(wav) |
| n_chunk, ctx = cfg["chunk_mel_frames"], 4 |
| state = self.encoder.initial_state(self.device) |
| blank, space = cfg["blank_id"], cfg["space_id"] |
| pred, h = self.prednet.step(blank, torch.zeros(1, cfg["pred_hidden"], device=self.device)) |
| pred_proj = self.joiner.pred_proj(pred) |
| words, chars, t0, t1, frame = [], "", 0, 0, 0 |
| fs = cfg["frame_seconds"] |
| for start in range(0, mel.shape[1], n_chunk): |
| |
| |
| lo = max(0, start - ctx) |
| x = self.frontend.subsample(mel[:, lo:start + n_chunk]) |
| x = x[:, 1:] if lo < start else x |
| enc, state = self.encoder(x, state) |
| enc_proj = self.joiner.enc_proj(enc)[0] |
| for e in enc_proj: |
| for _ in range(cfg["max_symbols_per_frame"]): |
| k = int(self.joiner.out(torch.tanh(e + pred_proj[0])).argmax()) |
| if k == blank: |
| break |
| if k == space: |
| if chars: |
| words.append(Word(chars, round((t0 + 1) * fs, 2), round((t1 + 1) * fs, 2))) |
| chars = "" |
| else: |
| if not chars: |
| t0 = frame |
| chars, t1 = chars + cfg["vocab"][k], frame |
| pred, h = self.prednet.step(k, h) |
| pred_proj = self.joiner.pred_proj(pred) |
| frame += 1 |
| if chars: |
| words.append(Word(chars, round((t0 + 1) * fs, 2), round((t1 + 1) * fs, 2))) |
| return Transcript(words) |
|
|
|
|
| def load_audio(path, sample_rate: int = 16000): |
| import numpy as np |
| import soundfile as sf |
| x, sr = sf.read(str(path), dtype="float32", always_2d=True) |
| x = x.mean(axis=1) |
| if sr != sample_rate: |
| from scipy.signal import resample_poly |
| g = math.gcd(sr, sample_rate) |
| x = resample_poly(x, sample_rate // g, sr // g).astype(np.float32) |
| return x |
|
|
|
|
| def main(): |
| ap = argparse.ArgumentParser(description=__doc__.split("\n")[0]) |
| ap.add_argument("audio") |
| ap.add_argument("--model-dir", default=str(Path(__file__).resolve().parent)) |
| ap.add_argument("--device", default=None) |
| ap.add_argument("--words", action="store_true", help="print words with start and end times") |
| a = ap.parse_args() |
| asr = Mateli.from_pretrained(a.model_dir, a.device) |
| t = asr.transcribe(load_audio(a.audio, asr.cfg["sample_rate"])) |
| if a.words: |
| for w in t.words: |
| print(f"{w.start:7.2f} {w.end:7.2f} {w.text}") |
| else: |
| print(t.text) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
