diff --git a/.gitignore b/.gitignore index 1f0492f..dc92ee6 100644 --- a/.gitignore +++ b/.gitignore @@ -5,6 +5,14 @@ build .DS_Store *tmp* onnx +mlx_convert/.venv +mlx_convert/gigaam-v3-ctc-mlx/ +mlx_convert/gigaam-v3-ctc-mlx-fp32/ +mlx_convert/gigaam-v3-rnnt-mlx/ +mlx_convert/gigaam-v3-e2e-rnnt-mlx/ +mlx_convert/eval-models/ +mlx_convert/eval/data/ +mlx_convert/eval/results/*.jsonl train_utils/data train_utils/checkpoints .venv diff --git a/mlx_convert/README.md b/mlx_convert/README.md new file mode 100644 index 0000000..634039a --- /dev/null +++ b/mlx_convert/README.md @@ -0,0 +1,260 @@ +# GigaAM v3 MLX — Russian ASR on Apple Silicon + +GigaAM v3 (Conformer encoder, 16 layers, 768d) converted to [MLX](https://github.com/ml-explore/mlx) for fast inference on Apple Silicon. +Supports **CTC**, **RNNT** and the end-to-end **E2E CTC / E2E RNNT** models (punctuated, normalized text). + +On 200 Golos Crowd clips MLX transcripts are identical to upstream PyTorch for all four models; +speed, WER and parity numbers are in [Evaluation](#evaluation) and are reproducible with `eval/`. + +## Quick Start + +### 1. Install dependencies + +```bash +uv venv .venv +uv pip install mlx safetensors numpy +# For streaming from microphone: +uv pip install sounddevice +# For conversion (step 2) the PyTorch package is needed once: +uv pip install -e "..[torch]" +``` + +### 2. Convert the model (one-time) + +```bash +# CTC model (fastest) +python convert_gigaam_to_mlx.py --model v3_ctc --output ./gigaam-v3-ctc-mlx + +# RNNT model (lower WER, sequential decode) +python convert_gigaam_to_mlx.py --model v3_rnnt --output ./gigaam-v3-rnnt-mlx + +# End-to-end models (SentencePiece output with punctuation/normalization) +python convert_gigaam_to_mlx.py --model v3_e2e_ctc --output ./gigaam-v3-e2e-ctc-mlx +python convert_gigaam_to_mlx.py --model v3_e2e_rnnt --output ./gigaam-v3-e2e-rnnt-mlx +``` + +This creates a directory with: +- `model.safetensors` — weights (≈421–425 MB, fp16) +- `config.json` — model configuration + vocabulary (for E2E: all SentencePiece pieces) +- `tokenizer.model` — SentencePiece model (E2E only; the runtime decodes from `config.json` and does not need the `sentencepiece` package) + +The upstream encoder weights are stored in fp16, so `--dtype float32` doubles the file without changing them (only the small CTC head is fp32 upstream). + +### 3. Transcribe + +```bash +# Single file +python gigaam-cli -f audio.wav + +# Streaming from file +python gigaam-stream --file audio.wav + +# Live microphone streaming +python gigaam-stream +``` + +--- + +## Python API + +### Basic transcription + +```python +from gigaam_mlx import load_model, load_audio + +model = load_model("./gigaam-v3-ctc-mlx") +audio = load_audio("audio.wav") # any format, resampled to 16kHz via ffmpeg +text = model.transcribe(audio) +print(text) +# → ничьих не требуя похвал счастлив уж я надеждой сладкой +``` + +### Streaming (pre-recorded file) + +Process audio incrementally, yielding results every N seconds: + +```python +from gigaam_mlx import load_model, load_audio, StreamingConfig + +model = load_model("./gigaam-v3-ctc-mlx") +audio = load_audio("audio.wav") + +config = StreamingConfig(step_duration=1.0) # yield every 1s + +for result in model.stream_generate(audio, config): + print(f"[{result.audio_position:.1f}s] {result.cumulative_text}") + # [1.0s] ничьих не требуя + # [2.0s] ничьих не требуя похвал + # [3.0s] ничьих не требуя похвал счастлив уж я надеж + # ... +``` + +`StreamingResult` fields: + +| Field | Type | Description | +|-------|------|-------------| +| `text` | `str` | New text since last emission | +| `cumulative_text` | `str` | Full transcription so far | +| `is_final` | `bool` | `True` if last chunk | +| `audio_position` | `float` | Current position in seconds | +| `audio_duration` | `float` | Total audio duration | +| `progress` | `float` | 0.0–1.0 | +| `language` | `str` | Always `"ru"` | + +### Streaming (live microphone) + +For real-time transcription, call `stream_live()` with a growing audio buffer: + +```python +import numpy as np +import mlx.core as mx +from gigaam_mlx import load_model + +model = load_model("./gigaam-v3-ctc-mlx") + +# Accumulate audio from microphone (16kHz float32 mono) +buffer = np.zeros(0, dtype=np.float32) + +# Called every N ms with new audio +def on_audio_chunk(chunk: np.ndarray): + global buffer + buffer = np.concatenate([buffer, chunk]) + + result = model.stream_live(mx.array(buffer)) + print(f"\r{result.cumulative_text}", end="", flush=True) +``` + +### StreamingConfig options + +```python +from gigaam_mlx import StreamingConfig + +config = StreamingConfig( + step_duration=1.0, # process every 1s (default: 2s) + chunk_duration=2.0, # unused for stream_generate (kept for compat) + context_duration=3.0, # unused for stream_generate (kept for compat) +) +``` + +--- + +## mlx-audio Compatibility + +The `StreamingResult` dataclass follows the same contract as [mlx-audio](https://github.com/Blaizzy/mlx-audio) Parakeet/Whisper streaming, making it straightforward to integrate GigaAM as an mlx-audio STT model. + +--- + +## CLI Tools + +### `gigaam-cli` — Single-file transcription + +```bash +python gigaam-cli -f audio.wav # default model +python gigaam-cli -f audio.wav -m /path/to/model # custom model path +python gigaam-cli -f audio.wav --no-prints # only text to stdout +``` + +### `gigaam-stream` — Real-time streaming + +```bash +# Live microphone +python gigaam-stream +python gigaam-stream --step 1000 # update every 1s +python gigaam-stream --step 500 # update every 0.5s + +# File streaming (simulates real-time) +python gigaam-stream --file audio.wav +python gigaam-stream --file audio.wav --step 1000 --no-overwrite +``` + +Options: + +| Flag | Default | Description | +|------|---------|-------------| +| `--step N` | 2000 | Process every N ms | +| `--file PATH` | — | File mode instead of microphone | +| `--model PATH` | auto | Model directory | +| `--no-overwrite` | off | Print incrementally (don't clear line) | +| `--vad-threshold` | 0.003 | Energy threshold for speech detection | + +### `gigaam-transcribe` — Shell wrapper + +```bash +# Uses bundled Python venv automatically +gigaam-transcribe -f audio.wav --no-prints + +# Symlink for PATH access +ln -s /path/to/mlx_convert/gigaam-transcribe /usr/local/bin/gigaam-transcribe +``` + +--- + +## Evaluation + +Measured on Apple M4 Pro (24 GB), macOS 26.6.2, mlx 0.32.2, torch 2.14.0, on 200 Golos Crowd test clips +(823.85 s) — the sample of [garrrikkotua/gigaam-vs-whisper-russian](https://github.com/garrrikkotua/gigaam-vs-whisper-russian), +verified by SHA-256. Full tables and method: [`eval/results/RESULTS.md`](eval/results/RESULTS.md); +to reproduce: [`eval/README.md`](eval/README.md). + +**Parity with PyTorch** (upstream model on CPU, fp32): identical transcripts on 200/200 clips for all four +models; encoder output max |Δ| ≤ 9e-5 per clip (outputs are O(1)); CTC frame labels 100% identical. + +**Accuracy** (corpus WER; E2E output normalized — lowercase, ё→е, numbers → words, no punctuation): + +| | v3_ctc | v3_rnnt | v3_e2e_ctc | v3_e2e_rnnt | +|---|---:|---:|---:|---:| +| WER | 1.70% | 1.50% | 6.49% | 5.69% | +| WER, 168 clips without digits/Latin in E2E output | 1.35% | 1.23% | 2.70% | 2.09% | + +E2E models write "HD", "4K", "15-й", phone numbers as digit groups where Golos references spell words out, +so their full-corpus WER is dominated by formatting. The v3_rnnt result (15 errors / 1001 words) matches the +archived PyTorch result of the garrrikkotua benchmark, hypothesis for hypothesis. + +**Speed** (RTFx = audio seconds / wall seconds, median of 3 passes after warm-up; `short` = the 200 clips one +by one, median 3.9 s; `long` = 38 segments of ≈20 s): + +| | MLX | PyTorch MPS | PyTorch CPU | +|---|---:|---:|---:| +| v3_ctc short / long | 201 / 285 | 201 / 268 | 10 / 41 | +| v3_rnnt short / long | 91 / 102 | 37 / 41 | 10 / 39 | +| v3_e2e_ctc short / long | 205 / 289 | 202 / 271 | 10 / 41 | +| v3_e2e_rnnt short / long | 99 / 112 | 45 / 48 | 10 / 41 | + +For CTC models MLX and PyTorch MPS are on par; for RNNT models MLX is 2.3–2.5× faster (the greedy +decode loop is cheaper). MLX peak memory ≈1.8 GB with the default fp32 upcast at load, ≈1.35 GB with +`load_model(..., dtype=None)` at 5–20% lower speed (see RESULTS.md). + +**Memory note.** MLX caches freed buffers for every new input shape; with varying audio lengths (and +the growing buffer of `gigaam-stream`) the cache can grow to most of unified memory. Cap it in +long-running processes: `mx.set_cache_limit(1 << 30)` (≈4% slower on short clips). `gigaam-stream` +does this. + +## Architecture + +All four models share the same Conformer encoder: + +``` +Audio (16kHz) → Log-Mel Spectrogram (64 bins) + → Conv1d Subsampling (4x stride) + → 16× Conformer Layers: + ├─ FFN₁ (half-step residual) + ├─ RoPE Multi-Head Self-Attention (16 heads) + ├─ Convolution Module (GLU + depthwise conv) + └─ FFN₂ (half-step residual) + → CTC Head (Conv1d → vocabulary + blank → greedy decode) + or + → RNNT Head (Joint + LSTM Decoder → greedy decode) +``` + +Key implementation details: +- **RoPE before projections**: GigaAM applies rotary embeddings to raw input *before* Q/K/V linear projections (non-standard) +- **RoPE base 5000**: upstream calls `RotaryPositionalEmbedding(d_model // n_heads, pos_emb_max_len)`, whose second positional parameter is `base`, so the base is 5000 rather than the usual 10000 +- **Exact mel filterbank**: Saved from PyTorch to avoid HTK recomputation differences +- **All Conv1d weights transposed**: `[out, in, K]` → `[out, K, in]` for MLX convention +- **RNNT LSTM weights**: PyTorch `(weight_ih, weight_hh, bias_ih, bias_hh)` mapped to MLX `(Wx, Wh, bias)` layout +- **SentencePiece without sentencepiece**: E2E pieces are exported to `config.json`; `decode_sentencepiece` reproduces `SentencePieceProcessor.decode` (byte fallback, ``, dummy prefix) — checked by `eval/test_decode.py` + +## License + +GigaAM model weights: [ai-sage/GigaAM](https://huggingface.co/ai-sage/GigaAM) — check their license. +MLX conversion code: MIT. diff --git a/mlx_convert/convert_gigaam_to_mlx.py b/mlx_convert/convert_gigaam_to_mlx.py new file mode 100644 index 0000000..112bca8 --- /dev/null +++ b/mlx_convert/convert_gigaam_to_mlx.py @@ -0,0 +1,246 @@ +""" +Convert GigaAM v3 PyTorch checkpoint to MLX safetensors format. + +Handles weight shape transpositions required by MLX conventions: +- PyTorch Conv1d: [out_ch, in_ch, kernel] → MLX Conv1d: [out_ch, kernel, in_ch] +- PyTorch Linear: same (no change needed, MLX nn.Linear uses same layout) +- RNNT LSTM weights properly mapped to MLX layout +- BatchNorm running stats need special handling + +Usage: + python convert_gigaam_to_mlx.py --model v3_ctc --output ./gigaam-v3-ctc-mlx + python convert_gigaam_to_mlx.py --model v3_rnnt --output ./gigaam-v3-rnnt-mlx +""" +import argparse +import json +import os +import sys + +sys.path.insert(0, "..") + +import numpy as np +import torch +from safetensors.numpy import save_file + +def transpose_conv1d_weight(w: np.ndarray) -> np.ndarray: + """PyTorch Conv1d: [out, in, kernel] → MLX Conv1d: [out, kernel, in]""" + # Safetensors requires contiguous arrays; preserve the transposed view's logical order. + return np.ascontiguousarray(np.transpose(w, (0, 2, 1))) + + +def sanitize_weights(state_dict: dict) -> dict: + """ + Convert PyTorch state_dict to MLX-compatible weight dict. + + Key transformations: + 1. Conv1d weights: transpose [out, in, K] → [out, K, in] + 2. Skip preprocessor (mel filterbank computed at runtime in MLX) + 3. Handle LayerNorm batch_norm naming + """ + mlx_weights = {} + + skipped = [] + for key, tensor in state_dict.items(): + w = tensor.detach().cpu().float().numpy() + + # Keep preprocessor mel filterbank and window for exact reproduction + if key == "preprocessor.featurizer.0.mel_scale.fb": + mlx_weights["mel_filterbank"] = w # [n_fft//2+1, n_mels] + skipped.append(key + " → mel_filterbank") + continue + if key == "preprocessor.featurizer.0.spectrogram.window": + mlx_weights["stft_window"] = w # [win_length] + skipped.append(key + " → stft_window") + continue + if key.startswith("preprocessor."): + skipped.append(key) + continue + + # ALL 3D weights are Conv1d weights and need transposition + # PyTorch Conv1d: [out, in, kernel] → MLX Conv1d: [out, kernel, in] + # This covers: encoder.pre_encode.conv.*, encoder.layers.*.conv.*, head.decoder_layers.* + if "weight" in key and len(w.shape) == 3: + w = transpose_conv1d_weight(w) + + # RNNT LSTM conversions + if "lstm.weight_ih" in key: + mlx_weights[key.replace(".weight_ih_l0", ".Wx")] = w + continue + if "lstm.weight_hh" in key: + mlx_weights[key.replace(".weight_hh_l0", ".Wh")] = w + continue + if "lstm.bias_ih" in key: + # We must add bias_ih and bias_hh + hh_key = key.replace(".bias_ih_l0", ".bias_hh_l0") + hh_w = state_dict[hh_key].detach().cpu().float().numpy() + mlx_weights[key.replace(".bias_ih_l0", ".bias")] = w + hh_w + continue + if "lstm.bias_hh" in key: + # Already handled with bias_ih + continue + + # RNNT Joint conversions + if "joint.joint_net.1.weight" in key: + mlx_weights[key.replace("joint_net.1.weight", "joint_net_linear.weight")] = w + continue + if "joint.joint_net.1.bias" in key: + mlx_weights[key.replace("joint_net.1.bias", "joint_net_linear.bias")] = w + continue + + # BatchNorm / LayerNorm: + # GigaAM v3 uses layer_norm for conv_norm_type, so batch_norm is actually LayerNorm + # The keys already use "batch_norm" name but the module is nn.LayerNorm + # In MLX, this maps to nn.LayerNorm with weight/bias — same key structure works + + mlx_weights[key] = w + + if skipped: + print(f"Skipped {len(skipped)} preprocessor keys: {skipped}") + + return mlx_weights + + +def build_config(model_name: str, cfg, tokenizer) -> dict: + """Build config.json for the MLX model.""" + from omegaconf import OmegaConf + cfg_dict = OmegaConf.to_container(cfg, resolve=True) + + config = { + "model_type": "gigaam", + "model_name": model_name, + "sample_rate": cfg_dict.get("sample_rate", 16000), + "preprocessor": { + "sample_rate": cfg_dict["preprocessor"].get("sample_rate", 16000), + "features": cfg_dict["preprocessor"].get("features", 64), + "win_length": cfg_dict["preprocessor"].get("win_length", 320), + "hop_length": cfg_dict["preprocessor"].get("hop_length", 160), + "n_fft": cfg_dict["preprocessor"].get("n_fft", 320), + "center": cfg_dict["preprocessor"].get("center", False), + }, + "encoder": { + "feat_in": cfg_dict["encoder"].get("feat_in", 64), + "n_layers": cfg_dict["encoder"].get("n_layers", 16), + "d_model": cfg_dict["encoder"].get("d_model", 768), + "subsampling": cfg_dict["encoder"].get("subsampling", "conv1d"), + "subs_kernel_size": cfg_dict["encoder"].get("subs_kernel_size", 5), + "subsampling_factor": cfg_dict["encoder"].get("subsampling_factor", 4), + "ff_expansion_factor": cfg_dict["encoder"].get("ff_expansion_factor", 4), + "self_attention_model": cfg_dict["encoder"].get("self_attention_model", "rotary"), + "pos_emb_max_len": cfg_dict["encoder"].get("pos_emb_max_len", 5000), + "n_heads": cfg_dict["encoder"].get("n_heads", 16), + "conv_kernel_size": cfg_dict["encoder"].get("conv_kernel_size", 5), + "conv_norm_type": cfg_dict["encoder"].get("conv_norm_type", "layer_norm"), + }, + } + + # Model-specific head config + if "ctc" in model_name: + config["head_type"] = "ctc" + config["head"] = { + "feat_in": cfg_dict["head"].get("feat_in", 768), + "num_classes": cfg_dict["head"].get("num_classes", 34), + } + elif "rnnt" in model_name: + config["head_type"] = "rnnt" + config["head"] = { + "decoder": cfg_dict["head"]["decoder"], + "joint": cfg_dict["head"]["joint"], + } + + config.update(build_vocabulary(tokenizer)) + return config + + +def build_vocabulary(tokenizer) -> dict: + """Vocabulary for config.json from the PyTorch model's decoding tokenizer. + + Char-wise models store their characters directly. E2E models use + SentencePiece: all pieces are exported so the MLX runtime can decode + without the sentencepiece package, plus the ids of control tokens that + decoding must drop. + """ + if tokenizer.charwise: + return {"vocabulary": list(tokenizer.vocab)} + sp = tokenizer.model + size = sp.get_piece_size() + return { + "vocabulary": [sp.id_to_piece(i) for i in range(size)], + "tokenizer_type": "sentencepiece", + "tokenizer_control_token_ids": [i for i in range(size) if sp.is_control(i)], + "tokenizer_model": "tokenizer.model", + } + return config + + +def main(): + parser = argparse.ArgumentParser(description="Convert GigaAM to MLX format") + parser.add_argument("--model", type=str, default="v3_ctc", + help="Model name: v3_ctc, v3_rnnt, v3_e2e_ctc, v3_e2e_rnnt, v3_ssl") + parser.add_argument("--output", type=str, default="./gigaam-v3-ctc-mlx", + help="Output directory") + parser.add_argument("--dtype", type=str, default="float16", + choices=["float16", "bfloat16", "float32"], + help="Output dtype") + args = parser.parse_args() + + print(f"Loading GigaAM model: {args.model}") + import gigaam + model = gigaam.load_model(args.model, device="cpu", fp16_encoder=False, use_flash=False) + + print("Extracting state dict...") + state_dict = model.state_dict() + print(f" Total PyTorch keys: {len(state_dict)}") + print(f" Total parameters: {sum(p.numel() for p in model.parameters()):,}") + + print("Sanitizing weights for MLX...") + mlx_weights = sanitize_weights(state_dict) + print(f" Total MLX keys: {len(mlx_weights)}") + + # Convert dtype + np_dtype = { + "float16": np.float16, + "bfloat16": np.float16, # safetensors doesn't support bf16 natively in numpy + "float32": np.float32, + }[args.dtype] + + # The mel filterbank and STFT window are fp32 in the checkpoint and tiny; keep them exact. + keep_fp32 = {"mel_filterbank", "stft_window"} + mlx_weights = { + k: v.astype(np.float32 if k in keep_fp32 else np_dtype) for k, v in mlx_weights.items() + } + + # Build config + print("Building config.json...") + tokenizer = model.decoding.tokenizer + config = build_config(args.model, model.cfg, tokenizer) + + # Save + os.makedirs(args.output, exist_ok=True) + + safetensors_path = os.path.join(args.output, "model.safetensors") + print(f"Saving weights to {safetensors_path}...") + save_file(mlx_weights, safetensors_path) + + config_path = os.path.join(args.output, "config.json") + print(f"Saving config to {config_path}...") + with open(config_path, "w") as f: + json.dump(config, f, indent=2, ensure_ascii=False) + + # SentencePiece models: ship the tokenizer next to the weights + if not tokenizer.charwise: + tokenizer_dst = os.path.join(args.output, "tokenizer.model") + with open(tokenizer_dst, "wb") as f: + f.write(tokenizer.model.serialized_model_proto()) + print(f"Saved tokenizer to {tokenizer_dst}") + + # Summary + total_bytes = os.path.getsize(safetensors_path) + print(f"\n✅ Conversion complete!") + print(f" Output: {args.output}") + print(f" Weights: {total_bytes / 1024 / 1024:.1f} MB ({args.dtype})") + print(f" Keys: {len(mlx_weights)}") + print(f" Config: {config_path}") + + +if __name__ == "__main__": + main() diff --git a/mlx_convert/eval/README.md b/mlx_convert/eval/README.md new file mode 100644 index 0000000..6786f38 --- /dev/null +++ b/mlx_convert/eval/README.md @@ -0,0 +1,41 @@ +# MLX evaluation: parity, accuracy, speed + +Checks the MLX port against upstream PyTorch GigaAM on a fixed, verifiable corpus. +Latest numbers: [`results/RESULTS.md`](results/RESULTS.md). + +## Corpus + +200 utterances (823.85 s) from the Golos Crowd test split — the sample published by +[garrrikkotua/gigaam-vs-whisper-russian](https://github.com/garrrikkotua/gigaam-vs-whisper-russian). +`golos_crowd_200.json` holds the row ids and per-clip SHA-256 from that repository; the audio is not +redistributed here (Golos licence) — `prepare_golos.py` downloads the Hugging Face mirror +`bond005/sberdevices_golos_10h_crowd` at a pinned revision and verifies every hash. + +References are lowercase Cyrillic, no punctuation, numbers spelled out. CTC/RNNT output already has +that form; E2E output (cased, punctuated, digits) is normalized before scoring, and the report adds +WER on the subset where no E2E transcript contains digits or Latin letters (see `common.normalize`). + +## Run + +```bash +cd mlx_convert +uv venv .venv && uv pip install --python .venv/bin/python -e "..[torch]" \ + mlx safetensors soundfile pyarrow num2words huggingface_hub pytest + +# models: fp16 (default) and fp32 copies +for m in v3_ctc v3_rnnt v3_e2e_ctc v3_e2e_rnnt; do + .venv/bin/python convert_gigaam_to_mlx.py --model $m --output eval-models/$m + .venv/bin/python convert_gigaam_to_mlx.py --model $m --output eval-models/$m-fp32 --dtype float32 +done + +.venv/bin/python eval/prepare_golos.py # corpus → eval/data/golos_crowd_200 +.venv/bin/python -m pytest eval/test_decode.py # SentencePiece decoder == sentencepiece +.venv/bin/python eval/run_eval.py # parity + WER/CER → results/summary.json +.venv/bin/python eval/bench_speed.py --repeats 3 # speed → results/speed.json +.venv/bin/python eval/make_report.py # → results/RESULTS.md +``` + +`run_eval.py` compares, per clip, the transcript and the encoder output of MLX with PyTorch on CPU +in fp32. `bench_speed.py` runs each model/backend in its own subprocess with MLX's buffer cache +capped at 1 GB and aborts above `--max-gb` of framework memory; without the cap MLX keeps freed +buffers for every new input length and can grow to most of unified memory. diff --git a/mlx_convert/eval/bench_speed.py b/mlx_convert/eval/bench_speed.py new file mode 100644 index 0000000..d6b556f --- /dev/null +++ b/mlx_convert/eval/bench_speed.py @@ -0,0 +1,200 @@ +#!/usr/bin/env python3 +"""Transcription speed: MLX vs PyTorch on the Golos Crowd sample. + +Two workloads, both from in-memory audio (no ffmpeg / file I/O timed): + +* ``short`` — the 200 clips one by one (1–13.5 s, median ≈ 3.9 s); +* ``long`` — consecutive clips concatenated into ≈20 s segments + (kept under the 25 s limit of upstream ``transcribe``). + +Each backend gets one warm-up pass over the workload, then ``--repeats`` timed +passes. Reported per pass: wall time and RTFx = audio seconds / wall seconds; +the summary uses the median pass. Model load time is measured separately. + +Backends: mlx_fp16 (fp16 file, weights upcast to fp32 at load — the +load_model default), mlx_fp16_stored (fp16 params kept), mlx_fp32 (fp32 file), +pytorch cpu fp32 and pytorch mps (upstream default for accelerators: fp16 +encoder under autocast). + +Every model/backend pair runs in its own subprocess so memory is returned to +the OS between runs. MLX keeps freed buffers in a cache that, by default, may +grow to most of unified memory when input lengths vary; the cache is capped at +1 GB here (costs ≈4% on the short workload). A run aborts if framework memory +exceeds --max-gb. Peak framework memory is reported: MLX peak allocation, +MPS driver allocation, CPU max RSS. +""" +import argparse +import json +import resource +import statistics +import subprocess +import sys +import time +from pathlib import Path + +import mlx.core as mx +import numpy as np +import torch + +import gigaam +from common import DEFAULT_DATA, HERE, load_corpus, machine_info +from gigaam_mlx import load_model as load_mlx + +MODELS = ["v3_ctc", "v3_rnnt", "v3_e2e_ctc", "v3_e2e_rnnt"] +BACKENDS = ["mlx_fp16", "mlx_fp16_stored", "mlx_fp32", "torch_cpu", "torch_mps"] + + +def long_segments(clips, target=20.0, limit=24.0): + segments, cur, dur = [], [], 0.0 + for c in clips: + if dur + c.duration > limit and cur: + segments.append(np.concatenate(cur)) + cur, dur = [], 0.0 + cur.append(c.audio) + dur += c.duration + if dur >= target: + segments.append(np.concatenate(cur)) + cur, dur = [], 0.0 + return segments # the last partial segment is dropped + + +def make_backend(name, backend, models_dir: Path): + t0 = time.perf_counter() + if backend.startswith("mlx"): + path = models_dir / (name + ("-fp32" if backend == "mlx_fp32" else "")) + # mlx_fp16: fp16 file, upcast at load (load_model default); _stored keeps fp16 params + model = load_mlx(path, dtype=None) if backend == "mlx_fp16_stored" else load_mlx(path) + + def run(audio): + return model.transcribe(mx.array(audio)) + + else: + device = "cpu" if backend == "torch_cpu" else "mps" + model = gigaam.load_model(name, device=device, fp16_encoder=device != "cpu", use_flash=False) + model.eval() + + def run(audio): + with torch.inference_mode(): + wav = torch.from_numpy(audio).to(model._device).to(model._dtype)[None] + length = torch.full([1], wav.shape[-1], device=model._device) + encoded, encoded_len = model.forward(wav, length) + text = model._decode(encoded, encoded_len, length)[0][0] + if device == "mps": + torch.mps.synchronize() + return text + + load_s = time.perf_counter() - t0 + return run, load_s + + +GB = 2**30 + + +def framework_bytes(backend): + if backend.startswith("mlx"): + return mx.get_active_memory() + mx.get_cache_memory() + if backend == "torch_mps": + return torch.mps.driver_allocated_memory() + return resource.getrusage(resource.RUSAGE_SELF).ru_maxrss # bytes on macOS + + +def peak_bytes(backend): + if backend.startswith("mlx"): + return mx.get_peak_memory() + if backend == "torch_mps": + return torch.mps.driver_allocated_memory() + return resource.getrusage(resource.RUSAGE_SELF).ru_maxrss + + +def timed_pass(run, audios, backend, max_bytes): + lat = [] + t0 = time.perf_counter() + for a in audios: + s = time.perf_counter() + run(a) + lat.append(time.perf_counter() - s) + if framework_bytes(backend) > max_bytes: + raise SystemExit(f"{backend}: framework memory above {max_bytes / GB:.1f} GB, aborting") + return time.perf_counter() - t0, lat + + +def run_one(name, backend, args, workloads, seconds): + mx.set_cache_limit(GB) + run, load_s = make_backend(name, backend, Path(args.models_dir)) + entry = {"load_seconds": load_s} + max_bytes = args.max_gb * GB + peak = 0 + for wl, audios in workloads.items(): + timed_pass(run, audios, backend, max_bytes) # warm-up + passes = [timed_pass(run, audios, backend, max_bytes) for _ in range(args.repeats)] + peak = max(peak, peak_bytes(backend)) + walls = [w for w, _ in passes] + med = statistics.median(walls) + lat = sorted(passes[walls.index(med)][1]) + entry[wl] = { + "wall_seconds": walls, + "median_wall_seconds": med, + "rtfx_median": seconds[wl] / med, + "latency_ms_p50": 1000 * lat[len(lat) // 2], + "latency_ms_p90": 1000 * lat[int(len(lat) * 0.9)], + } + entry["peak_memory_gb"] = peak / GB + return entry + + +def main(): + parser = argparse.ArgumentParser(description=__doc__.splitlines()[0]) + parser.add_argument("--data", default=str(DEFAULT_DATA)) + parser.add_argument("--models-dir", default=str(HERE.parent / "eval-models")) + parser.add_argument("--out", default=str(HERE / "results" / "speed.json")) + parser.add_argument("--models", nargs="+", default=MODELS) + parser.add_argument("--backends", nargs="+", default=BACKENDS) + parser.add_argument("--repeats", type=int, default=5) + parser.add_argument("--max-gb", type=float, default=6.0, help="abort a run above this framework memory") + parser.add_argument("--one", help="internal: run a single model/backend and print JSON") + args = parser.parse_args() + + clips = load_corpus(Path(args.data)) + workloads = {"short": [c.audio for c in clips], "long": long_segments(clips)} + seconds = {k: sum(len(a) for a in v) / 16000 for k, v in workloads.items()} + + if args.one: + name, backend = args.one.split("/") + print(json.dumps(run_one(name, backend, args, workloads, seconds))) + return + + out = Path(args.out) + out.parent.mkdir(parents=True, exist_ok=True) + result = json.loads(out.read_text()) if out.exists() else {"runs": {}} + result["machine"] = machine_info() + result["protocol"] = { + "repeats": args.repeats, + "warmup": "one full pass per workload", + "torch_threads": torch.get_num_threads(), + "mlx_cache_limit_gb": 1, + "isolation": "one subprocess per model/backend", + "workloads": {k: {"items": len(v), "audio_seconds": seconds[k]} for k, v in workloads.items()}, + } + + for name in args.models: + for backend in args.backends: + key = f"{name}/{backend}" + cmd = [sys.executable, __file__, "--one", key, "--repeats", str(args.repeats), + "--max-gb", str(args.max_gb), "--data", args.data, "--models-dir", args.models_dir] + proc = subprocess.run(cmd, capture_output=True, text=True) + if proc.returncode != 0: + print(f"{key}: FAILED\n{proc.stderr[-2000:]}", flush=True) + continue + entry = json.loads(proc.stdout.strip().splitlines()[-1]) + result["runs"][key] = entry + print( + f"{name:12s} {backend:9s} load {entry['load_seconds']:5.1f}s | short RTFx {entry['short']['rtfx_median']:6.1f} " + f"p50 {entry['short']['latency_ms_p50']:6.1f} ms | long RTFx {entry['long']['rtfx_median']:6.1f} " + f"p50 {entry['long']['latency_ms_p50']:6.1f} ms | peak {entry['peak_memory_gb']:.2f} GB", + flush=True, + ) + out.write_text(json.dumps(result, indent=2)) + + +if __name__ == "__main__": + main() diff --git a/mlx_convert/eval/common.py b/mlx_convert/eval/common.py new file mode 100644 index 0000000..e87fad1 --- /dev/null +++ b/mlx_convert/eval/common.py @@ -0,0 +1,130 @@ +"""Shared helpers for the MLX eval: corpus loading, text normalization, WER.""" +import json +import re +import sys +from dataclasses import dataclass +from pathlib import Path +from typing import Dict, List + +import numpy as np +import soundfile as sf +from num2words import num2words + +HERE = Path(__file__).parent +DEFAULT_DATA = HERE / "data" / "golos_crowd_200" +sys.path.insert(0, str(HERE.parent)) # gigaam_mlx + + +@dataclass +class Clip: + id: int + path: Path + duration: float + reference: str + audio: np.ndarray # float32, 16 kHz mono + + +def load_corpus(data_dir: Path = DEFAULT_DATA) -> List[Clip]: + clips = [] + for line in open(data_dir / "manifest.jsonl"): + rec = json.loads(line) + path = data_dir / rec["audio"] + audio, sr = sf.read(path, dtype="int16") + assert sr == 16000 and audio.ndim == 1, path + clips.append( + Clip( + id=rec["id"], + path=path, + duration=rec["duration"], + reference=rec["reference"], + # same scaling as gigaam.preprocess.load_audio (s16le / 32768) + audio=audio.astype(np.float32) / 32768.0, + ) + ) + return clips + + +# ─────────────────────────── normalization ─────────────────────────── +# Golos references are lowercase Cyrillic words, no punctuation, no "ё", +# numbers spelled out. CTC/RNNT output already looks like that; E2E models +# output cased text with punctuation and digits. Every hypothesis goes through +# the same function, so char-wise outputs are (near) unchanged. + +_NUMBER = re.compile(r"\d+(?:[   ]\d{3})*(?:[.,]\d+)?") + + +def _number_to_words(match: re.Match) -> str: + token = re.sub(r"[   ]", "", match.group(0)).replace(",", ".") + try: + value = float(token) if "." in token else int(token) + return " " + num2words(value, lang="ru") + " " + except (ValueError, OverflowError, NotImplementedError): + return " " + token + " " + + +def normalize(text: str) -> str: + text = text.lower().replace("ё", "е") + text = _NUMBER.sub(_number_to_words, text) + text = re.sub(r"[^\w\s]|_", " ", text) # punctuation, hyphens, symbols + return " ".join(text.split()) + + +# ─────────────────────────── metrics ─────────────────────────── + + +def edit_ops(ref: List[str], hyp: List[str]) -> Dict[str, int]: + """Levenshtein alignment counts: substitutions, insertions, deletions.""" + n, m = len(ref), len(hyp) + # dp[i][j] = (cost, S, I, D) + prev = [(j, 0, j, 0) for j in range(m + 1)] + for i in range(1, n + 1): + cur = [(i, 0, 0, i)] + for j in range(1, m + 1): + if ref[i - 1] == hyp[j - 1]: + best = prev[j - 1] + else: + c_sub, c_ins, c_del = prev[j - 1], cur[j - 1], prev[j] + best = min( + (c_sub[0] + 1, c_sub[1] + 1, c_sub[2], c_sub[3]), + (c_ins[0] + 1, c_ins[1], c_ins[2] + 1, c_ins[3]), + (c_del[0] + 1, c_del[1], c_del[2], c_del[3] + 1), + ) + cur.append(best) + prev = cur + cost, s, i_, d = prev[m] + return {"errors": cost, "substitutions": s, "insertions": i_, "deletions": d} + + +def corpus_error_rate(pairs, unit: str = "word") -> Dict[str, float]: + """Corpus-level rate: total edits / total reference units (not a mean of + per-utterance rates), matching the garrrikkotua benchmark's verify.py.""" + total = {"errors": 0, "substitutions": 0, "insertions": 0, "deletions": 0} + ref_units = 0 + for ref, hyp in pairs: + r = ref.split() if unit == "word" else list(ref.replace(" ", "")) + h = hyp.split() if unit == "word" else list(hyp.replace(" ", "")) + for k, v in edit_ops(r, h).items(): + total[k] += v + ref_units += len(r) + return {**total, "ref_units": ref_units, "rate": total["errors"] / ref_units} + + +def machine_info() -> Dict[str, str]: + import platform + import subprocess + from importlib.metadata import version + + def sysctl(key): + return subprocess.run( + ["sysctl", "-n", key], capture_output=True, text=True + ).stdout.strip() + + return { + "chip": sysctl("machdep.cpu.brand_string"), + "memory_gb": str(int(sysctl("hw.memsize") or 0) // 2**30), + "macos": platform.mac_ver()[0], + "python": platform.python_version(), + "mlx": version("mlx"), + "torch": version("torch"), + "torchaudio": version("torchaudio"), + } diff --git a/mlx_convert/eval/golos_crowd_200.json b/mlx_convert/eval/golos_crowd_200.json new file mode 100644 index 0000000..4eb9a41 --- /dev/null +++ b/mlx_convert/eval/golos_crowd_200.json @@ -0,0 +1,808 @@ +{ + "source": "https://github.com/garrrikkotua/gigaam-vs-whisper-russian @ 95825b4 (results.json, audioManifest)", + "dataset": "bond005/sberdevices_golos_10h_crowd", + "revision": "e634b6b810e4d30c81b4c6d8262379fe8b9f708c", + "file": "data/test-00000-of-00003-cad0cfaddbc8fa71.parquet", + "items": [ + { + "id": 0, + "sha256": "500d2e88a1634ee766fec46f83cb903f3a1c276b2cedb36f11db800a1c4c14d0" + }, + { + "id": 1, + "sha256": "63666d51b938f5c9acf4d92e9984251fff122d9425232c75ec45ee81cf7aab1c" + }, + { + "id": 2, + "sha256": "101105b34eb3dbad83eb27849218c1b56756b2265f85d1922f32983876462e1f" + }, + { + "id": 3, + "sha256": "d5e5002ab533245a61757ef267c75c298dfc7a22904fdb5385374ade2276a240" + }, + { + "id": 4, + "sha256": "606918e23e0894f7c330441d3e8e4eacb37374722f95fc0a73cf91515a4cc758" + }, + { + "id": 5, + "sha256": "ed0f1cfe5a6d48694e54d8a3357901acbc0f1113dc846ee01cb00a405a112c04" + }, + { + "id": 6, + "sha256": "eb3042fedddbe16c5c5857b3ae9443d464a42a2d0bbc829a61714af5e7b240f8" + }, + { + "id": 7, + "sha256": "df9e4058f32ed3fe8cacdf9e31e92c7b6c2f49cd8adf5f50f853086ec4207867" + }, + { + "id": 8, + "sha256": "23e2fcf0734d97acd6837bd85591ce998e385f5b32bd8a55ef217b689f00d8a7" + }, + { + "id": 9, + "sha256": "a23daa8fa6529d02a7e931fb507ac3be998de6c3e10475c693d4b7fc2f086b68" + }, + { + "id": 10, + "sha256": "fd681926b2d8087d1772246d6c95110695f253190b0cef581d25c80977477d81" + }, + { + "id": 11, + "sha256": "98e3db631196538177d6e1fc1db34da93bd68e6e820b1a5b73df3bff59177f72" + }, + { + "id": 12, + "sha256": "d9b6fea6d4b9cef70df3241d4280d5dcf0ceadc0a6bd2ca240e8ee30cbef5e4b" + }, + { + "id": 13, + "sha256": "a84b9e37c1474cc57a34a9a5e30215864c52c4e9e98271b634a42541163370c0" + }, + { + "id": 14, + "sha256": "5711ab2be5b07eb1c0c2af925f5586da5db3805cad0a0b7c1a8b84a03b894cb8" + }, + { + "id": 15, + "sha256": "ab3baa250e7ee404dec2a0bfd6f3ce58640ad8fca31cb27b8deaddc1a9273e87" + }, + { + "id": 16, + "sha256": "adda3e2b2ff797d722bf9efc1f2ad8376e6a2cf57fc81e306590a77c90a71dab" + }, + { + "id": 17, + "sha256": "b71d8b3ab34aa92fa7afff00030fd9f2f3babb5a7133a1cbcf223caf2fc5cb1b" + }, + { + "id": 18, + "sha256": "172d18b6e87cd0e5829f3d46fa3dcb646e3f81d23427d874d9eb2411b926c77a" + }, + { + "id": 19, + "sha256": "14d9ce4d3b9a34123898e0209adcb43b069f0dda5f7e5a330530b81dea28d253" + }, + { + "id": 20, + "sha256": "35d623f1eafe146eedb543942fc938604b11a234f9e2e41fb16c88dff6120dd9" + }, + { + "id": 21, + "sha256": "acba9cab8c22bd45947279c7fb60f57c975c6013e3461d988f6e088fed101398" + }, + { + "id": 22, + "sha256": "7a924fb233b586aa09e5e115819f3717881d2b8d8d8bba2371c5c77f29c70126" + }, + { + "id": 23, + "sha256": "f47a1d829a72b9a8179114dfa7010042886b1479a25a88fbe5106c1e22f392e4" + }, + { + "id": 24, + "sha256": "5d2607aebbd820aad85c254072419e45997f569f73b5faab31803d8f6f2d07be" + }, + { + "id": 25, + "sha256": "bf9b20cd89c9a5cb4c7e483fe88c49c8e6b319e8d5355d5538dcf9fd4c560dce" + }, + { + "id": 26, + "sha256": "f05aed608b6f2aa64864a8fbe32e779e48e4b2f67322942b25302debbc1cfa04" + }, + { + "id": 27, + "sha256": "5d0876bb2e953e8650fdfbd06b1d806e3bc2dbc8064c1d6267618caf6cada7d1" + }, + { + "id": 28, + "sha256": "d160f269f2c783dbac72f90276383a4a4f8f97f2844cf143e29d1728acb12ab3" + }, + { + "id": 29, + "sha256": "0fc1063dc3794309d112918552ab0d8bbf3263173adbc7842f5488afe65556e1" + }, + { + "id": 30, + "sha256": "4d905da4f73550e462ddac78e6c032ef75f31c4835eadd4173a973aff4f208b4" + }, + { + "id": 31, + "sha256": "2b6c5ea81234eb6987255e9ee14ca612e1057b86f3128d2dabe0f8d484f02c38" + }, + { + "id": 32, + "sha256": "2569924b9ef36b8ff2f46c989f352d07bd1cff0f0e3cca54e9bb7e1d6cfdc79b" + }, + { + "id": 33, + "sha256": "7a7483f4186d11cbd2f0fe3d09b4b24c72150851475378c0d6e559643cf2e6cf" + }, + { + "id": 34, + "sha256": "0fced7e1495bbae50b70d02fd6e1b13291b637564a07ba16342994c673140c0e" + }, + { + "id": 35, + "sha256": "5e4d0e24e6866c94a7891d5b95d3228037381ed1a515af490715477d0b0bfed1" + }, + { + "id": 36, + "sha256": "335f8e78705bc19c242bb2f120050df13295f0c8cb4eef294738a9fc7a4c4ed8" + }, + { + "id": 37, + "sha256": "f62b2037d23ceaeaa22968896f6adf8f7c690af6f013d2a91af1bf2056fcb826" + }, + { + "id": 38, + "sha256": "bf52b5a0327e31c6f41671323dad7b3ec4c15e898ba71d817c83598ac68bd2bb" + }, + { + "id": 39, + "sha256": "cd273cec2f341650985078bd5844be95e708717b42d88a811a61f8c6f3ccfddb" + }, + { + "id": 40, + "sha256": "b06ec3d91f55e25d0f5b09b22485e62821005b9cebb225dbc9bc5e2310a325bb" + }, + { + "id": 41, + "sha256": "b5b89a6404aa2eba7c65534ecc10c3f99cf750f49c826c3d7bd2fe2f2daa3ec6" + }, + { + "id": 42, + "sha256": "52c4c3d0c9518d5f2590324e992e796451eff231ccb55729afdc3fb42e97e55e" + }, + { + "id": 43, + "sha256": "de46b81590716605d89aa83b2a3b5c687f72c65a60ed3aabde26af786e8a77b0" + }, + { + "id": 44, + "sha256": "bd0db0c3eecaa28d151c50058b97ae4cdc5ea2ba701704c08854a57f33b42973" + }, + { + "id": 45, + "sha256": "d2cdfac855bb4e2c74773dfa5f9faf2756d5dddf6cad029275b4d75faeac551d" + }, + { + "id": 46, + "sha256": "223fccc831db4acc9f38e8f9c429bc7a13869e266fba6318405a058b0b6b1033" + }, + { + "id": 47, + "sha256": "fb44702ead2bb3700195aaa9d98a895070dddbbd08958294104bc44db70fb2ae" + }, + { + "id": 48, + "sha256": "b58333a800b2b8ea59526ec52755d8afefc84e2733dc7faa5dd982b334f449f5" + }, + { + "id": 49, + "sha256": "64b193332ced6c4a00deeeae8ef4d67b0ddc8d08f0a6ea41311ac8be57cedd34" + }, + { + "id": 50, + "sha256": "7f41b63239e2ff8ba16d7d839872d6d9ec3414fb5094a950492cb3209254c805" + }, + { + "id": 51, + "sha256": "1d167708585b04a0943968703095201fc21d3199c8dab18be888c358c1325991" + }, + { + "id": 52, + "sha256": "69f9df2a1257716296ade320eafde80360babf5a18489f717a8af0a8b4479108" + }, + { + "id": 53, + "sha256": "775a99e6b1b50f297e58a0591272bf04d8959ac4db68bb8c10d8028c5f9473b1" + }, + { + "id": 54, + "sha256": "d6dbce1fbc4e3763a0f7b324984c5ad2ee6f9f2407b07d64c30fd09d3314e077" + }, + { + "id": 55, + "sha256": "53e0a54538483c5532a42a2cc98146a9e3a1210a82a827e58390787982d11095" + }, + { + "id": 56, + "sha256": "fc858a35b1c661dc18f6be9cdb1631bf418984bfab725f5954ec535eeffd9914" + }, + { + "id": 57, + "sha256": "115fbc30fad01f95eb71a743f457b2f98dd1f600ba554bbbefc76fbfbcf84a6c" + }, + { + "id": 58, + "sha256": "ccf650c058437d0324f8ca2f221382c343e0acec3f54f834ea6887625ec0072b" + }, + { + "id": 59, + "sha256": "9c950b57332e1c4a4e8c4cb2bd42ddf939d3af41f75b47b438ded7dc0769e4b7" + }, + { + "id": 60, + "sha256": "2e8dc87a18d361b9165364ddccdafa34d516d801d7fb3161bbf02e935e3f69d4" + }, + { + "id": 61, + "sha256": "d387d682c367b488bee3ea5222930e1f6db45e318b0c7209cd5a905e8fa388d9" + }, + { + "id": 62, + "sha256": "74145db99dd89b7f860afbf7d2e71d0615facb5cd604d21f5082928e07436906" + }, + { + "id": 63, + "sha256": "0d6e30e25182487ea9256ea738c151fd560b58219016ac4c367a858388abb6fc" + }, + { + "id": 64, + "sha256": "5fbc796647964be8777c768eddd8999376dafce162b036380e6cbb9b617af674" + }, + { + "id": 65, + "sha256": "3e07cb730dfcb90405c67db6a28fc00fc618fb2dff9e2c4e7d7ff45ac8391ab1" + }, + { + "id": 66, + "sha256": "42a215f7adf7baf37453bf2bf96a79effac0e04c3fed2f7ecff5c33f58300337" + }, + { + "id": 67, + "sha256": "a6c04f518c5019998465f3e073adf15da3a2fb8ab155a55c4cd86f2e66585a24" + }, + { + "id": 68, + "sha256": "23e5ce5955e734a502ee84f467cf1eddc8f0cbd0447b56fe8e99a0701db6ce9e" + }, + { + "id": 69, + "sha256": "375398861454441e225bddb17926ed0231493e8810feffd9d593738c17da6ad0" + }, + { + "id": 70, + "sha256": "5aa0098d4e84bb15c239915c940ead456965a8529180cd555780efab9d6bda1a" + }, + { + "id": 71, + "sha256": "7464fcd09822037995c67343727e8d3af47623a08a71f61d3b3c92f1dd9434fd" + }, + { + "id": 72, + "sha256": "242e236289bba2c4bb47d177fa0efa905dd5c5c7b3385dc37364950f828d704a" + }, + { + "id": 73, + "sha256": "6480a503231954a9e8b232c27d8feccca27e90758ed403a621c2c1b1349f12cc" + }, + { + "id": 74, + "sha256": "931cc61cfd1c789b6e674d04a3d661feec63dc363b88b2552d44a991ba1d0a3a" + }, + { + "id": 75, + "sha256": "5fc13db2895ff8e2f38a1a131e0d2cb56de9aedcc7dea8f9a73ab70b6f439c23" + }, + { + "id": 76, + "sha256": "602a5b9b1bfe3b48d657d7a1317cf10b27b5655ba0573414927f5423bdd8630c" + }, + { + "id": 77, + "sha256": "fbac536c8f766e29bfa613515c988d33d5f86a4103e91588b6fea6c9b9365149" + }, + { + "id": 78, + "sha256": "981beeb891587cbaff1420e303c0d646640426aed72343bd52c5f5c17d72861d" + }, + { + "id": 79, + "sha256": "8ac8651bc9925c312d2e13b6d29ac78e6e8f6ed4ad7211e0623db105d461b097" + }, + { + "id": 80, + "sha256": "de8a0bc25d2f640857b0c3b5a9ec8035f8ea39eeccfd390bf825700dccb1ce23" + }, + { + "id": 81, + "sha256": "613b3ea0ded232963375cac3a34fda77cbc02f90a9ccfb83fdf2aae6cfed5480" + }, + { + "id": 82, + "sha256": "e2464ddc5cacd21d266a21c03d8b7dea965148e391b57ea8ddae452e037b0ed7" + }, + { + "id": 83, + "sha256": "c3f684a21e59efb9ec5ab44aa745e78be354bf1e120805ccf4e7eed6fd701b14" + }, + { + "id": 84, + "sha256": "288ce18300e81b934349728a44ee9ee3f5e98dc3156541a2b567767574afc361" + }, + { + "id": 85, + "sha256": "f802c8ef4d3f71284e7119ab1bc09096f8e2a77dae88272107a6d83c4ce28f2b" + }, + { + "id": 86, + "sha256": "b34dc4647771960e56bd2e1ae84edd1af166560de8bca5fe48cd6d0572eab306" + }, + { + "id": 87, + "sha256": "b10aa0bb225508f82783c601a39dc94d49be67eca30dbdf2e6f7e434061a6112" + }, + { + "id": 88, + "sha256": "fa0ac78b24fe1cd4a14bc4dc59c3f95ae700e5a80777c8cf51d59aec94c0fbf6" + }, + { + "id": 89, + "sha256": "5eddd9a9f041e76fb50d05a9df3e24a35a281e31c436038c01c102b6c3242e08" + }, + { + "id": 90, + "sha256": "51011926aa3bdda985041365e96f11a48018722b01f05e65d4abebaa02a7e015" + }, + { + "id": 91, + "sha256": "bb156e61da7c844b08e080b3875ce5ca6a8d893b05b1cd6376a91c239878fbc9" + }, + { + "id": 92, + "sha256": "92e36cf55f0813484df99e6fdd6bec8f0356136f98a9eada8f7f6dd8e1a4222b" + }, + { + "id": 93, + "sha256": "cae87195a11932ed11a25e0fc68147a95fe56aabc009db4709ba1743b524870f" + }, + { + "id": 94, + "sha256": "24d20244a042d2d54ed199dcb2b398c00a924834083aeaac1302dd958f4d09d0" + }, + { + "id": 95, + "sha256": "3515956bcccb860b6277a1c0a3b914c7594f38f6ceb7fdec7bb05674702912ef" + }, + { + "id": 96, + "sha256": "a6f4f74672616d45c64ecfd0a363f7f22d71703e0ce3e71c68d209abf2a1aeb5" + }, + { + "id": 97, + "sha256": "388d7c9c5893a8a0e0e0f68fa6323a27c165dabc434b055a9d12a65a8b7ea4f3" + }, + { + "id": 98, + "sha256": "ccd34d25a92098b2af2936960083c0284d69ad9bcc59c31a821a70b67d44500f" + }, + { + "id": 99, + "sha256": "6af3cbc1c0a5840d062458fa3a1055af6caf0fda524d768ab8e6ae4ac1e2359d" + }, + { + "id": 100, + "sha256": "3377d55ebd45fef0a46bd8a1a8f9efd83095bc5e60a7078f48c4c55488b8d4b4" + }, + { + "id": 101, + "sha256": "fe014dc975da21b01d072c460eb970493f0a84a307acfd049b46313fa6cb7d79" + }, + { + "id": 102, + "sha256": "eed9a10400b3707107001ccba01ef6b6f34b274c9f524ff43796620316849346" + }, + { + "id": 103, + "sha256": "821ce8af9dcec02d85ed4670e3c171f827b7346f56ba1e4f217bde06fe850dad" + }, + { + "id": 104, + "sha256": "a6421ee533f6a30dd64084292739e017a183103f008b9227593033aea31eb4ba" + }, + { + "id": 105, + "sha256": "f8fddb633987b9f3a96d7f13b78e34023a058c983f6dacd97aca060e83089012" + }, + { + "id": 106, + "sha256": "2b5198ae88e2a2b2968f205c623908ad0e21b78614fe9c3fb97e02f34305d082" + }, + { + "id": 107, + "sha256": "f487e760ad8f558aa03e273fa52ab594f625176327aa8e99dab783a4ddaf7202" + }, + { + "id": 108, + "sha256": "5b27bd8e37d42ba0309dff4d0a3555d85b0a00859b131ffaf2073b8312454628" + }, + { + "id": 109, + "sha256": "7c7e12e38ccb8e9f6888fbb5173b9ba67c4b0f2abc72b6d1d14eb2e23e7566e6" + }, + { + "id": 110, + "sha256": "17752a71b589e1bae94d068fbec85ca1bf4787fc2e39183e164b0292acca8b1a" + }, + { + "id": 111, + "sha256": "a15540ac78bebb46cebcdf7f59869edd9e16b1271d05d832c27da2938d8fd009" + }, + { + "id": 112, + "sha256": "022c11c89e4b0eb474ebe4fb2aa3d0d2cf4320bc2e2f4e6ed5ef066403cad08b" + }, + { + "id": 113, + "sha256": "280e922914a3b23f69ce7b0e95b1f826aa98e6ac11a00d06576ce48376504cea" + }, + { + "id": 114, + "sha256": "7e4a5bb81d7c80c38d6b8c41f87c0014d948425966f87e9d8e1cb78f376a775b" + }, + { + "id": 115, + "sha256": "9ef66c275b7638e18f01c88205165bd61bb05bbe94504fdfc1c4eb5ad7665ba9" + }, + { + "id": 116, + "sha256": "bfcd734ac8b58dcc703e1ee12ebafef0249842d1ef49b5fdd0f198fbdb15594b" + }, + { + "id": 117, + "sha256": "ef79fdebd30f70c02e0362db242e547c088132052451d37b500d0d0622e823bc" + }, + { + "id": 118, + "sha256": "020b6a33b9c24b5efeb591b20a94034fbbd83b0b1fd0bb6bd1f340e170c6706a" + }, + { + "id": 119, + "sha256": "2a8245f847283fab1a0b63185e6343099293bf385bde9513a149351038ba591c" + }, + { + "id": 120, + "sha256": "43b18dcaa5f0b046b6d23e2ceb2d779321d15307e70a01b003aaa066dd224769" + }, + { + "id": 121, + "sha256": "dbb767248239ee0cc05931dbfe98ea5d87e2b7e3186ef45de820dde4fb422952" + }, + { + "id": 122, + "sha256": "18acdd7647922ea9b5d5a1d28544394c190325c08da4159acfb5f903d9d04775" + }, + { + "id": 123, + "sha256": "d348be3d619fbc103323ed99f5d087c25a26961becd6f7b0e5718ec0ecf6bbc5" + }, + { + "id": 124, + "sha256": "4496f464d0595d7ed10ee0789ff262a70c6dbdb1816ba4b6e83e472d312f7b68" + }, + { + "id": 125, + "sha256": "6824eaa18939b0702d81c840a7e2b08c82add5a8784162f9b2c6770512e2c764" + }, + { + "id": 126, + "sha256": "00f8194139d934d083c593949168c49e3f5949cc9362ac38db7c874486a7b127" + }, + { + "id": 127, + "sha256": "10c4714d62b5bfb5eaba8bab16132d558dd9361208d5513fc8031f42b78b772b" + }, + { + "id": 128, + "sha256": "59813b3a52371ef53ce63ce0e3903fdf93696f0d6cfd6f140cd957565874b71c" + }, + { + "id": 129, + "sha256": "bb0bb58c6ec6525b02c80ecb5b37e472a4e00304e23efb9fd7ec8db5432c535b" + }, + { + "id": 130, + "sha256": "17df9daeb92e01e4ba09e24fd2ca0ffe10b59d756d3bc2c1097150f2c46dbc55" + }, + { + "id": 131, + "sha256": "5c66332eac01786fe110f666442ad9cbb6d7bd82bd60cb4e65e1c7de2c0b0abd" + }, + { + "id": 132, + "sha256": "ec404da97fbdf37172461aeea048838c5c4f4389e7009c62a8ca2113a7733c25" + }, + { + "id": 133, + "sha256": "bb4eba68d8c7250ab2acbdc920eb3ed46c04c8daa088bb2b0ea79a9cc80bbf6c" + }, + { + "id": 134, + "sha256": "d49dc288d1aaaa92241860e9d9b985736c4bdfb9594daeb04d4c2cfdffa395f6" + }, + { + "id": 135, + "sha256": "06bc5898b49d435a5ac86ecaf88f66e47e5b5aee568d68b56a3d8c844649762d" + }, + { + "id": 136, + "sha256": "8478f218474b682617fbed2b5f8db1c8ca451acdb3064ccf7131ae0c40ba018a" + }, + { + "id": 137, + "sha256": "6bb0a2ad2c261c2e87478735e550ad40cd2a0d5a64c32045afe281c98201891e" + }, + { + "id": 138, + "sha256": "d810a2938a082f363c20b0602337ea5f8302c3663c6b28b368822aac930fd540" + }, + { + "id": 139, + "sha256": "738a27d1bedb72881904f099bd04a79194dffcd8e2da46279ca87b10d46e4814" + }, + { + "id": 140, + "sha256": "bc424ca5503816c37e2ab58da42e47e2b9c38e2633cb057f3d4f14ddf18b6f56" + }, + { + "id": 141, + "sha256": "6f151a86a4667f79cc6c66121f1df1afbab71909a61c197a19a26a9bd2e7c524" + }, + { + "id": 142, + "sha256": "0125acb2fe6561bf2e525c5b7fad186170d1f463c21362a43fa53941434d616f" + }, + { + "id": 143, + "sha256": "a1a6fc8f769fd6e0e5e784b39ebfbb30eda27d31cf0c1607b5e00c0b06dfd311" + }, + { + "id": 144, + "sha256": "7ccf7841a7cf76dd9e74512c4bf4e7abcb16f9807f3c3e6c20d04bd260e78058" + }, + { + "id": 145, + "sha256": "4ed320dab5df2a4038aaeca382ee5981cdc2e9a3c6bec762507c3611f5a33b48" + }, + { + "id": 146, + "sha256": "8c07f1f2acd5238fded1267f837fdfcad3305b7491cf5430ba04394f3dc7a8e6" + }, + { + "id": 147, + "sha256": "c2144647c15742a314755ef009289906a890c17dac8ac7d810691b1de4daa4d7" + }, + { + "id": 148, + "sha256": "cd707b8ae97211495b40d0206abc7f19875ee423bb4ac9a1c04cfca73183d64d" + }, + { + "id": 149, + "sha256": "0965d1ae29d7ef2690b8df37d6740fed3adc05e454c819783b8bcee961309874" + }, + { + "id": 150, + "sha256": "20e4988fbf4c8fe56945961836138eea0a67e47c3aa55430c4d07d865800c48c" + }, + { + "id": 151, + "sha256": "c4a2aec91039222dce8ecdd5eee9f0ee8bfd8e8d9a0e251c53e84bfba6f23599" + }, + { + "id": 152, + "sha256": "6147909feab3c260822fb7d3cce5697db58ff1b2e8e4476b2d2ee4a7be90dd12" + }, + { + "id": 153, + "sha256": "618150ccc6cdcb26f721c0206cb5ad08b5bfe9375b615e404dbade768a3e0437" + }, + { + "id": 154, + "sha256": "392e724e247acacb601a12765dddc83c55e284a4dc695c1e0852a1bdd9ef4d57" + }, + { + "id": 155, + "sha256": "aa141f6dec367725efc039311ff80bf39855e9a81e07b8618a04b461664cfca0" + }, + { + "id": 156, + "sha256": "42fc47d3299c31ebd5bbcb7ec5b8bc52ee87b57785fb6e1e3a1db7ae98fa5f53" + }, + { + "id": 157, + "sha256": "ad26a985fb0be1b76a33b724a6eb4c121b512e73210d269c36d6584acc8d028b" + }, + { + "id": 158, + "sha256": "63b399300d48335543c157e92a083860f05c58bd488c7f9edb38cb4b2dc33ada" + }, + { + "id": 159, + "sha256": "c94c77bef4f3568160bc46422b2a2bfb03d09c295b1a1b5f2474ed99dec10149" + }, + { + "id": 160, + "sha256": "0f76366fad4b0e146294017cb504ab2e7be3e24bd164b804528f51eccf2701f2" + }, + { + "id": 161, + "sha256": "9b4abc98b7c5a770070b893fea46bf36928388d95616e02adbedb714c28cf1a4" + }, + { + "id": 162, + "sha256": "979e012bed5b269068448d41cfa6dae654fa2f316e47133ab7c7fcb7cc1aac58" + }, + { + "id": 163, + "sha256": "136d4b63d272ba9581e50bd896f850b6a7b7f7f1c03cec07d35b674d8cf8b0b8" + }, + { + "id": 164, + "sha256": "da954845412c8e28cbe2dbc5e27013cc8dcbeff2b67e917400541050d2fe65a3" + }, + { + "id": 165, + "sha256": "d63cc0b0e6dc95d08c0e7edb919b972a17eb771483a815acc41b814dbf1884a5" + }, + { + "id": 167, + "sha256": "f28cb82fd7347ec98bff944a72a20a74b6af4edf275c5cdd9af0aa21cf427b2d" + }, + { + "id": 168, + "sha256": "9b10eef08d4458fa3e05c91e7b0bce1e2e780affce9ab40fcb68b81ff6c2b270" + }, + { + "id": 169, + "sha256": "54903796d177c9ec9b949dcc21b60abbc1671b990851f414b5288e8603a3e1a3" + }, + { + "id": 170, + "sha256": "a315d46795cd40b42c90c0e0d5961cf4a607d9e9e191141322f04de1f1d6cf81" + }, + { + "id": 171, + "sha256": "11f22f5c2c669588192b70b40ea873af70370adcc7831d123d7c5643fb230bd3" + }, + { + "id": 172, + "sha256": "0f10a5592ccaea8a74647a58d4bb57048310794b620537b70120ecd0e6a3a621" + }, + { + "id": 173, + "sha256": "bb4b18c25600e1828988e9033ce7cb4bc4c936e8f81e684148065e1eba9259e9" + }, + { + "id": 174, + "sha256": "36cdd23abd4828c59a8fe622d8a9116f2f503f97c1aa51003872c97f4a404050" + }, + { + "id": 175, + "sha256": "5256233d69b95857bd2fd222dd40cf0236f9629d3c68adb502da1b6ce1c1b60d" + }, + { + "id": 176, + "sha256": "933387ddcc955352c2b540b44f96de93d7acfc1888fe679d302805f5990a5a00" + }, + { + "id": 177, + "sha256": "fd88b72d63607524de239bbe4b23a3940772cb9ec31402d640b5b885d22eb862" + }, + { + "id": 178, + "sha256": "c657cd9561418effa06d5c79b9c2ee0bc8529c9af6453ce2b20657e16ba43838" + }, + { + "id": 179, + "sha256": "15c0c9031544260f50f3aa124e516cfe34e6e010ba5cb2f9f0cf1d3dcc942f10" + }, + { + "id": 180, + "sha256": "96e069739b48fc113a02e68cce90cd342dd83937931dc4e0d3e2a4c16d1215fd" + }, + { + "id": 182, + "sha256": "56193432582553ad37f4f3e7028ab23f015dc9cfcd571b6fd516c90e7917f5ed" + }, + { + "id": 183, + "sha256": "416ea12f9f5aadc5b75556c8b6004b41faffd03f8ec800a98b1fa8ca67101353" + }, + { + "id": 184, + "sha256": "3ed67563f5aec76d4e4efcdcd80e274d8e5c1be721569dcd0a04c045e2599d76" + }, + { + "id": 185, + "sha256": "5a974272deed1eea2b6252abb883cb477d195de9af239b48c6bc185e33aa03f6" + }, + { + "id": 186, + "sha256": "6a01e81c5af56aa233392b6beb9fbcffa6e91cf556c187e5675a9f5bb0c60b06" + }, + { + "id": 187, + "sha256": "c2f1308dcb82f9e08aa2a1dbe3de005a650186da3d96e28313818eb97da16f60" + }, + { + "id": 188, + "sha256": "f3f4a42567ee495b379aaa0af03485e761bd65cc90474b0682e69cc51e5d5462" + }, + { + "id": 189, + "sha256": "086d54f8476f58657a102613959e9be23ac9bb28c0f18dd287ad4d0516cdc7c6" + }, + { + "id": 190, + "sha256": "17f44dcf39c96e8df13164c4d8855f31fc7689e252d6f2e981f99cec44c114dd" + }, + { + "id": 191, + "sha256": "bff129b0628d950f6f42b1e097de86f8d365b668bb0a48001c5b24fec413d440" + }, + { + "id": 192, + "sha256": "02af58a396d48f2ea6277d27abc878147396e0b32cbef677ad84db622ba35957" + }, + { + "id": 193, + "sha256": "97c3c18c378cf10fd8b5b3554b629b7b9d563a9aac806ebdd12f9dce5bb77e4f" + }, + { + "id": 194, + "sha256": "97ff2e998af995664a677385179bc83dd11844b8db31078f71e8d05cd8d5b97c" + }, + { + "id": 195, + "sha256": "1418835385838c077bd53de8af6c17645fe841e784827438873a6ca5a4226e84" + }, + { + "id": 196, + "sha256": "a6b4df5d085de18b379708d30e7aecc19ecd900e4d9187e116c2f36a1899c8ea" + }, + { + "id": 197, + "sha256": "94c7cb503d147d70631eb874772cd07d0a50530fa7cdbfcebfe03c96067499c9" + }, + { + "id": 198, + "sha256": "0f125052a379f5e3f1d0e135a6c594c5d2f51702d5f8ebfbcb63c34e6e5ad488" + }, + { + "id": 199, + "sha256": "4062735dd92ee8d076e9023c3e1a23c76255b2ba0d58fedf170570cee8423de9" + }, + { + "id": 200, + "sha256": "4e1ee45883f20ca771416f545be4a06524ba011c13f7c0a228b3b952c7f036c6" + }, + { + "id": 201, + "sha256": "0a16da35ab9d4b6800b13e608197a7b68bccd571545f1c6c606c1a35c91118d0" + } + ] +} \ No newline at end of file diff --git a/mlx_convert/eval/make_report.py b/mlx_convert/eval/make_report.py new file mode 100644 index 0000000..faf6c2d --- /dev/null +++ b/mlx_convert/eval/make_report.py @@ -0,0 +1,113 @@ +#!/usr/bin/env python3 +"""Render results/RESULTS.md from summary.json, speed.json and per-clip jsonl.""" +import json +import re +from pathlib import Path + +from common import HERE, corpus_error_rate, normalize + +MODELS = ["v3_ctc", "v3_rnnt", "v3_e2e_ctc", "v3_e2e_rnnt"] +FORMATTED = re.compile(r"[0-9a-zA-Z]") + + +def pct(x): + return f"{100 * x:.2f}%" + + +def main(): + res = HERE / "results" + summary = json.loads((res / "summary.json").read_text()) + speed = json.loads((res / "speed.json").read_text()) if (res / "speed.json").exists() else None + rows = {m: {r["id"]: r for r in map(json.loads, open(res / f"{m}.jsonl"))} for m in MODELS} + + # Common subset: clips where neither E2E transcript has digits or Latin letters. + ids = sorted(rows["v3_ctc"]) + plain = [i for i in ids if not any(FORMATTED.search(rows[m][i]["pytorch"]) for m in ("v3_e2e_ctc", "v3_e2e_rnnt"))] + + m = summary["machine"] + lines = [ + "# GigaAM MLX eval results", + "", + f"Machine: {m['chip']}, {m['memory_gb']} GB, macOS {m['macos']}; " + f"Python {m['python']}, mlx {m['mlx']}, torch {m['torch']}, torchaudio {m['torchaudio']}.", + "", + f"Corpus: Golos Crowd test, {summary['corpus']['clips']} clips, " + f"{summary['corpus']['seconds']:.2f} s ({summary['v3_ctc']['pytorch']['wer']['ref_units']} reference words); " + "the sample of github.com/garrrikkotua/gigaam-vs-whisper-russian, verified by SHA-256.", + "", + "## Parity: MLX vs PyTorch", + "", + "Reference: upstream PyTorch GigaAM on CPU in fp32. MLX: the fp16 file loaded with the `load_model`", + "default (weights upcast to fp32 at load). The upstream encoder weights are stored in fp16, so fp16", + "conversion loses nothing there; only the small CTC head weights are fp32 in the checkpoint.", + "", + "| Model | Identical transcripts | Encoder max \\|Δ\\| (median / max over clips) | CTC frame labels identical |", + "|---|---:|---:|---:|", + ] + for name in MODELS: + s = summary[name]["mlx_fp16"] + frames = pct(s["ctc_frame_agreement"]) if "ctc_frame_agreement" in s else "—" + lines.append( + f"| {name} | {s['exact_match_vs_pytorch']}/{summary[name]['clips']} | " + f"{s['enc_max_abs_diff']['median']:.1e} / {s['enc_max_abs_diff']['max']:.1e} | {frames} |" + ) + lines += [ + "", + f"Median max |encoder output| is {summary['v3_ctc']['enc_abs_max_median']:.2f} (v3_ctc) for scale.", + "", + "## Accuracy (MLX fp16; identical to PyTorch, see above)", + "", + "Normalization for every hypothesis: lowercase, ё→е, numbers → words (num2words), punctuation and", + "symbols removed. Corpus WER = total word edits / total reference words.", + "", + f"The last column restricts to the {len(plain)} clips where neither E2E transcript contains digits", + "or Latin letters (E2E writes \"HD\", \"4K\", \"15-й\", phone numbers as digit groups; Golos spells them", + "out in Cyrillic, and no normalizer maps one onto the other reliably).", + "", + "| Model | WER (S/I/D) | CER | WER, plain subset |", + "|---|---:|---:|---:|", + ] + for name in MODELS: + s = summary[name]["mlx_fp16"] + w = s["wer"] + sub = corpus_error_rate([(rows[name][i]["reference"], normalize(rows[name][i]["mlx_fp16"])) for i in plain]) + lines.append( + f"| {name} | {pct(w['rate'])} ({w['substitutions']}/{w['insertions']}/{w['deletions']}) | " + f"{pct(s['cer']['rate'])} | {pct(sub['rate'])} ({sub['errors']}/{sub['ref_units']}) |" + ) + + if speed: + p = speed["protocol"] + wl = p["workloads"] + lines += [ + "", + "## Speed", + "", + f"In-memory audio, one warm-up pass, then {p['repeats']} timed passes; median pass reported.", + f"`short`: {wl['short']['items']} clips one by one ({wl['short']['audio_seconds']:.1f} s). " + f"`long`: {wl['long']['items']} segments of ≈20 s made by concatenating consecutive clips " + f"({wl['long']['audio_seconds']:.1f} s). RTFx = audio seconds / wall seconds. " + f"PyTorch CPU uses {p['torch_threads']} threads; PyTorch MPS runs the upstream fp16-autocast encoder. " + "`mlx_fp16`: fp16 file, weights upcast to fp32 at load (default); `mlx_fp16_stored`: fp16 params kept " + "(`load_model(..., dtype=None)`); `mlx_fp32`: fp32 file. MLX buffer cache capped at 1 GB; one subprocess " + "per row. Peak memory: MLX peak allocation / MPS driver allocation / CPU max RSS — not directly comparable.", + "", + "| Model | Backend | Load, s | short RTFx | short p50, ms | long RTFx | long p50, ms | Peak mem, GB |", + "|---|---|---:|---:|---:|---:|---:|---:|", + ] + order = {b: i for i, b in enumerate(["mlx_fp16", "mlx_fp16_stored", "mlx_fp32", "torch_mps", "torch_cpu"])} + runs = sorted(speed["runs"].items(), key=lambda kv: (MODELS.index(kv[0].split("/")[0]), order[kv[0].split("/")[1]])) + for key, e in runs: + name, backend = key.split("/") + lines.append( + f"| {name} | {backend} | {e['load_seconds']:.1f} | {e['short']['rtfx_median']:.1f} | " + f"{e['short']['latency_ms_p50']:.1f} | {e['long']['rtfx_median']:.1f} | {e['long']['latency_ms_p50']:.1f} | " + f"{e.get('peak_memory_gb', float('nan')):.2f} |" + ) + + (res / "RESULTS.md").write_text("\n".join(lines) + "\n") + print((res / "RESULTS.md").read_text()) + + +if __name__ == "__main__": + main() diff --git a/mlx_convert/eval/prepare_golos.py b/mlx_convert/eval/prepare_golos.py new file mode 100644 index 0000000..7a5255f --- /dev/null +++ b/mlx_convert/eval/prepare_golos.py @@ -0,0 +1,67 @@ +#!/usr/bin/env python3 +"""Download the Golos Crowd eval sample and verify it byte-for-byte. + +The sample is the 200-utterance Golos Crowd test subset published by +https://github.com/garrrikkotua/gigaam-vs-whisper-russian. Its manifest stores +a SHA-256 per clip; `golos_crowd_200.json` keeps those ids and hashes. The +audio itself is not redistributed here (Golos licence), so this script pulls +the Hugging Face mirror, takes the listed rows and checks every hash. + +Output: /wav/.wav and /manifest.jsonl with +{"id", "audio", "duration", "reference"} per line. +""" +import argparse +import hashlib +import io +import json +from pathlib import Path + +import pyarrow.parquet as pq +import soundfile as sf +from huggingface_hub import hf_hub_download + +HERE = Path(__file__).parent + + +def main(): + parser = argparse.ArgumentParser(description=__doc__.splitlines()[0]) + parser.add_argument("--out", default=str(HERE / "data" / "golos_crowd_200")) + args = parser.parse_args() + + spec = json.loads((HERE / "golos_crowd_200.json").read_text()) + parquet = hf_hub_download( + spec["dataset"], spec["file"], repo_type="dataset", revision=spec["revision"] + ) + table = pq.read_table(parquet) + audio, text = table.column("audio"), table.column("transcription") + + out = Path(args.out) + (out / "wav").mkdir(parents=True, exist_ok=True) + with open(out / "manifest.jsonl", "w") as manifest: + for item in spec["items"]: + row = item["id"] + data = audio[row].as_py()["bytes"] + digest = hashlib.sha256(data).hexdigest() + if digest != item["sha256"]: + raise SystemExit(f"row {row}: sha256 {digest} != {item['sha256']}") + path = out / "wav" / f"{row:06d}.wav" + path.write_bytes(data) + info = sf.info(io.BytesIO(data)) + if info.samplerate != 16000 or info.channels != 1: + raise SystemExit(f"row {row}: {info.samplerate} Hz, {info.channels} ch") + record = { + "id": row, + "audio": str(path.relative_to(out)), + "duration": info.frames / info.samplerate, + "reference": text[row].as_py(), + } + manifest.write(json.dumps(record, ensure_ascii=False) + "\n") + + total = sum( + json.loads(line)["duration"] for line in open(out / "manifest.jsonl") + ) + print(f"{len(spec['items'])} clips verified, {total:.2f} s of audio → {out}") + + +if __name__ == "__main__": + main() diff --git a/mlx_convert/eval/results/RESULTS.md b/mlx_convert/eval/results/RESULTS.md new file mode 100644 index 0000000..012e05b --- /dev/null +++ b/mlx_convert/eval/results/RESULTS.md @@ -0,0 +1,64 @@ +# GigaAM MLX eval results + +Machine: Apple M4 Pro, 24 GB, macOS 26.6.2; Python 3.12.13, mlx 0.32.2, torch 2.14.0, torchaudio 2.11.0. + +Corpus: Golos Crowd test, 200 clips, 823.85 s (1001 reference words); the sample of github.com/garrrikkotua/gigaam-vs-whisper-russian, verified by SHA-256. + +## Parity: MLX vs PyTorch + +Reference: upstream PyTorch GigaAM on CPU in fp32. MLX: the fp16 file loaded with the `load_model` +default (weights upcast to fp32 at load). The upstream encoder weights are stored in fp16, so fp16 +conversion loses nothing there; only the small CTC head weights are fp32 in the checkpoint. + +| Model | Identical transcripts | Encoder max \|Δ\| (median / max over clips) | CTC frame labels identical | +|---|---:|---:|---:| +| v3_ctc | 200/200 | 3.0e-06 / 8.1e-05 | 100.00% | +| v3_rnnt | 200/200 | 3.2e-06 / 1.7e-05 | — | +| v3_e2e_ctc | 200/200 | 3.9e-06 / 8.7e-05 | 100.00% | +| v3_e2e_rnnt | 200/200 | 3.3e-06 / 1.0e-05 | — | + +Median max |encoder output| is 1.54 (v3_ctc) for scale. + +## Accuracy (MLX fp16; identical to PyTorch, see above) + +Normalization for every hypothesis: lowercase, ё→е, numbers → words (num2words), punctuation and +symbols removed. Corpus WER = total word edits / total reference words. + +The last column restricts to the 168 clips where neither E2E transcript contains digits +or Latin letters (E2E writes "HD", "4K", "15-й", phone numbers as digit groups; Golos spells them +out in Cyrillic, and no normalizer maps one onto the other reliably). + +| Model | WER (S/I/D) | CER | WER, plain subset | +|---|---:|---:|---:| +| v3_ctc | 1.70% (15/1/1) | 0.30% | 1.35% (11/815) | +| v3_rnnt | 1.50% (12/2/1) | 0.24% | 1.23% (10/815) | +| v3_e2e_ctc | 6.49% (46/13/6) | 3.14% | 2.70% (22/815) | +| v3_e2e_rnnt | 5.69% (42/10/5) | 3.19% | 2.09% (17/815) | + +## Speed + +In-memory audio, one warm-up pass, then 3 timed passes; median pass reported. +`short`: 200 clips one by one (823.8 s). `long`: 38 segments of ≈20 s made by concatenating consecutive clips (815.0 s). RTFx = audio seconds / wall seconds. PyTorch CPU uses 8 threads; PyTorch MPS runs the upstream fp16-autocast encoder. `mlx_fp16`: fp16 file, weights upcast to fp32 at load (default); `mlx_fp16_stored`: fp16 params kept (`load_model(..., dtype=None)`); `mlx_fp32`: fp32 file. MLX buffer cache capped at 1 GB; one subprocess per row. Peak memory: MLX peak allocation / MPS driver allocation / CPU max RSS — not directly comparable. + +| Model | Backend | Load, s | short RTFx | short p50, ms | long RTFx | long p50, ms | Peak mem, GB | +|---|---|---:|---:|---:|---:|---:|---:| +| v3_ctc | mlx_fp16 | 0.1 | 201.4 | 20.0 | 285.2 | 74.4 | 1.81 | +| v3_ctc | mlx_fp16_stored | 0.0 | 167.4 | 24.5 | 270.0 | 79.6 | 1.35 | +| v3_ctc | mlx_fp32 | 0.2 | 199.5 | 19.9 | 286.8 | 76.0 | 1.81 | +| v3_ctc | torch_mps | 1.0 | 200.6 | 20.2 | 268.0 | 80.9 | 1.63 | +| v3_ctc | torch_cpu | 1.0 | 10.0 | 402.6 | 41.2 | 507.2 | 2.01 | +| v3_rnnt | mlx_fp16 | 0.1 | 90.9 | 43.4 | 102.4 | 208.9 | 1.82 | +| v3_rnnt | mlx_fp16_stored | 0.0 | 72.8 | 54.6 | 87.8 | 249.2 | 1.35 | +| v3_rnnt | mlx_fp32 | 0.2 | 86.6 | 45.6 | 97.0 | 222.1 | 1.82 | +| v3_rnnt | torch_mps | 1.1 | 37.1 | 104.1 | 40.5 | 515.6 | 0.66 | +| v3_rnnt | torch_cpu | 1.0 | 9.8 | 412.7 | 39.0 | 542.1 | 2.03 | +| v3_e2e_ctc | mlx_fp16 | 0.1 | 204.5 | 19.7 | 288.8 | 74.9 | 1.81 | +| v3_e2e_ctc | mlx_fp16_stored | 0.0 | 168.4 | 24.1 | 266.1 | 80.8 | 1.35 | +| v3_e2e_ctc | mlx_fp32 | 0.2 | 198.7 | 20.7 | 279.3 | 77.2 | 1.81 | +| v3_e2e_ctc | torch_mps | 1.1 | 202.1 | 20.1 | 271.1 | 79.2 | 0.63 | +| v3_e2e_ctc | torch_cpu | 1.0 | 10.0 | 403.4 | 41.4 | 503.6 | 2.02 | +| v3_e2e_rnnt | mlx_fp16 | 0.1 | 98.8 | 41.4 | 112.4 | 191.2 | 1.82 | +| v3_e2e_rnnt | mlx_fp16_stored | 0.0 | 79.9 | 50.5 | 94.1 | 224.6 | 1.35 | +| v3_e2e_rnnt | mlx_fp32 | 0.2 | 92.2 | 43.1 | 104.3 | 206.2 | 1.82 | +| v3_e2e_rnnt | torch_mps | 1.1 | 44.5 | 87.7 | 48.2 | 445.9 | 0.66 | +| v3_e2e_rnnt | torch_cpu | 1.0 | 10.0 | 409.8 | 41.4 | 516.9 | 2.04 | diff --git a/mlx_convert/eval/results/speed.json b/mlx_convert/eval/results/speed.json new file mode 100644 index 0000000..6f59f35 --- /dev/null +++ b/mlx_convert/eval/results/speed.json @@ -0,0 +1,550 @@ +{ + "runs": { + "v3_ctc/mlx_fp16": { + "load_seconds": 0.10232454098877497, + "short": { + "wall_seconds": [ + 4.086858624999877, + 4.1288881250075065, + 4.0907160000060685 + ], + "median_wall_seconds": 4.0907160000060685, + "rtfx_median": 201.39431275570774, + "latency_ms_p50": 19.96616701944731, + "latency_ms_p90": 26.182749978033826 + }, + "long": { + "wall_seconds": [ + 2.8686550830025226, + 2.8571916249929927, + 2.82385716601857 + ], + "median_wall_seconds": 2.8571916249929927, + "rtfx_median": 285.24160940797896, + "latency_ms_p50": 74.41724999807775, + "latency_ms_p90": 83.50595901720226 + }, + "peak_memory_gb": 1.812965091317892 + }, + "v3_ctc/mlx_fp32": { + "load_seconds": 0.22649966698372737, + "short": { + "wall_seconds": [ + 4.128988791984739, + 4.151853167015361, + 3.93045429198537 + ], + "median_wall_seconds": 4.128988791984739, + "rtfx_median": 199.52753059036277, + "latency_ms_p50": 19.9359169928357, + "latency_ms_p90": 26.26841701567173 + }, + "long": { + "wall_seconds": [ + 2.84438450000016, + 2.8418748330150265, + 2.8209071249875706 + ], + "median_wall_seconds": 2.8418748330150265, + "rtfx_median": 286.77897000669583, + "latency_ms_p50": 75.97370800795034, + "latency_ms_p90": 82.05737499520183 + }, + "peak_memory_gb": 1.812965091317892 + }, + "v3_ctc/torch_cpu": { + "load_seconds": 0.9533450839808211, + "short": { + "wall_seconds": [ + 82.11337449998246, + 82.0158625410113, + 84.40155920898542 + ], + "median_wall_seconds": 82.11337449998246, + "rtfx_median": 10.03304202898367, + "latency_ms_p50": 402.62562499265186, + "latency_ms_p90": 421.77512499620207 + }, + "long": { + "wall_seconds": [ + 19.085255415993743, + 19.75800987498951, + 19.84372141599306 + ], + "median_wall_seconds": 19.75800987498951, + "rtfx_median": 41.248584379525354, + "latency_ms_p50": 507.21633300418034, + "latency_ms_p90": 526.5212499944028 + }, + "peak_memory_gb": 2.0136260986328125 + }, + "v3_ctc/torch_mps": { + "load_seconds": 1.0356653339986224, + "short": { + "wall_seconds": [ + 4.116375624987995, + 4.086937082989607, + 4.106905458000256 + ], + "median_wall_seconds": 4.106905458000256, + "rtfx_median": 200.60041457617325, + "latency_ms_p50": 20.247250009560958, + "latency_ms_p90": 25.665416003903374 + }, + "long": { + "wall_seconds": [ + 3.0429849999782164, + 3.0348832920135465, + 3.0414441670000087 + ], + "median_wall_seconds": 3.0414441670000087, + "rtfx_median": 267.9614987980799, + "latency_ms_p50": 80.88933399994858, + "latency_ms_p90": 90.72420798474923 + }, + "peak_memory_gb": 1.6261444091796875 + }, + "v3_rnnt/mlx_fp16": { + "load_seconds": 0.0937266249966342, + "short": { + "wall_seconds": [ + 8.820479834015714, + 9.16008470801171, + 9.061758958006976 + ], + "median_wall_seconds": 9.061758958006976, + "rtfx_median": 90.9146823831645, + "latency_ms_p50": 43.41783400741406, + "latency_ms_p90": 60.08016600389965 + }, + "long": { + "wall_seconds": [ + 8.035170332994312, + 7.961204249993898, + 7.95552766599576 + ], + "median_wall_seconds": 7.961204249993898, + "rtfx_median": 102.37018319189897, + "latency_ms_p50": 208.91670801211149, + "latency_ms_p90": 234.76270900573581 + }, + "peak_memory_gb": 1.8173071704804897 + }, + "v3_rnnt/mlx_fp32": { + "load_seconds": 0.22607033399981447, + "short": { + "wall_seconds": [ + 9.518382083013421, + 9.381419958983315, + 9.756669124995824 + ], + "median_wall_seconds": 9.518382083013421, + "rtfx_median": 86.55325351671306, + "latency_ms_p50": 45.588624983793125, + "latency_ms_p90": 63.879249995807186 + }, + "long": { + "wall_seconds": [ + 8.401950792002026, + 8.483831791003468, + 8.361744041991187 + ], + "median_wall_seconds": 8.401950792002026, + "rtfx_median": 97.00008458461863, + "latency_ms_p50": 222.0606249757111, + "latency_ms_p90": 238.56308398535475 + }, + "peak_memory_gb": 1.8173057399690151 + }, + "v3_rnnt/torch_cpu": { + "load_seconds": 0.9802879159979057, + "short": { + "wall_seconds": [ + 83.93764229200315, + 84.36142062500585, + 81.68046299999696 + ], + "median_wall_seconds": 83.93764229200315, + "rtfx_median": 9.81498782910762, + "latency_ms_p50": 412.71370797767304, + "latency_ms_p90": 437.1128329948988 + }, + "long": { + "wall_seconds": [ + 20.890096249990165, + 20.960050041991053, + 20.87172887500492 + ], + "median_wall_seconds": 20.890096249990165, + "rtfx_median": 39.013220798366774, + "latency_ms_p50": 542.1365420043003, + "latency_ms_p90": 587.5122500001453 + }, + "peak_memory_gb": 2.03131103515625 + }, + "v3_rnnt/torch_mps": { + "load_seconds": 1.1195401250151917, + "short": { + "wall_seconds": [ + 22.19301445799647, + 22.24447766697267, + 22.19432979199337 + ], + "median_wall_seconds": 22.19432979199337, + "rtfx_median": 37.11970333058689, + "latency_ms_p50": 104.08745802124031, + "latency_ms_p90": 148.4908330021426 + }, + "long": { + "wall_seconds": [ + 20.284873666998465, + 20.13470245798817, + 19.637342165980954 + ], + "median_wall_seconds": 20.13470245798817, + "rtfx_median": 40.4768801128553, + "latency_ms_p50": 515.5519999971148, + "latency_ms_p90": 585.4536249826197 + }, + "peak_memory_gb": 0.6553497314453125 + }, + "v3_e2e_ctc/mlx_fp16": { + "load_seconds": 0.09245475000352599, + "short": { + "wall_seconds": [ + 4.153190999990329, + 4.015187292010523, + 4.0294104579952545 + ], + "median_wall_seconds": 4.0294104579952545, + "rtfx_median": 204.458430355563, + "latency_ms_p50": 19.721083983313292, + "latency_ms_p90": 25.923959008650854 + }, + "long": { + "wall_seconds": [ + 2.9061026249837596, + 2.8086023340001702, + 2.821747999987565 + ], + "median_wall_seconds": 2.821747999987565, + "rtfx_median": 288.8244937193511, + "latency_ms_p50": 74.92437498876825, + "latency_ms_p90": 82.7202080108691 + }, + "peak_memory_gb": 1.8136096522212029 + }, + "v3_e2e_ctc/mlx_fp32": { + "load_seconds": 0.2305249159981031, + "short": { + "wall_seconds": [ + 4.145894540997688, + 4.1303812919941265, + 4.1583265000081155 + ], + "median_wall_seconds": 4.145894540997688, + "rtfx_median": 198.71391550199576, + "latency_ms_p50": 20.7422909734305, + "latency_ms_p90": 26.029250002466142 + }, + "long": { + "wall_seconds": [ + 2.935972415987635, + 2.918340333009837, + 2.916278290998889 + ], + "median_wall_seconds": 2.918340333009837, + "rtfx_median": 279.26487129739877, + "latency_ms_p50": 77.22762500634417, + "latency_ms_p90": 85.58816701406613 + }, + "peak_memory_gb": 1.8136067911982536 + }, + "v3_e2e_ctc/torch_cpu": { + "load_seconds": 1.0108446250087582, + "short": { + "wall_seconds": [ + 83.21497370800353, + 82.3341674169933, + 82.51597404197673 + ], + "median_wall_seconds": 82.51597404197673, + "rtfx_median": 9.984090317844403, + "latency_ms_p50": 403.36716698948294, + "latency_ms_p90": 428.0246249982156 + }, + "long": { + "wall_seconds": [ + 19.24095095798839, + 19.692850207997253, + 19.698533666000003 + ], + "median_wall_seconds": 19.692850207997253, + "rtfx_median": 41.3850676205841, + "latency_ms_p50": 503.59850001404993, + "latency_ms_p90": 523.0746250017546 + }, + "peak_memory_gb": 2.0199127197265625 + }, + "v3_e2e_ctc/torch_mps": { + "load_seconds": 1.106898166006431, + "short": { + "wall_seconds": [ + 4.074119083990809, + 4.097471874993062, + 4.075789124995936 + ], + "median_wall_seconds": 4.075789124995936, + "rtfx_median": 202.13188470608657, + "latency_ms_p50": 20.133209007326514, + "latency_ms_p90": 24.82254197821021 + }, + "long": { + "wall_seconds": [ + 3.5833492499950808, + 3.0046262920077424, + 3.0066322079801466 + ], + "median_wall_seconds": 3.0066322079801466, + "rtfx_median": 271.0640614229, + "latency_ms_p50": 79.21700002043508, + "latency_ms_p90": 89.37324999715202 + }, + "peak_memory_gb": 0.6260223388671875 + }, + "v3_e2e_rnnt/mlx_fp16": { + "load_seconds": 0.11242091699386947, + "short": { + "wall_seconds": [ + 8.301743292016909, + 8.337002540996764, + 8.716015999991214 + ], + "median_wall_seconds": 8.337002540996764, + "rtfx_median": 98.81812239456288, + "latency_ms_p50": 41.36545900837518, + "latency_ms_p90": 55.05166601506062 + }, + "long": { + "wall_seconds": [ + 7.787013332999777, + 7.248533332982333, + 7.233173917018576 + ], + "median_wall_seconds": 7.248533332982333, + "rtfx_median": 112.43515067960388, + "latency_ms_p50": 191.23850000323728, + "latency_ms_p90": 208.62295798724517 + }, + "peak_memory_gb": 1.8196912333369255 + }, + "v3_e2e_rnnt/mlx_fp32": { + "load_seconds": 0.23576770900399424, + "short": { + "wall_seconds": [ + 8.938225041987607, + 8.983470959006809, + 8.709973041026387 + ], + "median_wall_seconds": 8.938225041987607, + "rtfx_median": 92.17120106396425, + "latency_ms_p50": 43.10587499639951, + "latency_ms_p90": 58.26729198452085 + }, + "long": { + "wall_seconds": [ + 7.698811874986859, + 7.828786166995997, + 7.815913667000132 + ], + "median_wall_seconds": 7.815913667000132, + "rtfx_median": 104.27314991221054, + "latency_ms_p50": 206.24812500318512, + "latency_ms_p90": 224.4062919926364 + }, + "peak_memory_gb": 1.819689802825451 + }, + "v3_e2e_rnnt/torch_cpu": { + "load_seconds": 0.9726846250123344, + "short": { + "wall_seconds": [ + 82.60581683297642, + 82.63600333299837, + 82.29966808401514 + ], + "median_wall_seconds": 82.60581683297642, + "rtfx_median": 9.973231536052296, + "latency_ms_p50": 409.7996670170687, + "latency_ms_p90": 434.0387500124052 + }, + "long": { + "wall_seconds": [ + 19.61619995799265, + 19.704344458994456, + 19.668803999986267 + ], + "median_wall_seconds": 19.668803999986267, + "rtfx_median": 41.43566316999087, + "latency_ms_p50": 516.8940409785137, + "latency_ms_p90": 533.971582975937 + }, + "peak_memory_gb": 2.0365753173828125 + }, + "v3_e2e_rnnt/torch_mps": { + "load_seconds": 1.0527325840084814, + "short": { + "wall_seconds": [ + 18.126744542008964, + 18.50924837499042, + 18.89126333297463 + ], + "median_wall_seconds": 18.50924837499042, + "rtfx_median": 44.510015793681646, + "latency_ms_p50": 87.6647499972023, + "latency_ms_p90": 124.20445898897015 + }, + "long": { + "wall_seconds": [ + 16.704064791003475, + 16.91423220897559, + 16.981200874986826 + ], + "median_wall_seconds": 16.91423220897559, + "rtfx_median": 48.183679130733644, + "latency_ms_p50": 445.9148339810781, + "latency_ms_p90": 504.9896249838639 + }, + "peak_memory_gb": 0.6553497314453125 + }, + "v3_ctc/mlx_fp16_stored": { + "load_seconds": 0.020814209012314677, + "short": { + "wall_seconds": [ + 4.898272165999515, + 4.933290082990425, + 4.922680165997008 + ], + "median_wall_seconds": 4.922680165997008, + "rtfx_median": 167.35739672681808, + "latency_ms_p50": 24.491249991115183, + "latency_ms_p90": 30.245792004279792 + }, + "long": { + "wall_seconds": [ + 3.0179245000181254, + 3.037745707988506, + 3.0157417919836007 + ], + "median_wall_seconds": 3.0179245000181254, + "rtfx_median": 270.049809892562, + "latency_ms_p50": 79.56420900882222, + "latency_ms_p90": 87.22962500178255 + }, + "peak_memory_gb": 1.3466144055128098 + }, + "v3_rnnt/mlx_fp16_stored": { + "load_seconds": 0.030107124999631196, + "short": { + "wall_seconds": [ + 11.318400541000301, + 11.271531167003559, + 11.334681915992405 + ], + "median_wall_seconds": 11.318400541000301, + "rtfx_median": 72.78828263018777, + "latency_ms_p50": 54.63833297835663, + "latency_ms_p90": 73.1724170036614 + }, + "long": { + "wall_seconds": [ + 9.22983816600754, + 9.298193167021964, + 9.28689704099088 + ], + "median_wall_seconds": 9.28689704099088, + "rtfx_median": 87.75696919032963, + "latency_ms_p50": 249.17695901240222, + "latency_ms_p90": 272.87974997307174 + }, + "peak_memory_gb": 1.3487971276044846 + }, + "v3_e2e_ctc/mlx_fp16_stored": { + "load_seconds": 0.02963579201605171, + "short": { + "wall_seconds": [ + 4.893074207997415, + 4.920103124983143, + 4.891598458983935 + ], + "median_wall_seconds": 4.893074207997415, + "rtfx_median": 168.37000676455614, + "latency_ms_p50": 24.09483300289139, + "latency_ms_p90": 30.001624982105568 + }, + "long": { + "wall_seconds": [ + 3.0630069590115454, + 3.0731088329921477, + 3.060329416010063 + ], + "median_wall_seconds": 3.0630069590115454, + "rtfx_median": 266.07511781919135, + "latency_ms_p50": 80.78724998631515, + "latency_ms_p90": 88.67341699078679 + }, + "peak_memory_gb": 1.3469352554529905 + }, + "v3_e2e_rnnt/mlx_fp16_stored": { + "load_seconds": 0.030824666988337412, + "short": { + "wall_seconds": [ + 10.474925499991514, + 10.280821041989839, + 10.305550166987814 + ], + "median_wall_seconds": 10.305550166987814, + "rtfx_median": 79.9420626895847, + "latency_ms_p50": 50.53154198685661, + "latency_ms_p90": 66.24508398817852 + }, + "long": { + "wall_seconds": [ + 8.663679290999426, + 8.624599916976877, + 8.714920916012488 + ], + "median_wall_seconds": 8.663679290999426, + "rtfx_median": 94.06972605122648, + "latency_ms_p50": 224.59279198665172, + "latency_ms_p90": 247.93300000601448 + }, + "peak_memory_gb": 1.349992023780942 + } + }, + "machine": { + "chip": "Apple M4 Pro", + "memory_gb": "24", + "macos": "26.6.2", + "python": "3.12.13", + "mlx": "0.32.2", + "torch": "2.14.0", + "torchaudio": "2.11.0" + }, + "protocol": { + "repeats": 3, + "warmup": "one full pass per workload", + "torch_threads": 8, + "mlx_cache_limit_gb": 1, + "isolation": "one subprocess per model/backend", + "workloads": { + "short": { + "items": 200, + "audio_seconds": 823.8469375 + }, + "long": { + "items": 38, + "audio_seconds": 814.9899375 + } + } + } +} \ No newline at end of file diff --git a/mlx_convert/eval/results/summary.json b/mlx_convert/eval/results/summary.json new file mode 100644 index 0000000..5f31e3c --- /dev/null +++ b/mlx_convert/eval/results/summary.json @@ -0,0 +1,391 @@ +{ + "machine": { + "chip": "Apple M4 Pro", + "memory_gb": "24", + "macos": "26.6.2", + "python": "3.12.13", + "mlx": "0.32.2", + "torch": "2.14.0", + "torchaudio": "2.11.0" + }, + "corpus": { + "clips": 200, + "seconds": 823.8469375 + }, + "v3_ctc": { + "pytorch": { + "wer": { + "errors": 17, + "substitutions": 15, + "insertions": 1, + "deletions": 1, + "ref_units": 1001, + "rate": 0.016983016983016984 + }, + "cer": { + "errors": 18, + "substitutions": 6, + "insertions": 2, + "deletions": 10, + "ref_units": 5924, + "rate": 0.003038487508440243 + }, + "wer_plain_subset": { + "errors": 17, + "substitutions": 15, + "insertions": 1, + "deletions": 1, + "ref_units": 1001, + "rate": 0.016983016983016984 + } + }, + "mlx_fp32": { + "wer": { + "errors": 17, + "substitutions": 15, + "insertions": 1, + "deletions": 1, + "ref_units": 1001, + "rate": 0.016983016983016984 + }, + "cer": { + "errors": 18, + "substitutions": 6, + "insertions": 2, + "deletions": 10, + "ref_units": 5924, + "rate": 0.003038487508440243 + }, + "wer_plain_subset": { + "errors": 17, + "substitutions": 15, + "insertions": 1, + "deletions": 1, + "ref_units": 1001, + "rate": 0.016983016983016984 + }, + "exact_match_vs_pytorch": 200, + "enc_max_abs_diff": { + "median": 2.9653310775756836e-06, + "max": 8.106231689453125e-05 + }, + "ctc_frame_agreement": 1.0 + }, + "mlx_fp16": { + "wer": { + "errors": 17, + "substitutions": 15, + "insertions": 1, + "deletions": 1, + "ref_units": 1001, + "rate": 0.016983016983016984 + }, + "cer": { + "errors": 18, + "substitutions": 6, + "insertions": 2, + "deletions": 10, + "ref_units": 5924, + "rate": 0.003038487508440243 + }, + "wer_plain_subset": { + "errors": 17, + "substitutions": 15, + "insertions": 1, + "deletions": 1, + "ref_units": 1001, + "rate": 0.016983016983016984 + }, + "exact_match_vs_pytorch": 200, + "enc_max_abs_diff": { + "median": 2.9653310775756836e-06, + "max": 8.106231689453125e-05 + }, + "ctc_frame_agreement": 1.0 + }, + "enc_abs_max_median": 1.5383418798446655, + "clips": 200, + "plain_subset_clips": 200 + }, + "v3_rnnt": { + "pytorch": { + "wer": { + "errors": 15, + "substitutions": 12, + "insertions": 2, + "deletions": 1, + "ref_units": 1001, + "rate": 0.014985014985014986 + }, + "cer": { + "errors": 14, + "substitutions": 3, + "insertions": 3, + "deletions": 8, + "ref_units": 5924, + "rate": 0.002363268062120189 + }, + "wer_plain_subset": { + "errors": 15, + "substitutions": 12, + "insertions": 2, + "deletions": 1, + "ref_units": 1001, + "rate": 0.014985014985014986 + } + }, + "mlx_fp32": { + "wer": { + "errors": 15, + "substitutions": 12, + "insertions": 2, + "deletions": 1, + "ref_units": 1001, + "rate": 0.014985014985014986 + }, + "cer": { + "errors": 14, + "substitutions": 3, + "insertions": 3, + "deletions": 8, + "ref_units": 5924, + "rate": 0.002363268062120189 + }, + "wer_plain_subset": { + "errors": 15, + "substitutions": 12, + "insertions": 2, + "deletions": 1, + "ref_units": 1001, + "rate": 0.014985014985014986 + }, + "exact_match_vs_pytorch": 200, + "enc_max_abs_diff": { + "median": 3.2032839953899384e-06, + "max": 1.7091631889343262e-05 + } + }, + "mlx_fp16": { + "wer": { + "errors": 15, + "substitutions": 12, + "insertions": 2, + "deletions": 1, + "ref_units": 1001, + "rate": 0.014985014985014986 + }, + "cer": { + "errors": 14, + "substitutions": 3, + "insertions": 3, + "deletions": 8, + "ref_units": 5924, + "rate": 0.002363268062120189 + }, + "wer_plain_subset": { + "errors": 15, + "substitutions": 12, + "insertions": 2, + "deletions": 1, + "ref_units": 1001, + "rate": 0.014985014985014986 + }, + "exact_match_vs_pytorch": 200, + "enc_max_abs_diff": { + "median": 3.2032839953899384e-06, + "max": 1.7091631889343262e-05 + } + }, + "enc_abs_max_median": 2.671436071395874, + "clips": 200, + "plain_subset_clips": 200 + }, + "v3_e2e_ctc": { + "pytorch": { + "wer": { + "errors": 65, + "substitutions": 46, + "insertions": 13, + "deletions": 6, + "ref_units": 1001, + "rate": 0.06493506493506493 + }, + "cer": { + "errors": 186, + "substitutions": 78, + "insertions": 75, + "deletions": 33, + "ref_units": 5924, + "rate": 0.03139770425388251 + }, + "wer_plain_subset": { + "errors": 25, + "substitutions": 19, + "insertions": 3, + "deletions": 3, + "ref_units": 825, + "rate": 0.030303030303030304 + } + }, + "mlx_fp32": { + "wer": { + "errors": 65, + "substitutions": 46, + "insertions": 13, + "deletions": 6, + "ref_units": 1001, + "rate": 0.06493506493506493 + }, + "cer": { + "errors": 186, + "substitutions": 78, + "insertions": 75, + "deletions": 33, + "ref_units": 5924, + "rate": 0.03139770425388251 + }, + "wer_plain_subset": { + "errors": 25, + "substitutions": 19, + "insertions": 3, + "deletions": 3, + "ref_units": 825, + "rate": 0.030303030303030304 + }, + "exact_match_vs_pytorch": 200, + "enc_max_abs_diff": { + "median": 3.933906555175781e-06, + "max": 8.713826537132263e-05 + }, + "ctc_frame_agreement": 1.0 + }, + "mlx_fp16": { + "wer": { + "errors": 65, + "substitutions": 46, + "insertions": 13, + "deletions": 6, + "ref_units": 1001, + "rate": 0.06493506493506493 + }, + "cer": { + "errors": 186, + "substitutions": 78, + "insertions": 75, + "deletions": 33, + "ref_units": 5924, + "rate": 0.03139770425388251 + }, + "wer_plain_subset": { + "errors": 25, + "substitutions": 19, + "insertions": 3, + "deletions": 3, + "ref_units": 825, + "rate": 0.030303030303030304 + }, + "exact_match_vs_pytorch": 200, + "enc_max_abs_diff": { + "median": 3.933906555175781e-06, + "max": 8.713826537132263e-05 + }, + "ctc_frame_agreement": 1.0 + }, + "enc_abs_max_median": 2.340498208999634, + "clips": 200, + "plain_subset_clips": 171 + }, + "v3_e2e_rnnt": { + "pytorch": { + "wer": { + "errors": 57, + "substitutions": 42, + "insertions": 10, + "deletions": 5, + "ref_units": 1001, + "rate": 0.056943056943056944 + }, + "cer": { + "errors": 189, + "substitutions": 79, + "insertions": 73, + "deletions": 37, + "ref_units": 5924, + "rate": 0.031904118838622555 + }, + "wer_plain_subset": { + "errors": 18, + "substitutions": 16, + "insertions": 1, + "deletions": 1, + "ref_units": 826, + "rate": 0.021791767554479417 + } + }, + "mlx_fp32": { + "wer": { + "errors": 57, + "substitutions": 42, + "insertions": 10, + "deletions": 5, + "ref_units": 1001, + "rate": 0.056943056943056944 + }, + "cer": { + "errors": 189, + "substitutions": 79, + "insertions": 73, + "deletions": 37, + "ref_units": 5924, + "rate": 0.031904118838622555 + }, + "wer_plain_subset": { + "errors": 18, + "substitutions": 16, + "insertions": 1, + "deletions": 1, + "ref_units": 826, + "rate": 0.021791767554479417 + }, + "exact_match_vs_pytorch": 200, + "enc_max_abs_diff": { + "median": 3.310851752758026e-06, + "max": 9.998679161071777e-06 + } + }, + "mlx_fp16": { + "wer": { + "errors": 57, + "substitutions": 42, + "insertions": 10, + "deletions": 5, + "ref_units": 1001, + "rate": 0.056943056943056944 + }, + "cer": { + "errors": 189, + "substitutions": 79, + "insertions": 73, + "deletions": 37, + "ref_units": 5924, + "rate": 0.031904118838622555 + }, + "wer_plain_subset": { + "errors": 18, + "substitutions": 16, + "insertions": 1, + "deletions": 1, + "ref_units": 826, + "rate": 0.021791767554479417 + }, + "exact_match_vs_pytorch": 200, + "enc_max_abs_diff": { + "median": 3.310851752758026e-06, + "max": 9.998679161071777e-06 + } + }, + "enc_abs_max_median": 2.4387344121932983, + "clips": 200, + "plain_subset_clips": 170 + } +} \ No newline at end of file diff --git a/mlx_convert/eval/run_eval.py b/mlx_convert/eval/run_eval.py new file mode 100644 index 0000000..b64caa5 --- /dev/null +++ b/mlx_convert/eval/run_eval.py @@ -0,0 +1,142 @@ +#!/usr/bin/env python3 +"""Accuracy and PyTorch↔MLX parity on the Golos Crowd sample. + +For every model it runs the upstream PyTorch GigaAM (CPU, fp32 — the +reference) and the converted MLX model (fp32 and fp16) on each clip and records: + +* transcripts of all three, raw and normalized; +* corpus WER / CER against the Golos references; +* exact transcript agreement MLX vs PyTorch; +* max |encoder_pt − encoder_mlx| per clip (and CTC frame-label agreement). + +Models are read from / and /-fp32, as produced by +convert_gigaam_to_mlx.py. Results go to /.jsonl and /summary.json. +""" +import argparse +import json +import re +import time +from pathlib import Path + +import mlx.core as mx +import numpy as np +import torch + +import gigaam +from common import HERE, DEFAULT_DATA, corpus_error_rate, load_corpus, machine_info, normalize +from gigaam_mlx import load_model as load_mlx + +MODELS = ["v3_ctc", "v3_rnnt", "v3_e2e_ctc", "v3_e2e_rnnt"] + + +def mlx_encode(model, audio: np.ndarray): + mel, lengths = model._compute_features(mx.array(audio)) + encoded, enc_len = model(mel, lengths) + mx.eval(encoded, enc_len) + return encoded, int(enc_len[0]) + + +def evaluate(name: str, clips, models_dir: Path): + pt = gigaam.load_model(name, device="cpu", fp16_encoder=False, use_flash=False) + pt.eval() + variants = {"mlx_fp32": load_mlx(models_dir / f"{name}-fp32"), "mlx_fp16": load_mlx(models_dir / name)} + is_ctc = variants["mlx_fp32"].cfg.head_type == "ctc" + + rows = [] + for clip in clips: + row = {"id": clip.id, "duration": clip.duration, "reference": clip.reference} + with torch.no_grad(): + row["pytorch"] = pt.transcribe(str(clip.path)).text + wav = torch.from_numpy(clip.audio)[None] + enc_t, len_pt = pt.forward(wav, torch.tensor([wav.shape[-1]])) # [1, D, T] + logits_pt = pt.head(enc_t) if is_ctc else None # [1, T, C] + enc_pt = enc_t[0].T.numpy() # [T, D] + for key, model in variants.items(): + row[key] = model.transcribe(mx.array(clip.audio)) + enc, n = mlx_encode(model, clip.audio) + enc_np = np.array(enc[0].astype(mx.float32)) + assert n == int(len_pt[0]) and enc_np.shape == enc_pt.shape, (clip.id, n, enc_np.shape, enc_pt.shape) + row[f"{key}_enc_max_abs_diff"] = float(np.abs(enc_np - enc_pt).max()) + if is_ctc: + labels_mlx = np.array(mx.argmax(model.head(enc), axis=-1)[0]) + labels_pt = logits_pt[0].argmax(-1).numpy() + row[f"{key}_ctc_frames_agree"] = int((labels_mlx == labels_pt).sum()) + row["ctc_frames"] = int(labels_pt.shape[0]) + row["enc_abs_max"] = float(np.abs(enc_pt).max()) + rows.append(row) + return rows + + +_FORMATTED = re.compile(r"[0-9a-zA-Z]") + + +def summarize(rows): + """Corpus metrics. `wer_plain_subset` restricts to clips whose PyTorch + transcript has no digits or Latin letters: E2E models write "HD", "4K", + "15-й", phone numbers as digit groups, while Golos references spell them + in Cyrillic words, and no normalizer maps one onto the other reliably.""" + out = {} + plain = [r for r in rows if not _FORMATTED.search(r["pytorch"])] + for key in ["pytorch", "mlx_fp32", "mlx_fp16"]: + pairs = [(r["reference"], normalize(r[key])) for r in rows] + out[key] = { + "wer": corpus_error_rate(pairs, "word"), + "cer": corpus_error_rate(pairs, "char"), + "wer_plain_subset": corpus_error_rate( + [(r["reference"], normalize(r[key])) for r in plain], "word" + ), + } + if key != "pytorch": + out[key]["exact_match_vs_pytorch"] = sum(r[key] == r["pytorch"] for r in rows) + diffs = [r[f"{key}_enc_max_abs_diff"] for r in rows] + out[key]["enc_max_abs_diff"] = {"median": float(np.median(diffs)), "max": float(np.max(diffs))} + if "ctc_frames" in rows[0]: + agree = sum(r[f"{key}_ctc_frames_agree"] for r in rows) + out[key]["ctc_frame_agreement"] = agree / sum(r["ctc_frames"] for r in rows) + out["enc_abs_max_median"] = float(np.median([r["enc_abs_max"] for r in rows])) + out["clips"] = len(rows) + out["plain_subset_clips"] = len(plain) + return out + + +def main(): + parser = argparse.ArgumentParser(description=__doc__.splitlines()[0]) + parser.add_argument("--data", default=str(DEFAULT_DATA)) + parser.add_argument("--models-dir", default=str(HERE.parent / "eval-models")) + parser.add_argument("--out", default=str(HERE / "results")) + parser.add_argument("--models", nargs="+", default=MODELS) + parser.add_argument("--limit", type=int, default=None, help="first N clips only") + parser.add_argument("--summarize-only", action="store_true", help="recompute summary.json from saved /.jsonl") + args = parser.parse_args() + + mx.set_cache_limit(2**30) # see bench_speed.py: MLX's buffer cache grows with every new input length + clips = load_corpus(Path(args.data))[: args.limit] + out = Path(args.out) + out.mkdir(parents=True, exist_ok=True) + summary_path = out / "summary.json" + summary = json.loads(summary_path.read_text()) if summary_path.exists() else {} + summary["machine"] = machine_info() + summary["corpus"] = {"clips": len(clips), "seconds": sum(c.duration for c in clips)} + + for name in args.models: + t0 = time.time() + if args.summarize_only: + rows = [json.loads(line) for line in open(out / f"{name}.jsonl")] + else: + rows = evaluate(name, clips, Path(args.models_dir)) + with open(out / f"{name}.jsonl", "w") as f: + for r in rows: + f.write(json.dumps(r, ensure_ascii=False) + "\n") + summary[name] = summarize(rows) + s = summary[name] + print( + f"{name}: WER pt {s['pytorch']['wer']['rate']:.2%} | mlx32 {s['mlx_fp32']['wer']['rate']:.2%} " + f"| mlx16 {s['mlx_fp16']['wer']['rate']:.2%} | exact32 {s['mlx_fp32']['exact_match_vs_pytorch']}/{len(rows)} " + f"exact16 {s['mlx_fp16']['exact_match_vs_pytorch']}/{len(rows)} ({time.time() - t0:.0f}s)", + flush=True, + ) + summary_path.write_text(json.dumps(summary, ensure_ascii=False, indent=2)) + + +if __name__ == "__main__": + main() diff --git a/mlx_convert/eval/test_decode.py b/mlx_convert/eval/test_decode.py new file mode 100644 index 0000000..00edd7c --- /dev/null +++ b/mlx_convert/eval/test_decode.py @@ -0,0 +1,32 @@ +"""decode_sentencepiece must reproduce SentencePieceProcessor.decode exactly. + +Run: pytest mlx_convert/eval/test_decode.py (needs converted E2E models in +mlx_convert/eval-models/, see convert_gigaam_to_mlx.py). +""" +import json +import random +import sys +from pathlib import Path + +import pytest + +ROOT = Path(__file__).parent.parent +sys.path.insert(0, str(ROOT)) +from gigaam_mlx import decode_sentencepiece # noqa: E402 + + +@pytest.mark.parametrize("name", ["v3_e2e_ctc", "v3_e2e_rnnt"]) +def test_matches_sentencepiece(name): + spm = pytest.importorskip("sentencepiece") + model_dir = ROOT / "eval-models" / name + if not (model_dir / "tokenizer.model").exists(): + pytest.skip(f"{model_dir} not converted") + sp = spm.SentencePieceProcessor(model_file=str(model_dir / "tokenizer.model")) + cfg = json.loads((model_dir / "config.json").read_text()) + pieces, control = cfg["vocabulary"], cfg["tokenizer_control_token_ids"] + assert len(pieces) == sp.get_piece_size() + + rng = random.Random(0) + for _ in range(20000): + ids = [rng.randrange(len(pieces)) for _ in range(rng.randint(1, 30))] + assert decode_sentencepiece(ids, pieces, control) == sp.decode(ids), ids diff --git a/mlx_convert/gigaam-cli b/mlx_convert/gigaam-cli new file mode 100755 index 0000000..a8acb07 --- /dev/null +++ b/mlx_convert/gigaam-cli @@ -0,0 +1,64 @@ +#!/usr/bin/env python3 +""" +gigaam-cli — GigaAM v3 transcription CLI for MLX (CTC and RNNT). + +Usage: + gigaam-cli -f audio.wav [-m /path/to/model_dir] + gigaam-cli --help +""" +import argparse +import os +import sys +import time + +# Resolve default model dir relative to this script +SCRIPT_DIR = os.path.dirname(os.path.abspath(__file__)) +DEFAULT_MODEL = os.path.join(SCRIPT_DIR, "gigaam-v3-ctc-mlx") + + +def main(): + parser = argparse.ArgumentParser(description="GigaAM v3 transcription (MLX)") + parser.add_argument("-f", "--file", required=True, help="Audio file path") + parser.add_argument("-m", "--model", default=DEFAULT_MODEL, help="Model directory") + parser.add_argument("-l", "--language", default="ru", help="Language (ignored, always ru)") + parser.add_argument("-nt", "--no-timestamps", action="store_true", help="(compat flag, ignored)") + parser.add_argument("--no-prints", action="store_true", help="Suppress info messages") + args = parser.parse_args() + + if not os.path.isfile(args.file): + print(f"Error: file not found: {args.file}", file=sys.stderr) + sys.exit(1) + + if not os.path.isdir(args.model): + print(f"Error: model dir not found: {args.model}", file=sys.stderr) + sys.exit(1) + + # Add script directory to path so gigaam_mlx can be imported + sys.path.insert(0, SCRIPT_DIR) + from gigaam_mlx import load_model, load_audio + + if not args.no_prints: + print(f"gigaam-cli: loading model from {args.model}", file=sys.stderr) + + t0 = time.time() + model = load_model(args.model) + load_time = time.time() - t0 + + if not args.no_prints: + print(f"gigaam-cli: model loaded in {load_time:.2f}s", file=sys.stderr) + + audio = load_audio(args.file) + + t0 = time.time() + text = model.transcribe(audio) + transcribe_time = time.time() - t0 + + if not args.no_prints: + print(f"gigaam-cli: transcribed in {transcribe_time:.2f}s", file=sys.stderr) + + # Output ONLY the text to stdout (like whisper-cli -nt) + print(text) + + +if __name__ == "__main__": + main() diff --git a/mlx_convert/gigaam-stream b/mlx_convert/gigaam-stream new file mode 100755 index 0000000..67629ee --- /dev/null +++ b/mlx_convert/gigaam-stream @@ -0,0 +1,208 @@ +#!/usr/bin/env python3 +""" +gigaam-stream — Real-time speech recognition using GigaAM v3 on MLX. + +Growing-buffer pseudo-streaming: captures audio, transcribes from start +up to current position each step. GigaAM is fast enough on Apple Silicon +that even 30s audio processes in well under a second. + +Usage: + gigaam-stream # live microphone + gigaam-stream --file audio.wav # file streaming + gigaam-stream --step 1000 --context 5000 # custom timing + +Requirements: + pip install sounddevice (for live microphone capture) +""" +import argparse +import os +import sys +import time +import threading +import signal + +import numpy as np +import mlx.core as mx + +# Every step transcribes a buffer of a new length; without a cap MLX keeps the +# freed buffers of each shape cached, growing towards most of unified memory. +mx.set_cache_limit(1 << 30) + +SCRIPT_DIR = os.path.dirname(os.path.abspath(__file__)) +DEFAULT_MODEL = os.path.join(SCRIPT_DIR, "gigaam-v3-ctc-mlx") +SAMPLE_RATE = 16000 + + +def parse_args(): + p = argparse.ArgumentParser(description="GigaAM v3 real-time streaming (MLX)") + p.add_argument("-m", "--model", default=DEFAULT_MODEL, help="Model directory") + p.add_argument("--step", type=int, default=2000, + help="Process every N ms (default: 2000)") + p.add_argument("--file", "-f", type=str, default=None, + help="Transcribe file instead of microphone") + p.add_argument("--no-overwrite", action="store_true", + help="Print incrementally instead of overwriting line") + p.add_argument("--vad-threshold", type=float, default=0.003, + help="VAD energy threshold (default: 0.003)") + return p.parse_args() + + +def clear_line(): + sys.stdout.write("\33[2K\r") + sys.stdout.flush() + + +def stream_from_file(model, audio, args): + """Stream a pre-recorded file with growing buffer.""" + from gigaam_mlx import StreamingConfig + + config = StreamingConfig(step_duration=args.step / 1000.0) + total_dur = audio.shape[0] / SAMPLE_RATE + + print(f"\033[90m[Streaming {total_dur:.1f}s audio, step={args.step}ms]\033[0m\n") + + t0 = time.time() + for result in model.stream_generate(audio, config=config): + elapsed = time.time() - t0 + if args.no_overwrite: + if result.text: + sys.stdout.write(result.text) + sys.stdout.flush() + else: + clear_line() + # Show current transcription with position indicator + pos = f"[{result.audio_position:.1f}s]" + sys.stdout.write(f"\r\033[90m{pos}\033[0m {result.cumulative_text}") + sys.stdout.flush() + + if result.is_final: + print() + break + + total_time = time.time() - t0 + print(f"\n\033[90m[Done in {total_time:.2f}s | {total_dur/total_time:.0f}x realtime]\033[0m") + + +def stream_from_microphone(model, args): + """Live streaming from microphone with growing buffer.""" + try: + import sounddevice as sd + except ImportError: + print("Error: sounddevice not installed. Run: uv pip install sounddevice", + file=sys.stderr) + sys.exit(1) + + step_ms = args.step + step_sec = step_ms / 1000.0 + + # Shared state + audio_buffer = np.zeros(0, dtype=np.float32) + buffer_lock = threading.Lock() + running = True + prev_text = "" + + def audio_callback(indata, frames, time_info, status): + nonlocal audio_buffer + if status: + print(f"\033[90m[audio: {status}]\033[0m", file=sys.stderr) + chunk = indata[:, 0].copy() + with buffer_lock: + audio_buffer = np.concatenate([audio_buffer, chunk]) + # Cap at 30s + max_buf = SAMPLE_RATE * 30 + if len(audio_buffer) > max_buf: + audio_buffer = audio_buffer[-max_buf:] + + def signal_handler(sig, frame): + nonlocal running + running = False + + signal.signal(signal.SIGINT, signal_handler) + + print(f"\033[90m[GigaAM streaming — step={step_ms}ms, Ctrl+C to stop]\033[0m") + print(f"\033[90m[Start speaking...]\033[0m\n") + + stream = sd.InputStream( + samplerate=SAMPLE_RATE, + channels=1, + dtype="float32", + blocksize=int(SAMPLE_RATE * 0.1), + callback=audio_callback, + ) + stream.start() + + last_process_time = 0 + silence_count = 0 + MAX_SILENCE = 5 # stop after 5 consecutive silent steps + + try: + while running: + now = time.time() + if now - last_process_time < step_sec: + time.sleep(0.05) + continue + + last_process_time = now + + with buffer_lock: + if len(audio_buffer) < int(SAMPLE_RATE * 0.3): + continue + buf_copy = audio_buffer.copy() + + # Simple energy VAD on latest step + latest = buf_copy[-min(int(step_sec * SAMPLE_RATE), len(buf_copy)):] + energy = np.sqrt(np.mean(latest ** 2)) + + if energy < args.vad_threshold: + silence_count += 1 + continue + else: + silence_count = 0 + + # Transcribe entire buffer (growing buffer approach) + t0 = time.time() + audio_mx = mx.array(buf_copy) + result = model.stream_live(audio_mx) + dt = time.time() - t0 + + text = result.cumulative_text + + if text and text != prev_text: + clear_line() + dur = len(buf_copy) / SAMPLE_RATE + sys.stdout.write( + f"\r\033[90m[{dur:.1f}s {dt*1000:.0f}ms]\033[0m {text}" + ) + sys.stdout.flush() + prev_text = text + + except KeyboardInterrupt: + pass + finally: + stream.stop() + stream.close() + print(f"\n") + if prev_text: + print(f"\033[1mFinal:\033[0m {prev_text}") + + +def main(): + args = parse_args() + + sys.path.insert(0, SCRIPT_DIR) + from gigaam_mlx import load_model, load_audio + + print(f"\033[90mLoading GigaAM v3 MLX...\033[0m", file=sys.stderr) + t0 = time.time() + model = load_model(args.model) + print(f"\033[90mLoaded in {time.time() - t0:.2f}s\033[0m", file=sys.stderr) + + if args.file: + audio = load_audio(args.file) + stream_from_file(model, audio, args) + else: + stream_from_microphone(model, args) + + +if __name__ == "__main__": + main() diff --git a/mlx_convert/gigaam-transcribe b/mlx_convert/gigaam-transcribe new file mode 100755 index 0000000..af387fa --- /dev/null +++ b/mlx_convert/gigaam-transcribe @@ -0,0 +1,15 @@ +#!/bin/bash +# gigaam-transcribe — shell wrapper for gigaam-cli +# Activates the correct Python venv and runs gigaam-cli +# Install: ln -s /path/to/gigaam-transcribe /usr/local/bin/gigaam-transcribe + +SCRIPT_DIR="$(cd "$(dirname "$0")" && pwd)" +PYTHON="${SCRIPT_DIR}/.venv/bin/python" + +if [ ! -f "$PYTHON" ]; then + echo "Error: Python venv not found at ${SCRIPT_DIR}/.venv" >&2 + echo "Run: cd ${SCRIPT_DIR} && uv venv .venv && uv pip install mlx safetensors numpy" >&2 + exit 1 +fi + +exec "$PYTHON" "${SCRIPT_DIR}/gigaam-cli" "$@" diff --git a/mlx_convert/gigaam_mlx.py b/mlx_convert/gigaam_mlx.py new file mode 100644 index 0000000..f16828d --- /dev/null +++ b/mlx_convert/gigaam_mlx.py @@ -0,0 +1,790 @@ +""" +GigaAM v3 CTC model in MLX. + +Architecture: + - FeatureExtractor: log-mel spectrogram (computed via numpy/MLX at runtime) + - ConformerEncoder: conv1d subsampling + rotary positional embeddings + 16 conformer layers + - CTCHead: Conv1d(768, num_classes, kernel=1) + - CTC greedy decoding + +Supports pseudo-streaming via sliding window (like whisper-stream / mlx-audio Parakeet). +""" +import json +import math +from dataclasses import dataclass, field +from pathlib import Path +from typing import Generator, List, Optional, Tuple + +import mlx.core as mx +import mlx.nn as nn +import numpy as np + + +# ─────────────────────────── Config ─────────────────────────── + +@dataclass +class GigaAMConfig: + model_type: str = "gigaam" + model_name: str = "v3_ctc" + sample_rate: int = 16000 + # preprocessor + features: int = 64 + win_length: int = 320 + hop_length: int = 160 + n_fft: int = 320 + center: bool = False + # encoder + feat_in: int = 64 + n_layers: int = 16 + d_model: int = 768 + subsampling: str = "conv1d" + subs_kernel_size: int = 5 + subsampling_factor: int = 4 + ff_expansion_factor: int = 4 + self_attention_model: str = "rotary" + pos_emb_max_len: int = 5000 + n_heads: int = 16 + conv_kernel_size: int = 5 + conv_norm_type: str = "layer_norm" + # head + head_type: str = "ctc" + num_classes: int = 34 + + # rnnt specific + rnnt_pred_hidden: int = 320 + rnnt_joint_hidden: int = 320 + rnnt_max_symbols: int = 10 + + # vocabulary + vocabulary: Optional[List[str]] = None + tokenizer_type: Optional[str] = None # None = char-wise, "sentencepiece" = E2E models + tokenizer_control_token_ids: List[int] = field(default_factory=list) + + @classmethod + def from_file(cls, path: str) -> "GigaAMConfig": + with open(path) as f: + d = json.load(f) + enc = d.get("encoder", {}) + pre = d.get("preprocessor", {}) + head = d.get("head", {}) + return cls( + model_type=d.get("model_type", "gigaam"), + model_name=d.get("model_name", "v3_ctc"), + sample_rate=d.get("sample_rate", 16000), + features=pre.get("features", 64), + win_length=pre.get("win_length", 320), + hop_length=pre.get("hop_length", 160), + n_fft=pre.get("n_fft", 320), + center=pre.get("center", False), + feat_in=enc.get("feat_in", 64), + n_layers=enc.get("n_layers", 16), + d_model=enc.get("d_model", 768), + subsampling=enc.get("subsampling", "conv1d"), + subs_kernel_size=enc.get("subs_kernel_size", 5), + subsampling_factor=enc.get("subsampling_factor", 4), + ff_expansion_factor=enc.get("ff_expansion_factor", 4), + self_attention_model=enc.get("self_attention_model", "rotary"), + pos_emb_max_len=enc.get("pos_emb_max_len", 5000), + n_heads=enc.get("n_heads", 16), + conv_kernel_size=enc.get("conv_kernel_size", 5), + conv_norm_type=enc.get("conv_norm_type", "layer_norm"), + head_type=d.get("head_type", "ctc"), + num_classes=head.get("num_classes", 34) if "num_classes" in head else head.get("decoder", {}).get("num_classes", 34), + rnnt_pred_hidden=head.get("decoder", {}).get("pred_hidden", 320), + rnnt_joint_hidden=head.get("joint", {}).get("joint_hidden", 320), + vocabulary=d.get("vocabulary"), + tokenizer_type=d.get("tokenizer_type"), + tokenizer_control_token_ids=d.get("tokenizer_control_token_ids", []), + ) + + +# ─────────────────────── Audio preprocessing ─────────────────────── + +def mel_filters(sr: int, n_fft: int, n_mels: int) -> mx.array: + """HTK mel filterbank (matching torchaudio default for GigaAM).""" + def hz_to_mel(f): + return 2595.0 * math.log10(1.0 + f / 700.0) + def mel_to_hz(m): + return 700.0 * (10.0 ** (m / 2595.0) - 1.0) + + f_min, f_max = 0.0, sr / 2.0 + mel_min, mel_max = hz_to_mel(f_min), hz_to_mel(f_max) + mel_points = np.linspace(mel_min, mel_max, n_mels + 2) + hz_points = np.array([mel_to_hz(m) for m in mel_points]) + + bins = np.floor((n_fft + 1) * hz_points / sr).astype(int) + + fb = np.zeros((n_fft // 2 + 1, n_mels), dtype=np.float32) + for i in range(n_mels): + lo, mid, hi = bins[i], bins[i + 1], bins[i + 2] + for k in range(lo, mid): + if mid != lo: + fb[k, i] = (k - lo) / (mid - lo) + for k in range(mid, hi): + if hi != mid: + fb[k, i] = (hi - k) / (hi - mid) + return mx.array(fb) + + +def hanning_window(size: int) -> mx.array: + """Hanning window.""" + n = np.arange(size, dtype=np.float32) + return mx.array(0.5 - 0.5 * np.cos(2.0 * np.pi * n / size)) + + +def stft(signal: mx.array, n_fft: int, hop: int, win_len: int, window: mx.array) -> mx.array: + """Simple STFT using mx operations.""" + # Pad if necessary + pad_amount = n_fft // 2 + # For center=False, no padding + length = signal.shape[-1] + # Number of frames + n_frames = 1 + (length - win_len) // hop + + # Build frames via strided indexing + indices = mx.arange(win_len)[None, :] + (mx.arange(n_frames) * hop)[:, None] + frames = signal[indices] * window[None, :] + + # Zero-pad to n_fft if needed + if win_len < n_fft: + pad_size = n_fft - win_len + frames = mx.pad(frames, ((0, 0), (0, pad_size))) + + # Real FFT + spectrum = mx.fft.rfft(frames) + return spectrum # [n_frames, n_fft//2 + 1] + + +def log_mel_spectrogram(audio: mx.array, cfg: GigaAMConfig, + mel_fb: Optional[mx.array] = None, + stft_win: Optional[mx.array] = None) -> mx.array: + """Compute log-mel spectrogram matching GigaAM FeatureExtractor.""" + window = stft_win if stft_win is not None else hanning_window(cfg.win_length) + spec = stft(audio, cfg.n_fft, cfg.hop_length, cfg.win_length, window) + # Power spectrum + power = mx.square(mx.abs(spec)) # [T, n_fft//2+1] + # Mel filterbank + if mel_fb is not None: + filters = mel_fb + else: + filters = mel_filters(cfg.sample_rate, cfg.n_fft, cfg.features) # [n_fft//2+1, n_mels] + mel = power @ filters # [T, n_mels] + # Log + log_mel = mx.log(mx.clip(mel, 1e-9, 1e9)) + return log_mel # [T, n_mels] + + +# ─────────────────────── Model layers ─────────────────────── + +class Conv1dSubsampling(nn.Module): + """Conv1d striding subsampling: 2 conv1d layers with stride=2, ReLU.""" + def __init__(self, cfg: GigaAMConfig): + super().__init__() + ks = cfg.subs_kernel_size + pad = (ks - 1) // 2 + n_subs = int(math.log2(cfg.subsampling_factor)) + + layers = [] + in_ch = cfg.feat_in + for _ in range(n_subs): + layers.append(nn.Conv1d(in_ch, cfg.d_model, kernel_size=ks, stride=2, padding=pad)) + layers.append(nn.ReLU()) + in_ch = cfg.d_model + self.conv = layers + self._n_subs = n_subs + self._ks = ks + self._pad = pad + + def __call__(self, x: mx.array, lengths: mx.array) -> Tuple[mx.array, mx.array]: + """x: [B, T, feat_in] → [B, T', d_model]""" + for layer in self.conv: + x = layer(x) + # Compute output lengths + for _ in range(self._n_subs): + lengths = mx.floor((lengths.astype(mx.float32) + 2 * self._pad - self._ks) / 2 + 1).astype(mx.int32) + return x, lengths + + +class RotaryPositionalEmbedding(nn.Module): + """Rotary positional embeddings (RoPE).""" + def __init__(self, dim: int, base: int = 10000, max_len: int = 5000): + super().__init__() + self.dim = dim + self.base = base + inv_freq = 1.0 / (base ** (np.arange(0, dim, 2, dtype=np.float32) / dim)) + self._inv_freq = mx.array(inv_freq) + + def __call__(self, seq_len: int) -> Tuple[mx.array, mx.array]: + t = mx.arange(seq_len).astype(mx.float32) + freqs = mx.outer(t, self._inv_freq) + emb = mx.concatenate([freqs, freqs], axis=-1) + return mx.cos(emb), mx.sin(emb) + + +def rotate_half(x: mx.array) -> mx.array: + """Rotates half the hidden dims of the input.""" + d = x.shape[-1] // 2 + x1, x2 = x[..., :d], x[..., d:] + return mx.concatenate([-x2, x1], axis=-1) + + +def apply_rotary_emb(q: mx.array, k: mx.array, cos: mx.array, sin: mx.array) -> Tuple[mx.array, mx.array]: + """Apply rotary embeddings. q,k: [B, H, T, D], cos,sin: [T, D]""" + cos = cos[None, None, :, :] # [1, 1, T, D] + sin = sin[None, None, :, :] + q_rot = q * cos + rotate_half(q) * sin + k_rot = k * cos + rotate_half(k) * sin + return q_rot, k_rot + + +class RotaryMultiHeadAttention(nn.Module): + """Multi-head attention with rotary positional embeddings. + + GigaAM applies RoPE BEFORE linear projections: + 1. Reshape raw input to multi-head: [B, T, D] → [B, T, H, d_k] + 2. Apply RoPE to query/key heads + 3. Reshape back to [B, T, D] + 4. Project through linear_q/k/v + 5. Standard scaled dot-product attention + """ + def __init__(self, d_model: int, n_heads: int): + super().__init__() + self.n_heads = n_heads + self.d_k = d_model // n_heads + self.linear_q = nn.Linear(d_model, d_model) + self.linear_k = nn.Linear(d_model, d_model) + self.linear_v = nn.Linear(d_model, d_model) + self.linear_out = nn.Linear(d_model, d_model) + + def __call__(self, x: mx.array, cos: mx.array, sin: mx.array, + mask: Optional[mx.array] = None) -> mx.array: + B, T, D = x.shape + H, d_k = self.n_heads, self.d_k + + # 1. Reshape raw input to multi-head for RoPE: [B, T, H, d_k] + x_heads = x.reshape(B, T, H, d_k) + + # 2. Apply RoPE to query and key (same input for self-attention) + # cos, sin: [T, d_k] → [1, T, 1, d_k] + cos_e = cos[None, :, None, :] + sin_e = sin[None, :, None, :] + q_rot = x_heads * cos_e + rotate_half(x_heads) * sin_e + k_rot = x_heads * cos_e + rotate_half(x_heads) * sin_e + + # 3. Reshape back to [B, T, D] + q_rot = q_rot.reshape(B, T, D) + k_rot = k_rot.reshape(B, T, D) + + # 4. Project through linear layers + q = self.linear_q(q_rot).reshape(B, T, H, d_k).transpose(0, 2, 1, 3) # [B, H, T, d_k] + k = self.linear_k(k_rot).reshape(B, T, H, d_k).transpose(0, 2, 1, 3) + v = self.linear_v(x).reshape(B, T, H, d_k).transpose(0, 2, 1, 3) # value uses original x + + # 5. Scaled dot-product attention + scale = math.sqrt(d_k) + attn = (q @ k.transpose(0, 1, 3, 2)) / scale + if mask is not None: + attn = attn + mask + attn = mx.softmax(attn, axis=-1) + out = attn @ v # [B, H, T, d_k] + out = out.transpose(0, 2, 1, 3).reshape(B, T, -1) + return self.linear_out(out) + + +class ConformerConvolution(nn.Module): + """Conformer convolution module with LayerNorm.""" + def __init__(self, d_model: int, kernel_size: int): + super().__init__() + pad = (kernel_size - 1) // 2 + self.pointwise_conv1 = nn.Conv1d(d_model, d_model * 2, kernel_size=1) + self.depthwise_conv = nn.Conv1d(d_model, d_model, kernel_size=kernel_size, + padding=pad, groups=d_model) + # GigaAM v3 uses layer_norm for conv_norm_type + self.batch_norm = nn.LayerNorm(d_model) + self.activation = nn.SiLU() + self.pointwise_conv2 = nn.Conv1d(d_model, d_model, kernel_size=1) + + def __call__(self, x: mx.array) -> mx.array: + # x: [B, T, D] + x = self.pointwise_conv1(x) # → [B, T, 2*D] + # GLU: split last dim in half, apply sigmoid gate + half = x.shape[-1] // 2 + x = x[..., :half] * mx.sigmoid(x[..., half:]) # GLU + x = self.depthwise_conv(x) # → [B, T, D] + x = self.batch_norm(x) + x = self.activation(x) + x = self.pointwise_conv2(x) # → [B, T, D] + return x + + +class ConformerFeedForward(nn.Module): + """Conformer feed-forward module.""" + def __init__(self, d_model: int, d_ff: int): + super().__init__() + self.linear1 = nn.Linear(d_model, d_ff) + self.activation = nn.SiLU() + self.linear2 = nn.Linear(d_ff, d_model) + + def __call__(self, x: mx.array) -> mx.array: + return self.linear2(self.activation(self.linear1(x))) + + +class ConformerLayer(nn.Module): + """Single conformer layer.""" + def __init__(self, cfg: GigaAMConfig): + super().__init__() + d = cfg.d_model + d_ff = d * cfg.ff_expansion_factor + + self.norm_feed_forward1 = nn.LayerNorm(d) + self.feed_forward1 = ConformerFeedForward(d, d_ff) + self.norm_self_att = nn.LayerNorm(d) + self.self_attn = RotaryMultiHeadAttention(d, cfg.n_heads) + self.norm_conv = nn.LayerNorm(d) + self.conv = ConformerConvolution(d, cfg.conv_kernel_size) + self.norm_feed_forward2 = nn.LayerNorm(d) + self.feed_forward2 = ConformerFeedForward(d, d_ff) + self.norm_out = nn.LayerNorm(d) + + def __call__(self, x: mx.array, cos: mx.array, sin: mx.array, + mask: Optional[mx.array] = None) -> mx.array: + # FF1 + residual = x + x = self.norm_feed_forward1(x) + x = self.feed_forward1(x) + residual = residual + x * 0.5 + + # Self-attention + x = self.norm_self_att(residual) + x = self.self_attn(x, cos, sin, mask=mask) + residual = residual + x + + # Conv + x = self.norm_conv(residual) + x = self.conv(x) + residual = residual + x + + # FF2 + x = self.norm_feed_forward2(residual) + x = self.feed_forward2(x) + residual = residual + x * 0.5 + + return self.norm_out(residual) + + +class ConformerEncoder(nn.Module): + """GigaAM Conformer encoder.""" + def __init__(self, cfg: GigaAMConfig): + super().__init__() + self.pre_encode = Conv1dSubsampling(cfg) + # Upstream builds RotaryPositionalEmbedding(d_model // n_heads, pos_emb_max_len), + # whose second positional parameter is `base`: the RoPE base is 5000, not 10000. + self.pos_enc = RotaryPositionalEmbedding( + cfg.d_model // cfg.n_heads, + base=cfg.pos_emb_max_len, + max_len=cfg.pos_emb_max_len, + ) + self.layers = [ConformerLayer(cfg) for _ in range(cfg.n_layers)] + + def __call__(self, features: mx.array, lengths: mx.array) -> Tuple[mx.array, mx.array]: + """features: [B, T, feat_in] → encoded: [B, T', D], lengths: [B]""" + x, lengths = self.pre_encode(features, lengths) + T = x.shape[1] + cos, sin = self.pos_enc(T) + for layer in self.layers: + x = layer(x, cos, sin) + return x, lengths + + +class RNNTJoint(nn.Module): + def __init__(self, enc_hidden: int, pred_hidden: int, joint_hidden: int, num_classes: int): + super().__init__() + self.enc = nn.Linear(enc_hidden, joint_hidden) + self.pred = nn.Linear(pred_hidden, joint_hidden) + self.joint_net_linear = nn.Linear(joint_hidden, num_classes) + + def __call__(self, encoder_out: mx.array, decoder_out: mx.array) -> mx.array: + enc = mx.expand_dims(self.enc(encoder_out), 2) + pred = mx.expand_dims(self.pred(decoder_out), 1) + joint_rep = mx.maximum(enc + pred, 0.0) + return self.joint_net_linear(joint_rep) + + +class RNNTDecoder(nn.Module): + def __init__(self, num_classes: int, pred_hidden: int): + super().__init__() + self.blank_id = num_classes - 1 + self.pred_hidden = pred_hidden + self.embed = nn.Embedding(num_classes, pred_hidden) + self.lstm = nn.LSTM(pred_hidden, pred_hidden) + + def predict(self, x: Optional[int], state: Optional[Tuple[mx.array, mx.array]], batch_size: int = 1): + if x is not None: + emb = self.embed(mx.array([[x]])) + else: + emb = mx.zeros((batch_size, 1, self.pred_hidden)) + + h, c = state if state is not None else (None, None) + out, cell = self.lstm(emb, hidden=h, cell=c) + return out, (out[:, -1, :], cell[:, -1, :]) + + +class RNNTHead(nn.Module): + def __init__(self, feat_in: int, pred_hidden: int, joint_hidden: int, num_classes: int): + super().__init__() + self.decoder = RNNTDecoder(num_classes, pred_hidden) + self.joint = RNNTJoint(feat_in, pred_hidden, joint_hidden, num_classes) + + +class CTCHead(nn.Module): + """CTC decoder head: Conv1d(d_model, num_classes, kernel=1).""" + def __init__(self, feat_in: int, num_classes: int): + super().__init__() + self.decoder_layers = [nn.Conv1d(feat_in, num_classes, kernel_size=1)] + + def __call__(self, x: mx.array) -> mx.array: + """x: [B, T, D] → log_probs: [B, T, C]""" + logits = self.decoder_layers[0](x) # [B, T, C] + return mx.softmax(logits, axis=-1) # we'll use log later + + +# ─────────────────────── Streaming dataclasses ─────────────────────── + +@dataclass +class StreamingConfig: + """Configuration for pseudo-streaming transcription. + + Attributes: + chunk_duration: Duration of each audio chunk in seconds. + context_duration: Extra audio from before the chunk to keep as context. + The model always sees [context + chunk] for better accuracy. + step_duration: How much to advance between steps. If None, equals chunk_duration. + """ + chunk_duration: float = 2.0 + context_duration: float = 3.0 + step_duration: Optional[float] = None + + def __post_init__(self): + if self.step_duration is None: + self.step_duration = self.chunk_duration + + +@dataclass +class StreamingResult: + """Result from streaming transcription — mlx-audio compatible. + + Attributes: + text: New text since last emission. + is_final: True if this is the final result. + start_time: Start time in seconds. + end_time: End time in seconds. + progress: Progress 0.0–1.0. + audio_position: Current position in audio (seconds). + audio_duration: Total audio duration (seconds), 0 for live. + cumulative_text: Full accumulated transcription so far. + language: Language code. + """ + text: str + is_final: bool + start_time: float + end_time: float + progress: float = 0.0 + audio_position: float = 0.0 + audio_duration: float = 0.0 + cumulative_text: str = "" + language: str = "ru" + + +class GigaAM(nn.Module): + """Full GigaAM model (supports CTC and RNNT).""" + def __init__(self, cfg: GigaAMConfig): + super().__init__() + self.cfg = cfg + self.encoder = ConformerEncoder(cfg) + + if cfg.head_type == "rnnt": + self.head = RNNTHead( + feat_in=cfg.d_model, + pred_hidden=cfg.rnnt_pred_hidden, + joint_hidden=cfg.rnnt_joint_hidden, + num_classes=cfg.num_classes + ) + else: + self.head = CTCHead(cfg.d_model, cfg.num_classes) + + self.mel_filterbank: Optional[mx.array] = None + self.stft_window: Optional[mx.array] = None + + def __call__(self, features: mx.array, lengths: mx.array) -> Tuple[mx.array, mx.array]: + """Returns encoded features and their lengths.""" + encoded, enc_lengths = self.encoder(features, lengths) + return encoded, enc_lengths + + def _ctc_decode(self, log_probs: mx.array, enc_length: int) -> str: + """CTC greedy decode: collapse repeated + remove blanks.""" + vocab = self.cfg.vocabulary + blank_id = len(vocab) + labels = mx.argmax(log_probs, axis=-1) + labels_list = labels.tolist() + + result = [] + prev = -1 + for t, label in enumerate(labels_list): + if t >= enc_length: + break + if label == blank_id: + prev = label + continue + if label == prev: + continue + result.append(label) + prev = label + + return self._decode_tokens(result) + + def _decode_tokens(self, token_ids: List[int]) -> str: + """Token ids → text. Char-wise vocabularies are joined as is; + SentencePiece pieces are decoded like SentencePieceProcessor.decode.""" + vocab = self.cfg.vocabulary + if self.cfg.tokenizer_type != "sentencepiece": + return "".join(vocab[i] for i in token_ids) + return decode_sentencepiece(token_ids, vocab, self.cfg.tokenizer_control_token_ids) + + def _compute_features(self, audio: mx.array) -> Tuple[mx.array, mx.array]: + """Audio → mel features [1, T, features] + lengths [1].""" + mel = log_mel_spectrogram(audio, self.cfg, + mel_fb=self.mel_filterbank, + stft_win=self.stft_window) + mel = mx.expand_dims(mel, 0) + lengths = mx.array([mel.shape[1]]) + return mel, lengths + + def transcribe(self, audio: mx.array) -> str: + """Transcribe raw audio waveform → text.""" + mel, lengths = self._compute_features(audio) + encoded, enc_lengths = self(mel, lengths) + mx.eval(encoded, enc_lengths) + + if self.cfg.head_type == "rnnt": + return self._rnnt_decode(encoded[0], int(enc_lengths[0])) + else: + log_probs = self.head(encoded) + return self._ctc_decode(log_probs[0], int(enc_lengths[0])) + + def _rnnt_decode(self, encoded: mx.array, enc_length: int) -> str: + """Greedy decode for RNN-T.""" + vocab = self.cfg.vocabulary + blank_id = len(vocab) + max_symbols = self.cfg.rnnt_max_symbols + + hyp = [] + dec_state = None + last_label = None + + for t in range(enc_length): + # encoded: [T, D] -> [1, 1, D] + f = encoded[t:t+1, :].reshape(1, 1, -1) + + not_blank = True + new_symbols = 0 + while not_blank and new_symbols < max_symbols: + g, hidden = self.head.decoder.predict(last_label, dec_state, batch_size=1) + # g: [1, 1, D] + # joint output is [1, 1, 1, num_classes], we want argmax over classes + logits = self.head.joint(f, g) # [1, 1, 1, C] + k = int(mx.argmax(logits[0, 0, 0, :]).item()) + + if k == blank_id: + not_blank = False + else: + hyp.append(k) + dec_state = hidden + last_label = k + new_symbols += 1 + + return self._decode_tokens(hyp) + + def transcribe_chunk(self, audio: mx.array) -> str: + """Transcribe a single chunk (for streaming). Same as transcribe but clearer name.""" + return self.transcribe(audio) + + def stream_generate( + self, + audio: mx.array, + config: Optional[StreamingConfig] = None, + ) -> Generator[StreamingResult, None, None]: + """Pseudo-streaming transcription over pre-recorded audio. + + Uses growing buffer approach: each step transcribes from the start + up to the current position, so the cost of a step grows with the + buffer; see README → Evaluation for measured speed. + + For very long audio (>30s), falls back to sliding window of last 30s. + + Args: + audio: Raw audio waveform, mx.array, 16kHz mono. + config: StreamingConfig with step duration. + + Yields: + StreamingResult with incremental text for each step. + """ + if config is None: + config = StreamingConfig() + + sr = self.cfg.sample_rate + total_samples = audio.shape[0] + audio_duration = total_samples / sr + + step_samples = int(config.step_duration * sr) + max_window = int(30.0 * sr) # cap at 30s for memory/speed + + previous_text = "" + position = step_samples # start with first step + + while position <= total_samples: + is_last = position >= total_samples + # Growing buffer from start, capped at 30s + window_start = max(0, position - max_window) + window = audio[window_start:position] + + current_text = self.transcribe(window) + + # Incremental text + new_text = _incremental_text(previous_text, current_text) + previous_text = current_text + + audio_pos = position / sr + yield StreamingResult( + text=new_text, + is_final=is_last, + start_time=window_start / sr, + end_time=audio_pos, + progress=position / total_samples, + audio_position=audio_pos, + audio_duration=audio_duration, + cumulative_text=current_text, + language="ru", + ) + + if is_last: + break + + position = min(position + step_samples, total_samples) + + def stream_live( + self, + audio_buffer: mx.array, + ) -> StreamingResult: + """Transcribe a growing audio buffer for live microphone use. + + Call this repeatedly as new audio arrives. Transcribes the full buffer + (capped at 30s from the end). + + Args: + audio_buffer: The full accumulated audio so far. + + Returns: + StreamingResult with current full transcription. + """ + sr = self.cfg.sample_rate + total_samples = audio_buffer.shape[0] + max_window = int(30.0 * sr) + + start = max(0, total_samples - max_window) + window = audio_buffer[start:] + + text = self.transcribe(window) + + return StreamingResult( + text=text, + is_final=False, + start_time=start / sr, + end_time=total_samples / sr, + progress=0.0, + audio_position=total_samples / sr, + audio_duration=0.0, + cumulative_text=text, + language="ru", + ) + + +def decode_sentencepiece(token_ids: List[int], pieces: List[str], + control_ids: List[int] = ()) -> str: + """Decode SentencePiece token ids without the sentencepiece package. + + Control tokens are dropped, ``<0xNN>`` byte-fallback pieces become raw + bytes, ``▁`` marks a word boundary and ```` becomes `` ⁇ ``. + """ + control = set(control_ids) + out = bytearray() + for i in token_ids: + if i in control or not 0 <= i < len(pieces): + continue + piece = pieces[i] + if piece == "": + out += " \u2047 ".encode() + elif len(piece) == 6 and piece.startswith("<0x") and piece.endswith(">"): + out.append(int(piece[3:5], 16)) + else: + if not out: # SentencePiece drops the dummy-prefix "▁" of the first piece + piece = piece.removeprefix("\u2581") + out += piece.replace("\u2581", " ").encode() + return out.decode("utf-8", errors="replace") + + +def _incremental_text(previous: str, current: str) -> str: + """Find new text added to current vs previous. + + If current starts with previous → return the new suffix. + If model corrected → return full current (with marker). + """ + if not previous: + return current + if current.startswith(previous): + return current[len(previous):] + # Model self-corrected — return full updated text + return current + + +def load_model(model_dir: str, dtype: Optional[mx.Dtype] = mx.float32) -> GigaAM: + """Load converted GigaAM MLX model from directory. + + dtype: parameter dtype after loading. The default float32 upcasts the fp16 + file once: activations are float32 either way (the log-mel input is fp32), + and with fp16 weights every op pays for a cast — about 20% slower on M4 Pro. + Pass None to keep the stored dtype (≈0.4 GB less memory). + """ + model_dir = Path(model_dir) + cfg = GigaAMConfig.from_file(str(model_dir / "config.json")) + model = GigaAM(cfg) + weights = mx.load(str(model_dir / "model.safetensors")) + + # Extract preprocessing weights + mel_fb = weights.pop("mel_filterbank", None) + stft_win = weights.pop("stft_window", None) + + if dtype is not None: + weights = {k: v.astype(dtype) for k, v in weights.items()} + + # Load model weights + model.load_weights(list(weights.items()), strict=False) + + if mel_fb is not None: + model.mel_filterbank = mel_fb.astype(mx.float32) + if stft_win is not None: + model.stft_window = stft_win.astype(mx.float32) + mx.eval(model.parameters()) + return model + + +def load_audio(path: str, sample_rate: int = 16000) -> mx.array: + """Load audio via ffmpeg → mx.array.""" + import subprocess + cmd = [ + "ffmpeg", "-nostdin", "-threads", "0", "-i", path, + "-f", "s16le", "-ac", "1", "-acodec", "pcm_s16le", + "-ar", str(sample_rate), "-" + ] + result = subprocess.run(cmd, capture_output=True, check=True) + audio_np = np.frombuffer(result.stdout, dtype=np.int16).astype(np.float32) / 32768.0 + return mx.array(audio_np)