GigaAM Multilingual

README.md

@@ -0,0 +1,85 @@
+---
+license: mit
+language:
+- ru
+- en
+- kk
+- ky
+- uz
+pipeline_tag: automatic-speech-recognition
+---
+
+# GigaAM Multilingual
+
+GigaAM Multilingual is a family of Conformer-based foundation models (220M / 600M parameters) pre-trained with a HuBERT-style objective on **2M hours** of speech across **70+ languages** and fine-tuned for speech recognition with character-wise CTC decoders on 50K hours.
+
+The models provide best-in-class open-source quality on Russian, Kazakh, Kyrgyz, and Uzbek, and moderate quality on English.
+
+GigaAM Multilingual includes the following model variants:
+- `ssl` — 220M self-supervised encoder
+- `ctc` — 220M ASR model with a character-wise CTC decoder
+- `large_ssl` — 600M self-supervised encoder
+- `large_ctc` — 600M ASR model with a character-wise CTC decoder
+
+## Model Performance
+
+Word Error Rate (%) on [Common Voice](https://commonvoice.mozilla.org) (CV), [FLEURS](https://huggingface.co/datasets/google/fleurs), and internal in-the-wild test sets. Utterances longer than 30 s and references containing digits are excluded; references/hypotheses are normalized (lowercasing, punctuation removal, numerals→words); greedy decoding. Best per row in **bold**.
+
+| Language | Dataset | GigaAM Multilingual | GigaAM Multilingual Large | Omnilingual 1B (LLM) | Seamless M4T large v2 | Whisper large v3 |
+|:--------|:--------|------------:|------------:|---------------------:|----------------------:|-----------------:|
+| English | CV       |        26.0 |        21.5 |                 24.7 |              **16.2** |             20.0 |
+| English | FLEURS   |        12.2 |         9.4 |                  7.1 |                   5.8 |          **3.9** |
+| Russian | CV       |         7.1 |     **5.1** |                 13.6 |                    9.2 |              9.1 |
+| Russian | FLEURS   |         4.4 |     **3.0** |                  6.4 |                    4.6 |              3.1 |
+| Russian | Internal |         7.6 |     **6.0** |                 14.6 |                   16.1 |             10.1 |
+| Kazakh  | CV       |        17.2 |    **13.8** |                 23.7 |                   23.8 |             57.8 |
+| Kazakh  | FLEURS   |         5.2 |     **4.4** |                  6.6 |                    6.8 |             32.4 |
+| Kazakh  | Internal |        18.8 |    **15.8** |                 32.2 |                   62.9 |             65.2 |
+| Kyrgyz  | CV       |        12.5 |    **10.2** |                 21.6 |                   14.3 |             95.2 |
+| Kyrgyz  | FLEURS   |         7.0 |     **5.5** |                  8.1 |                    9.5 |             86.3 |
+| Kyrgyz  | Internal |        11.1 |     **9.8** |                 25.0 |                   78.3 |            102.2 |
+| Uzbek   | CV       |        11.3 |     **9.2** |                 32.8 |                   25.1 |            109.9 |
+| Uzbek   | FLEURS   |        10.0 |     **7.3** |                 15.4 |                   11.9 |            105.4 |
+| Uzbek   | Internal |        13.8 |    **12.7** |                 30.2 |                   40.0 |            120.6 |
+
+
+## Usage
+
+```python
+from transformers import AutoModel
+
+revision = "ctc"  # any variant: ssl, ctc, large_ssl, large_ctc
+model = AutoModel.from_pretrained(
+    "ai-sage/GigaAM-Multilingual",
+    revision=revision,
+    trust_remote_code=True,
+)
+
+transcription = model.transcribe("example.wav")
+print(transcription)
+```
+
+Recommended versions:
+- `torch==2.10.*`, `torchaudio==2.10.*`
+- `transformers==5.*`
+- (any) `hydra-core`, `omegaconf`
+
+Full usage guide can be found in the [example](https://github.com/salute-developers/GigaAM/blob/main/colab_example.ipynb).
+
+## Fine-tuning to a new language
+
+The `ssl` / `large_ssl` backbones can be adapted to a new language — see the [fine-tuning guide](https://github.com/salute-developers/GigaAM/blob/main/train_utils/README.md) and the [example notebook](https://github.com/salute-developers/GigaAM/blob/main/train_utils/example.ipynb).
+
+## Citation
+
+```bibtex
+@misc{gigaam_multilingual,
+      title={GigaAM Multilingual: Foundation Model for Underrepresented Languages},
+      author={Andrei Kuzmenko and Alexandr Maximenko and Aleksandr Kutsakov and Georgii Gospodinov and Dmitrii Bolotov and Oleg Kutuzov and Pavel Bogomolov and Fyodor Minkin},
+      year={2026},
+      eprint={2607.10371},
+      archivePrefix={arXiv},
+      primaryClass={eess.AS},
+      url={https://arxiv.org/abs/2607.10371}
+}
+```

config.json

@@ -0,0 +1,122 @@
+{
+  "model_type": "gigaam",
+  "auto_map": {
+    "AutoConfig": "modeling_gigaam.GigaAMConfig",
+    "AutoModel": "modeling_gigaam.GigaAMModel"
+  },
+  "cfg": {
+    "model": {
+      "cfg": {
+        "model_class": "ctc",
+        "sample_rate": 16000,
+        "model_name": "multilingual_ctc",
+        "preprocessor": {
+          "_target_": "modeling_gigaam.FeatureExtractor",
+          "sample_rate": 16000,
+          "features": 64,
+          "win_length": 320,
+          "hop_length": 160,
+          "n_fft": 320,
+          "center": false
+        },
+        "encoder": {
+          "_target_": "modeling_gigaam.ConformerEncoder",
+          "feat_in": 64,
+          "n_layers": 16,
+          "d_model": 768,
+          "subsampling": "conv1d",
+          "subs_kernel_size": 5,
+          "subsampling_factor": 4,
+          "ff_expansion_factor": 4,
+          "self_attention_model": "rotary",
+          "pos_emb_max_len": 5000,
+          "n_heads": 16,
+          "conv_norm_type": "layer_norm",
+          "conv_kernel_size": 5,
+          "flash_attn": false
+        },
+        "head": {
+          "_target_": "modeling_gigaam.CTCHead",
+          "feat_in": 768,
+          "num_classes": 71
+        },
+        "decoding": {
+          "_target_": "modeling_gigaam.CTCGreedyDecoding",
+          "vocabulary": [
+            " ",
+            "'",
+            "a",
+            "b",
+            "c",
+            "d",
+            "e",
+            "f",
+            "g",
+            "h",
+            "i",
+            "j",
+            "k",
+            "l",
+            "m",
+            "n",
+            "o",
+            "p",
+            "q",
+            "r",
+            "s",
+            "t",
+            "u",
+            "v",
+            "w",
+            "x",
+            "y",
+            "z",
+            "а",
+            "б",
+            "в",
+            "г",
+            "д",
+            "е",
+            "ж",
+            "з",
+            "и",
+            "й",
+            "к",
+            "л",
+            "м",
+            "н",
+            "о",
+            "п",
+            "р",
+            "с",
+            "т",
+            "у",
+            "ф",
+            "х",
+            "ц",
+            "ч",
+            "ш",
+            "щ",
+            "ъ",
+            "ы",
+            "ь",
+            "э",
+            "ю",
+            "я",
+            "ё",
+            "і",
+            "ғ",
+            "қ",
+            "ң",
+            "ү",
+            "ұ",
+            "һ",
+            "ә",
+            "ө"
+          ]
+        }
+      },
+      "_target_": "modeling_gigaam.GigaAMASR"
+    }
+  }
+}

modeling_gigaam.py

@@ -0,0 +1,2149 @@
+"""Self-contained GigaAM modeling file (generated by tools/build_hf_modeling.py)."""
+
+from __future__ import annotations
+
+import csv
+import math
+import os
+import sys
+import unicodedata
+import warnings
+from abc import ABC, abstractmethod
+from collections.abc import Iterable
+from contextlib import contextmanager
+from dataclasses import dataclass
+from pathlib import Path
+from subprocess import CalledProcessError, run
+from typing import TYPE_CHECKING, Dict, List, Optional, Tuple, Union, cast
+
+import hydra
+import numpy as np
+import omegaconf
+import soundfile as sf
+import torch
+import torch.nn.functional as F
+import torchaudio
+from huggingface_hub import snapshot_download
+from huggingface_hub.errors import LocalEntryNotFoundError
+from hydra.utils import instantiate
+from sentencepiece import SentencePieceProcessor
+from torch import Tensor, nn
+from torch.jit import TracerWarning
+from torch.torch_version import TorchVersion
+from torch.utils.checkpoint import checkpoint
+from torch.utils.data import DataLoader
+from transformers import PretrainedConfig, PreTrainedModel
+from transformers.utils import cached_file
+
+if TYPE_CHECKING:
+    try:
+        from pyannote.audio import Model, Pipeline
+    except ImportError:
+        pass
+
+DIR_NAME = os.path.dirname(os.path.abspath(__file__))
+sys.path.append(DIR_NAME)  # enable hydra targets like modeling_gigaam.<ClassName>
+
+# ==== gigaam/types.py ====
+
+
+@dataclass
+class AudioDatasetSample:
+    item: Union[str, np.ndarray, Tensor]
+    duration: float
+    text: Optional[str] = None
+    tokens: Optional[List[int]] = None
+
+
+@dataclass
+class Word:
+    text: str
+    start: float
+    end: float
+
+
+@dataclass
+class TranscriptionResult:
+    text: str
+    words: Optional[List[Word]] = None
+
+    def __str__(self) -> str:
+        return self.text
+
+
+@dataclass
+class Segment:
+    text: str
+    start: float
+    end: float
+    words: Optional[List[Word]] = None
+
+
+@dataclass
+class LongformTranscriptionResult:
+    segments: List[Segment]
+
+    @property
+    def words(self) -> List[Word]:
+        """Flatten all words from all segments."""
+        result = []
+        for seg in self.segments:
+            if seg.words:
+                result.extend(seg.words)
+        return result
+
+    @property
+    def has_word_timestamps(self) -> bool:
+        return bool(self.segments) and self.segments[0].words is not None
+
+    @property
+    def text(self) -> str:
+        return " ".join(s.text for s in self.segments)
+
+    def __str__(self) -> str:
+        return self.text
+
+    def __iter__(self):
+        return iter(self.segments)
+
+    def __len__(self) -> int:
+        return len(self.segments)
+
+
+# ==== gigaam/preprocess.py ====
+
+SAMPLE_RATE = 16000
+
+
+def load_audio(audio_path: str, sample_rate: int = SAMPLE_RATE) -> Tensor:
+    """
+    Load an audio file and resample it to the specified sample rate.
+    """
+    cmd = [
+        "ffmpeg",
+        "-nostdin",
+        "-threads",
+        "0",
+        "-i",
+        audio_path,
+        "-f",
+        "s16le",
+        "-ac",
+        "1",
+        "-acodec",
+        "pcm_s16le",
+        "-ar",
+        str(sample_rate),
+        "-",
+    ]
+    try:
+        audio = run(cmd, capture_output=True, check=True).stdout
+    except CalledProcessError as exc:
+        raise RuntimeError("Failed to load audio") from exc
+
+    with warnings.catch_warnings():
+        warnings.simplefilter("ignore", category=UserWarning)
+        return torch.frombuffer(audio, dtype=torch.int16).float() / 32768.0
+
+
+class SpecScaler(nn.Module):
+    """
+    Module that applies logarithmic scaling to spectrogram values.
+    This module clamps the input values within a certain range and then applies a natural logarithm.
+    """
+
+    def forward(self, x: Tensor) -> Tensor:
+        return torch.log(x.clamp_(1e-9, 1e9))
+
+
+class FeatureExtractor(nn.Module):
+    """
+    Module for extracting Log-mel spectrogram features from raw audio signals.
+    This module uses Torchaudio's MelSpectrogram transform to extract features
+    and applies logarithmic scaling.
+    """
+
+    def __init__(self, sample_rate: int, features: int, **kwargs):
+        super().__init__()
+        self.hop_length = kwargs.get("hop_length", sample_rate // 100)
+        self.win_length = kwargs.get("win_length", sample_rate // 40)
+        self.n_fft = kwargs.get("n_fft", sample_rate // 40)
+        self.center = kwargs.get("center", True)
+        self.featurizer = nn.Sequential(
+            torchaudio.transforms.MelSpectrogram(
+                sample_rate=sample_rate,
+                n_mels=features,
+                win_length=self.win_length,
+                hop_length=self.hop_length,
+                n_fft=self.n_fft,
+                center=self.center,
+            ),
+            SpecScaler(),
+        )
+
+    def out_len(self, input_lengths: Tensor) -> Tensor:
+        """
+        Calculates the output length after the feature extraction process.
+        """
+        if self.center:
+            return (
+                input_lengths.div(self.hop_length, rounding_mode="floor").add(1).long()
+            )
+        else:
+            return (
+                (input_lengths - self.win_length)
+                .div(self.hop_length, rounding_mode="floor")
+                .add(1)
+                .long()
+            )
+
+    def forward(self, input_signal: Tensor, length: Tensor) -> Tuple[Tensor, Tensor]:
+        """
+        Extract Log-mel spectrogram features from the input audio signal.
+        """
+        return self.featurizer(input_signal), self.out_len(length)
+
+
+# ==== gigaam/utils.py ====
+
+
+def normalize_raw_text(text: str) -> str:
+    """
+    Script-agnostic normalization: lowercase, collapse whitespace,
+    and keep only alphanumeric characters plus spaces — i.e. drop punctuation/symbols.
+
+    Word-internal apostrophes are preserved (e.g. ``don't`` stays ``don't``).
+    Apostrophe variants are normalized to ASCII ``'``.
+    Standalone quotes are dropped (not between two alphanumerics).
+    Hyphens/dashes (Unicode category Pd) split words (``word-internal`` -> ``word internal``).
+    """
+    text = text.replace("ё", "е").replace("Ё", "Е").lower()
+    for quote in ("’", "‘", "ʻ", "ʼ"):
+        text = text.replace(quote, "'")
+    out = []
+    for i, c in enumerate(text):
+        if c == "'":
+            if (
+                0 < i < len(text) - 1
+                and text[i - 1].isalnum()
+                and text[i + 1].isalnum()
+            ):
+                out.append(c)
+        elif c.isalnum() or c.isspace():
+            out.append(c)
+        elif unicodedata.category(c) == "Pd":
+            out.append(" ")
+    return " ".join("".join(out).split())
+
+
+def onnx_converter(
+    model_name: str,
+    module: torch.nn.Module,
+    out_dir: str,
+    inputs: Optional[Tuple[Tensor, ...]] = None,
+    input_names: Optional[List[str]] = None,
+    output_names: Optional[List[str]] = None,
+    dynamic_axes: Optional[
+        Union[Dict[str, List[int]], Dict[str, Dict[int, str]]]
+    ] = None,
+    opset_version: int = 17,
+    export_dtype: torch.dtype = torch.float32,
+):
+    """
+    Export a submodule to ONNX: casts inputs and ``module`` to ``export_dtype`` for tracing,
+    then restores the module to float32 via ``module.float()`` so the model stays usable.
+    """
+    if inputs is None:
+        inputs = module.input_example()  # type: ignore[operator]
+    if input_names is None:
+        input_names = module.input_names()  # type: ignore[operator]
+    if output_names is None:
+        output_names = module.output_names()  # type: ignore[operator]
+
+    inputs = tuple(
+        x.to(export_dtype) if x.dtype == torch.float32 else x for x in inputs
+    )
+
+    Path(out_dir).mkdir(exist_ok=True, parents=True)
+    out_path = str(Path(out_dir) / f"{model_name}.onnx")
+    with warnings.catch_warnings(), torch.no_grad():
+        warnings.simplefilter("ignore", category=UserWarning)
+        warnings.simplefilter("ignore", category=TracerWarning)
+        torch.onnx.export(
+            module.to(export_dtype),
+            inputs,
+            out_path,
+            input_names=input_names,
+            output_names=output_names,
+            dynamic_axes=dynamic_axes,
+            opset_version=opset_version,
+            dynamo=False,
+        )
+    print(f"Successfully ported onnx {model_name} to {out_path}.")
+    # We force the whole module to float32 to avoid fp16 preprocessing issues
+    module.float()
+
+
+def format_time(seconds: float) -> str:
+    """
+    Formats time in seconds to HH:MM:SS:mm format.
+    """
+    hours = int(seconds // 3600)
+    minutes = int((seconds % 3600) // 60)
+    seconds = seconds % 60
+    full_seconds = int(seconds)
+    milliseconds = int((seconds - full_seconds) * 100)
+
+    if hours > 0:
+        return f"{hours:02}:{minutes:02}:{full_seconds:02}:{milliseconds:02}"
+    return f"{minutes:02}:{full_seconds:02}:{milliseconds:02}"
+
+
+def rtt_half(x: Tensor) -> Tensor:
+    x1, x2 = x[..., : x.shape[-1] // 2], x[..., x.shape[-1] // 2 :]
+    return torch.cat([-x2, x1], dim=x1.ndim - 1)
+
+
+def apply_rotary_pos_emb(
+    q: Tensor, k: Tensor, cos: Tensor, sin: Tensor, offset: int = 0
+) -> Tuple[Tensor, Tensor]:
+    """
+    Applies Rotary Position Embeddings to query and key tensors.
+    """
+    cos, sin = (
+        cos[offset : q.shape[0] + offset, ...],
+        sin[offset : q.shape[0] + offset, ...],
+    )
+    cos = cos.to(dtype=q.dtype)
+    sin = sin.to(dtype=q.dtype)
+    return (q * cos) + (rtt_half(q) * sin), (k * cos) + (rtt_half(k) * sin)
+
+
+def apply_masked_flash_attn(
+    q: Tensor,
+    k: Tensor,
+    v: Tensor,
+    mask: Tensor,
+    h: int,
+    d_k: int,
+) -> Tensor:
+    """
+    Applies Flash Attention with padding masks.
+    """
+
+    try:
+        from einops import rearrange
+        from flash_attn import flash_attn_varlen_func
+        from flash_attn.bert_padding import pad_input, unpad_input
+    except ImportError as err:
+        raise RuntimeError("flash_attn and einops are required") from err
+
+    pad_mask = ~mask[:, 0, :]
+    b, t = pad_mask.shape
+    q = q.view(b, t, h * d_k)
+    k = k.view(b, t, h * d_k)
+    v = v.view(b, t, h * d_k)
+
+    q_unpad, indices_q, _, max_seqlen_q = unpad_input(q, pad_mask)[:4]
+    q_unpad = rearrange(q_unpad, "nnz (h d) -> nnz h d", h=h)
+
+    k_unpad = unpad_input(k, pad_mask)[0]
+    k_unpad = rearrange(k_unpad, "nnz (h d) -> nnz h d", h=h)
+
+    v_unpad = unpad_input(v, pad_mask)[0]
+    v_unpad = rearrange(v_unpad, "nnz (h d) -> nnz h d", h=h)
+
+    lengths_q = pad_mask.sum(1).to(torch.int32).to(q.device)
+    cu_seqlens_q = F.pad(lengths_q.cumsum(0), (1, 0), value=0).to(torch.int32)
+    max_seqlen_q = torch.max(lengths_q)
+
+    output_unpad = flash_attn_varlen_func(
+        q_unpad,
+        k_unpad,
+        v_unpad,
+        cu_seqlens_q,
+        cu_seqlens_q,
+        max_seqlen_q,
+        max_seqlen_q,
+    )
+
+    scores = pad_input(
+        rearrange(output_unpad, "nnz h d -> nnz (h d)"),
+        indices_q,
+        b,
+        t,
+    )
+
+    return scores
+
+
+def download_short_audio() -> str:
+    """Download test audio file if not exists"""
+    audio_file = "example.wav"
+    if not os.path.exists(audio_file):
+        os.system(
+            'wget -O example.wav "https://cdn.chatwm.opensmodel.sberdevices.ru/GigaAM/example.wav"'
+        )
+    assert os.path.exists(audio_file), "Short audio file not found"
+    return audio_file
+
+
+def download_long_audio() -> str:
+    """Download test audio file if not exists"""
+    audio_file = "long_example.wav"
+    if not os.path.exists(audio_file):
+        os.system(
+            'wget -O long_example.wav "https://cdn.chatwm.opensmodel.sberdevices.ru/GigaAM/long_example.wav"'
+        )
+    assert os.path.exists(audio_file), "Long audio file not found"
+    return audio_file
+
+
+class AudioDataset(torch.utils.data.Dataset):
+    """
+    Unified dataset class for training and inference.
+    Supports loading from manifest file or an iterable of audio paths / waveforms.
+    Provides min / max duration filtering, text normalization, and pre-tokenization.
+    """
+
+    def __init__(
+        self,
+        data: Union[str, Iterable[Union[str, np.ndarray, torch.Tensor]]],
+        tokenizer=None,
+        max_duration: Optional[float] = None,
+        min_duration: float = 0.0,
+        raw_text: bool = False,
+        return_tokens: bool = False,
+    ):
+        self.raw_text = raw_text
+        self.return_tokens = return_tokens
+        self.tokenizer = tokenizer
+        self.samples: List[AudioDatasetSample] = []
+
+        if return_tokens and tokenizer is None:
+            raise ValueError("tokenizer is required when return_tokens=True")
+
+        self.encode = self._make_encoder(tokenizer)
+
+        if isinstance(data, str):
+            self._load_manifest(data, min_duration, max_duration)
+        elif isinstance(data, Iterable) and not isinstance(
+            data, (str, bytes, bytearray)
+        ):
+            self._load_iterable(data, min_duration, max_duration)
+        else:
+            raise TypeError(f"Unsupported data type: {type(data)}")
+
+        if not self.samples:
+            raise ValueError("No valid samples found after filtering")
+
+    def _make_encoder(self, tokenizer):
+        if tokenizer is None:
+            return None
+
+        if getattr(tokenizer, "charwise", False):
+            c2i = {c: i for i, c in enumerate(tokenizer.vocab)}
+            return lambda text: [c2i[c] for c in text if c in c2i]
+
+        return tokenizer.model.encode
+
+    def normalize_text(self, text: str) -> str:
+        if not self.raw_text:
+            return text
+
+        text = normalize_raw_text(text)
+
+        if self.tokenizer is not None and getattr(self.tokenizer, "charwise", False):
+            vocab = set(self.tokenizer.vocab)
+            return "".join(c for c in text if c in vocab)
+
+        return text
+
+    @staticmethod
+    def _get_duration(item: Union[str, np.ndarray, Tensor]) -> float:
+        if isinstance(item, str):
+            with sf.SoundFile(item) as f:
+                return f.frames / f.samplerate
+        if isinstance(item, np.ndarray):
+            return len(item) / SAMPLE_RATE
+        if isinstance(item, torch.Tensor):
+            return item.numel() / SAMPLE_RATE
+        raise TypeError(f"Unexpected sample type: {type(item)}")
+
+    def _duration_ok(
+        self, duration: float, min_duration: float, max_duration: Optional[float]
+    ) -> bool:
+        if duration < min_duration:
+            return False
+        if max_duration is not None and duration > max_duration:
+            return False
+        return True
+
+    @staticmethod
+    def _print_filtered(
+        n_total: int, dur_total: float, n_filt: int, dur_filt: float
+    ) -> None:
+        if n_total == 0:
+            return
+        pn = 100.0 * n_filt / n_total
+        pd = 100.0 * dur_filt / dur_total if dur_total > 0 else 0.0
+        h_filt, h_total = dur_filt / 3600.0, dur_total / 3600.0
+        print(
+            f"filtered by duration: {n_filt}/{n_total} samples ({pn:.1f}%), "
+            f"{h_filt:.2f}/{h_total:.2f} h ({pd:.1f}%)"
+        )
+
+    def _append_sample(
+        self,
+        item: Union[str, np.ndarray, Tensor],
+        duration: float,
+        text: Optional[str] = None,
+    ) -> None:
+        norm_text: Optional[str] = None
+        tokens: Optional[List[int]] = None
+        if text is not None:
+            norm_text = self.normalize_text(text.strip())
+            if self.return_tokens:
+                assert self.encode is not None
+                tokens = self.encode(norm_text)
+        self.samples.append(
+            AudioDatasetSample(
+                item=item, duration=duration, text=norm_text, tokens=tokens
+            )
+        )
+
+    def _load_manifest(
+        self, manifest_path: str, min_duration: float, max_duration: Optional[float]
+    ):
+        data_dir = Path(manifest_path).resolve().parent
+        n_total = n_filt = 0
+        dur_total = dur_filt = 0.0
+
+        with open(manifest_path) as f:
+            for row in csv.DictReader(f, delimiter="\t"):
+                duration = float(row["duration"])
+                n_total += 1
+                dur_total += duration
+                if not self._duration_ok(duration, min_duration, max_duration):
+                    n_filt += 1
+                    dur_filt += duration
+                    continue
+
+                pth = Path(row["path"])
+                path = str((pth if pth.is_absolute() else data_dir / pth).resolve())
+                text = row["transcription"] if "transcription" in row else None
+                self._append_sample(path, duration, text=text)
+
+        self._print_filtered(n_total, dur_total, n_filt, dur_filt)
+
+    def _load_iterable(
+        self,
+        data: Iterable[Union[str, np.ndarray, torch.Tensor]],
+        min_duration: float,
+        max_duration: Optional[float],
+    ):
+        n_total = n_filt = 0
+        dur_total = dur_filt = 0.0
+        for item in data:
+            if not isinstance(item, (str, np.ndarray, torch.Tensor)):
+                raise TypeError(f"Unexpected dtype: {type(item)}")
+
+            duration = self._get_duration(item)
+            n_total += 1
+            dur_total += duration
+            if not self._duration_ok(duration, min_duration, max_duration):
+                n_filt += 1
+                dur_filt += duration
+                continue
+
+            self._append_sample(item, duration)
+
+        self._print_filtered(n_total, dur_total, n_filt, dur_filt)
+
+    def __len__(self) -> int:
+        return len(self.samples)
+
+    @staticmethod
+    def _load_audio(item: Union[str, np.ndarray, Tensor]) -> Tensor:
+        if isinstance(item, str):
+            wav, sr = torchaudio.load(item)
+            if wav.shape[0] > 1:
+                wav = wav.mean(dim=0, keepdim=True)
+            wav = wav.squeeze(0)
+            if sr != SAMPLE_RATE:
+                wav = torchaudio.functional.resample(wav, sr, SAMPLE_RATE)
+            return wav
+        if isinstance(item, np.ndarray):
+            return torch.from_numpy(item)
+        if isinstance(item, torch.Tensor):
+            return item
+        raise TypeError(f"Unexpected sample type: {type(item)}")
+
+    def __getitem__(self, idx: int) -> Union[Tensor, Tuple[Tensor, Tensor]]:
+        sample = self.samples[idx]
+        wav = self._load_audio(sample.item)
+
+        if self.return_tokens:
+            assert sample.tokens is not None
+            return wav, torch.tensor(sample.tokens, dtype=torch.long)
+
+        return wav
+
+    @staticmethod
+    def collate(wavs: List[Tensor]) -> Tuple[Tensor, Tensor]:
+        lengths = torch.tensor([len(w) for w in wavs], dtype=torch.long)
+        max_len = int(lengths.max().item())
+
+        batch = torch.zeros(len(wavs), max_len, dtype=wavs[0].dtype)
+        for i, wav in enumerate(wavs):
+            batch[i, : wav.shape[-1]] = wav.squeeze()
+
+        return batch, lengths
+
+    def collate_fn(
+        self, batch: List[Union[Tensor, Tuple[Tensor, Tensor]]]
+    ) -> Union[Tuple[Tensor, Tensor], Tuple[Tensor, Tensor, Tensor, Tensor]]:
+        if not self.return_tokens:
+            return self.collate(cast(List[Tensor], batch))
+
+        wavs, tokens = zip(*cast(List[Tuple[Tensor, Tensor]], batch))
+        wav_pad, wav_lens = self.collate(list(wavs))
+        tok_pad, tok_lens = self.collate(list(tokens))
+
+        return wav_pad, wav_lens, tok_pad, tok_lens
+
+
+# ==== gigaam/timestamps_utils.py ====
+
+
+def compute_frame_shift(audio_length_samples: int, seq_len: int) -> float:
+    """Compute frame shift (seconds per encoder frame)."""
+    return audio_length_samples / SAMPLE_RATE / seq_len
+
+
+def frames_to_words(
+    tokenizer: Tokenizer,
+    token_ids: List[int],
+    token_frames: List[int],
+    frame_shift: float,
+) -> List[Word]:
+    """
+    Convert token-level frame indices to word-level timestamps.
+    Groups tokens into words at word boundaries (space or sentencepiece '▁' prefix).
+    """
+    words: List[Word] = []
+    current_chars: List[str] = []
+    current_frames: List[int] = []
+
+    def commit():
+        if not current_chars:
+            return
+        text = "".join(current_chars).strip()
+        if not text:
+            current_chars.clear()
+            current_frames.clear()
+            return
+        start = current_frames[0] * frame_shift
+        end = (current_frames[-1] + 1) * frame_shift
+        words.append(Word(text=text, start=start, end=end))
+        current_chars.clear()
+        current_frames.clear()
+
+    for token_id, frame in zip(token_ids, token_frames):
+        char = tokenizer.id_to_str(token_id)
+        if char.startswith("▁"):
+            commit()
+            char = char[1:]
+        elif char == " ":
+            commit()
+            continue
+        current_chars.append(char)
+        current_frames.append(frame)
+
+    commit()
+    return words
+
+
+# ==== gigaam/vad_utils.py ====
+
+_PIPELINE = None
+
+
+def resolve_local_segmentation_path(model_id: str) -> str:
+    """
+    Finds the local path to the segmentation model.
+    """
+    try:
+        return snapshot_download(
+            repo_id=model_id,
+            local_files_only=True,
+        )
+    except LocalEntryNotFoundError:
+        pass
+
+    hf_token = os.getenv("HF_TOKEN")
+    if not hf_token:
+        raise RuntimeError(
+            f"Model {model_id} was not found locally, "
+            f"and no HF_TOKEN was provided to download it."
+        )
+
+    return snapshot_download(
+        repo_id=model_id,
+        token=hf_token,
+    )
+
+
+def load_segmentation_model(model_id: str) -> Model:
+    """
+    Loads the segmentation model from a local snapshot.
+    If it doesn’t exist, it first creates (downloads) the snapshot.
+    """
+    from pyannote.audio import Model
+    from pyannote.audio.core.task import Problem, Resolution, Specifications
+
+    local_path = resolve_local_segmentation_path(model_id=model_id)
+
+    with torch.serialization.safe_globals(
+        [
+            TorchVersion,
+            Problem,
+            Specifications,
+            Resolution,
+        ]
+    ):
+        return Model.from_pretrained(local_path)
+
+
+def get_pipeline(
+    device: torch.device, model_id: str = "pyannote/segmentation-3.0"
+) -> Pipeline:
+    """
+    Retrieves a PyAnnote voice activity detection pipeline and moves it to the specified device.
+    The pipeline is loaded only once and reused across subsequent calls.
+    It requires the Hugging Face API token to be set in the HF_TOKEN environment variable.
+    """
+    from pyannote.audio.pipelines import VoiceActivityDetection
+
+    global _PIPELINE
+    if _PIPELINE is not None:
+        return _PIPELINE.to(device)
+
+    model = load_segmentation_model(model_id=model_id)
+
+    _PIPELINE = VoiceActivityDetection(segmentation=model)
+    _PIPELINE.instantiate({"min_duration_on": 0.0, "min_duration_off": 0.0})
+
+    return _PIPELINE.to(device)
+
+
+def segment_audio_file(
+    wav_file: str,
+    sr: int,
+    max_duration: float = 22.0,
+    min_duration: float = 15.0,
+    strict_limit_duration: float = 30.0,
+    new_chunk_threshold: float = 0.2,
+    device: torch.device = torch.device("cpu"),
+) -> Tuple[List[torch.Tensor], List[Tuple[float, float]]]:
+    """
+    Segments an audio waveform into smaller chunks based on speech activity.
+    The segmentation is performed using a PyAnnote voice activity detection pipeline.
+    """
+    from pyannote.core import Annotation
+
+    audio = load_audio(wav_file)
+    pipeline = get_pipeline(device)
+    sad_segments = cast(Annotation, pipeline(wav_file))
+
+    segments: List[torch.Tensor] = []
+    curr_duration = 0.0
+    curr_start = 0.0
+    curr_end = 0.0
+    boundaries: List[Tuple[float, float]] = []
+
+    def _update_segments(curr_start: float, curr_end: float, curr_duration: float):
+        if curr_duration > strict_limit_duration:
+            max_segments = int(curr_duration / strict_limit_duration) + 1
+            segment_duration = curr_duration / max_segments
+            curr_end = curr_start + segment_duration
+            for _ in range(max_segments - 1):
+                segments.append(audio[int(curr_start * sr) : int(curr_end * sr)])
+                boundaries.append((curr_start, curr_end))
+                curr_start = curr_end
+                curr_end += segment_duration
+        segments.append(audio[int(curr_start * sr) : int(curr_end * sr)])
+        boundaries.append((curr_start, curr_end))
+
+    # Concat segments from pipeline into chunks for asr according to max/min duration
+    # Segments longer than strict_limit_duration are split manually
+    for segment in sad_segments.get_timeline().support():
+        start = max(0, segment.start)
+        end = min(audio.shape[0] / sr, segment.end)
+        if curr_duration == 0.0:
+            curr_start = start
+        elif curr_duration > new_chunk_threshold and (
+            curr_duration + (end - curr_end) > max_duration
+            or curr_duration > min_duration
+        ):
+            _update_segments(curr_start, curr_end, curr_duration)
+            curr_start = start
+        curr_end = end
+        curr_duration = curr_end - curr_start
+
+    if curr_duration > new_chunk_threshold:
+        _update_segments(curr_start, curr_end, curr_duration)
+
+    return segments, boundaries
+
+
+# ==== gigaam/decoder.py ====
+
+
+class CTCHead(nn.Module):
+    """
+    CTC Head module for Connectionist Temporal Classification.
+    """
+
+    def __init__(self, feat_in: int, num_classes: int):
+        super().__init__()
+        self.decoder_layers = torch.nn.Sequential(
+            torch.nn.Conv1d(feat_in, num_classes, kernel_size=1)
+        )
+
+    def forward(self, encoder_output: Tensor) -> Tensor:
+        return torch.nn.functional.log_softmax(
+            self.decoder_layers(encoder_output).transpose(1, 2), dim=-1
+        )
+
+
+class RNNTJoint(nn.Module):
+    """
+    RNN-Transducer Joint Network Module.
+    This module combines the outputs of the encoder and the prediction network using
+    a linear transformation followed by ReLU activation and another linear projection.
+    """
+
+    def __init__(
+        self, enc_hidden: int, pred_hidden: int, joint_hidden: int, num_classes: int
+    ):
+        super().__init__()
+        self.enc_hidden = enc_hidden
+        self.pred_hidden = pred_hidden
+        self.pred = nn.Linear(pred_hidden, joint_hidden)
+        self.enc = nn.Linear(enc_hidden, joint_hidden)
+        self.joint_net = nn.Sequential(nn.ReLU(), nn.Linear(joint_hidden, num_classes))
+
+    def joint(self, encoder_out: Tensor, decoder_out: Tensor) -> Tensor:
+        """
+        Combine the encoder and prediction network outputs into a joint representation.
+        """
+        enc = self.enc(encoder_out).unsqueeze(2)
+        pred = self.pred(decoder_out).unsqueeze(1)
+        return self.joint_net(enc + pred).log_softmax(-1)
+
+    def input_example(self, batch_size: int = 8) -> Tuple[Tensor, Tensor]:
+        device = next(self.parameters()).device
+        enc = torch.zeros(batch_size, self.enc_hidden, 1)
+        dec = torch.zeros(batch_size, self.pred_hidden, 1)
+        return enc.float().to(device), dec.float().to(device)
+
+    def input_names(self) -> List[str]:
+        return ["enc", "dec"]
+
+    def output_names(self) -> List[str]:
+        return ["joint"]
+
+    def dynamic_axes(self) -> Dict[str, Dict[int, str]]:
+        return {
+            "enc": {0: "batch_size"},
+            "dec": {0: "batch_size"},
+            "joint": {0: "batch_size"},
+        }
+
+    def forward(self, enc: Tensor, dec: Tensor) -> Tensor:
+        return self.joint(enc.transpose(1, 2), dec.transpose(1, 2))
+
+
+class RNNTDecoder(nn.Module):
+    """
+    RNN-Transducer Decoder Module.
+    This module handles the prediction network part of the RNN-Transducer architecture.
+    """
+
+    def __init__(self, pred_hidden: int, pred_rnn_layers: int, num_classes: int):
+        super().__init__()
+        self.blank_id = num_classes - 1
+        self.pred_hidden = pred_hidden
+        self.embed = nn.Embedding(num_classes, pred_hidden, padding_idx=self.blank_id)
+        self.lstm = nn.LSTM(pred_hidden, pred_hidden, pred_rnn_layers)
+
+    def predict(
+        self,
+        x: Optional[Tensor],
+        state: Optional[Tensor],
+        batch_size: int = 1,
+    ) -> Tuple[Tensor, Tensor]:
+        """
+        Make predictions based on the current input and previous states.
+        If no input is provided, use zeros as the initial input.
+        """
+        if x is not None:
+            emb: Tensor = self.embed(x)
+        else:
+            emb = torch.zeros(
+                (batch_size, 1, self.pred_hidden), device=next(self.parameters()).device
+            )
+        g, hid = self.lstm(emb.transpose(0, 1), state)
+        return g.transpose(0, 1), hid
+
+    def input_example(self, batch_size: int = 8) -> Tuple[Tensor, Tensor, Tensor]:
+        device = next(self.parameters()).device
+        label = torch.zeros(batch_size, 1, dtype=torch.long).to(device)
+        hidden_h = torch.zeros(self.lstm.num_layers, batch_size, self.pred_hidden).to(
+            device
+        )
+        hidden_c = torch.zeros(self.lstm.num_layers, batch_size, self.pred_hidden).to(
+            device
+        )
+        return label, hidden_h, hidden_c
+
+    def input_names(self) -> List[str]:
+        return ["x", "hi", "ci"]
+
+    def output_names(self) -> List[str]:
+        return ["dec", "ho", "co"]
+
+    def dynamic_axes(self) -> Dict[str, Dict[int, str]]:
+        return {
+            "x": {0: "batch_size"},
+            "hi": {1: "batch_size"},
+            "ci": {1: "batch_size"},
+            "dec": {0: "batch_size"},
+            "ho": {1: "batch_size"},
+            "co": {1: "batch_size"},
+        }
+
+    def forward(self, x: Tensor, h: Tensor, c: Tensor) -> Tuple[Tensor, Tensor, Tensor]:
+        """
+        ONNX-specific forward with x, state = (h, c) -> x, h, c.
+        """
+        emb = self.embed(x)
+        g, (h, c) = self.lstm(emb.transpose(0, 1), (h, c))
+        return g.transpose(0, 1), h, c
+
+
+class RNNTHead(nn.Module):
+    """
+    RNN-Transducer Head Module.
+    This module combines the decoder and joint network components of the RNN-Transducer architecture.
+    """
+
+    def __init__(self, decoder: Dict[str, int], joint: Dict[str, int]):
+        super().__init__()
+        self.decoder = RNNTDecoder(**decoder)
+        self.joint = RNNTJoint(**joint)
+
+
+# ==== gigaam/decoding.py ====
+
+
+class Tokenizer:
+    """
+    Tokenizer for converting between text and token IDs.
+    The tokenizer can operate either character-wise or using a pre-trained SentencePiece model.
+    """
+
+    def __init__(self, vocab: List[str], model_path: Optional[str] = None):
+        self.charwise = model_path is None
+        if self.charwise:
+            self.vocab = vocab
+        else:
+            self.model = SentencePieceProcessor()
+            self.model.load(model_path)
+
+    def decode(self, tokens: List[int]) -> str:
+        """
+        Convert a list of token IDs back to a string.
+        """
+        if self.charwise:
+            return "".join(self.vocab[tok] for tok in tokens)
+        return self.model.decode(tokens)
+
+    def __len__(self):
+        """
+        Get the total number of tokens in the vocabulary.
+        """
+        return len(self.vocab) if self.charwise else len(self.model)
+
+    def id_to_str(self, token_id: int) -> str:
+        """
+        Convert a single token ID to its string representation.
+        """
+        if self.charwise:
+            return self.vocab[token_id]
+        return self.model.IdToPiece(token_id)
+
+
+class CTCGreedyDecoding:
+    """
+    Class for performing greedy decoding of CTC outputs.
+    """
+
+    def __init__(self, vocabulary: List[str], model_path: Optional[str] = None):
+        self.tokenizer = Tokenizer(vocabulary, model_path)
+        self.blank_id = len(self.tokenizer)
+
+    @torch.inference_mode()
+    def decode(
+        self,
+        head: "CTCHead",
+        encoded: Tensor,
+        lengths: Tensor,
+    ) -> List[Tuple[str, List[int], List[int]]]:
+        """
+        CTC greedy decode: returns (text, token_ids, token_frames) per sample.
+        Token frames are time indices (0..T-1) where a token is emitted.
+        """
+        log_probs = head(encoder_output=encoded)
+        C = log_probs.shape[-1]
+        assert (
+            C == len(self.tokenizer) + 1
+        ), f"Num classes {C} != len(vocab)+1 {len(self.tokenizer) + 1}"
+        labels = log_probs.argmax(dim=-1)
+
+        B, T = labels.shape
+        device = labels.device
+        lengths = lengths.to(device=device).clamp(min=0, max=T)
+
+        skip_mask = labels != self.blank_id
+        skip_mask[:, 1:] &= labels[:, 1:] != labels[:, :-1]
+
+        time = torch.arange(T, device=device)[None, :]
+        skip_mask &= time < lengths[:, None]
+
+        idx = skip_mask.nonzero(as_tuple=False)
+        batch_idx = idx[:, 0]
+        token_frames_flat = idx[:, 1]
+        token_ids_flat = labels[skip_mask]
+
+        counts = torch.bincount(batch_idx, minlength=B).cpu().tolist()
+        ids_splits = token_ids_flat.cpu().split(counts)
+        fr_splits = token_frames_flat.cpu().split(counts)
+
+        return [
+            (self.tokenizer.decode(ids.tolist()), ids.tolist(), fr.tolist())
+            for ids, fr in zip(ids_splits, fr_splits)
+        ]
+
+
+class RNNTGreedyDecoding:
+    """
+    Class for performing greedy decoding of RNN-T outputs.
+    """
+
+    def __init__(
+        self,
+        vocabulary: List[str],
+        model_path: Optional[str] = None,
+        max_symbols_per_step: int = 10,
+    ):
+        self.tokenizer = Tokenizer(vocabulary, model_path)
+        self.blank_id = len(self.tokenizer)
+        self.max_symbols = max_symbols_per_step
+
+    @staticmethod
+    def _cat_states(states):
+        """Pack per-sample LSTM states into batched (h, c)."""
+        hs = [s[0] for s in states]
+        cs = [s[1] for s in states]
+        return torch.cat(hs, dim=1), torch.cat(cs, dim=1)
+
+    @staticmethod
+    def _split_state(state):
+        """Unpack batched (h, c) into per-sample states."""
+        h, c = state
+        b = h.shape[1]
+        return [(h[:, i : i + 1], c[:, i : i + 1]) for i in range(b)]
+
+    @torch.inference_mode()
+    def decode(
+        self,
+        head: "RNNTHead",
+        encoded: Tensor,
+        enc_len: Tensor,
+    ) -> List[Tuple[str, List[int], List[int]]]:
+        """
+        RNN-T greedy decode: returns (text, token_ids, token_frames) per sample.
+        Token frames are encoder time indices where tokens are emitted.
+        """
+        x = encoded.transpose(1, 2)  # [B, T, D]
+        B, T, _ = x.shape
+        device = x.device
+
+        hyps: List[List[int]] = [[] for _ in range(B)]
+        token_frames: List[List[int]] = [[] for _ in range(B)]
+        last_label: List[Optional[Tensor]] = [None] * B
+        dec_state: List[Optional[Tuple[Tensor, Tensor]]] = [None] * B
+
+        def emit_batch(batch_idx: List[int], t: int, fresh: bool) -> List[int]:
+            """One batched predictor+joint step; returns samples that emitted non-blank."""
+            idx = torch.tensor(batch_idx, device=device, dtype=torch.long)
+            f = x[idx, t : t + 1, :]  # [b, 1, D]
+
+            if fresh:
+                g, hidden = head.decoder.predict(None, None, batch_size=len(batch_idx))
+            else:
+                labels = torch.cat([last_label[i] for i in batch_idx], dim=0)  # [b, 1]
+                state = self._cat_states([dec_state[i] for i in batch_idx])
+                g, hidden = head.decoder.predict(
+                    labels, state, batch_size=len(batch_idx)
+                )
+
+            k = head.joint.joint(f, g)[:, 0, 0, :].argmax(dim=-1)  # [b]
+            emit = k.ne(self.blank_id)
+
+            if not emit.any():
+                return []
+
+            hidden_parts = self._split_state(hidden)
+            out = []
+
+            for p in emit.nonzero(as_tuple=False).squeeze(1).tolist():
+                bi = batch_idx[p]
+                tok = int(k[p])
+
+                hyps[bi].append(tok)
+                token_frames[bi].append(t)
+                last_label[bi] = k[p : p + 1].view(1, 1)
+                dec_state[bi] = hidden_parts[p]
+                out.append(bi)
+
+            return out
+
+        enc_len = enc_len.cpu()
+        for t in range(T):
+            active = (t < enc_len).nonzero(as_tuple=False).squeeze(1).tolist()
+            if not active:
+                break
+
+            for _ in range(self.max_symbols):
+                if not active:
+                    break
+
+                fresh = [i for i in active if dec_state[i] is None]
+                stateful = [i for i in active if dec_state[i] is not None]
+
+                next_active = []
+                if fresh:
+                    next_active.extend(emit_batch(fresh, t, fresh=True))
+                if stateful:
+                    next_active.extend(emit_batch(stateful, t, fresh=False))
+
+                if not next_active:
+                    break
+
+                active = next_active
+
+        return [(self.tokenizer.decode(h), h, tf) for h, tf in zip(hyps, token_frames)]
+
+
+# ==== gigaam/encoder.py ====
+
+try:
+    from flash_attn import flash_attn_func
+
+    IMPORT_FLASH = True
+except Exception as err:
+    IMPORT_FLASH = False
+    IMPORT_FLASH_ERR = err
+
+
+def _conformer_layer_fwd(
+    layer: nn.Module,
+    x: Tensor,
+    pos_emb: Union[Tensor, List[Tensor]],
+    att_mask: Optional[Tensor],
+    pad_mask: Optional[Tensor],
+) -> Tensor:
+    return layer(x=x, pos_emb=pos_emb, att_mask=att_mask, pad_mask=pad_mask)
+
+
+class StridingSubsampling(nn.Module):
+    """
+    Strided Subsampling layer used to reduce the sequence length.
+    """
+
+    def __init__(
+        self,
+        subsampling: str,
+        kernel_size: int,
+        subsampling_factor: int,
+        feat_in: int,
+        feat_out: int,
+        conv_channels: int,
+    ):
+        super().__init__()
+        self.subsampling_type = subsampling
+        assert self.subsampling_type in ["conv1d", "conv2d"]
+        self._sampling_num = int(math.log(subsampling_factor, 2))
+        self._stride = 2
+        self._kernel_size = kernel_size
+        self._padding = (self._kernel_size - 1) // 2
+
+        layers: List[nn.Module] = []
+        in_channels = 1 if self.subsampling_type == "conv2d" else feat_in
+        subs_conv_class = (
+            torch.nn.Conv2d if self.subsampling_type == "conv2d" else torch.nn.Conv1d
+        )
+        for _ in range(self._sampling_num):
+            layers.append(
+                subs_conv_class(
+                    in_channels=in_channels,
+                    out_channels=conv_channels,
+                    kernel_size=self._kernel_size,
+                    stride=self._stride,
+                    padding=self._padding,
+                )
+            )
+            layers.append(nn.ReLU())
+            in_channels = conv_channels
+
+        out_length = self.calc_output_length(torch.tensor(feat_in))
+        if self.subsampling_type == "conv2d":
+            self.out = torch.nn.Linear(conv_channels * int(out_length), feat_out)
+        self.conv = torch.nn.Sequential(*layers)
+
+    def calc_output_length(
+        self, lengths: Tensor, num_stages: Optional[int] = None
+    ) -> Tensor:
+        """
+        Valid length after applying ``num_stages`` strided subsampling conv
+        stages (defaults to all of them, i.e. the full subsampling output).
+        """
+        if num_stages is None:
+            num_stages = self._sampling_num
+        add_pad = 2 * self._padding - self._kernel_size
+        lengths = lengths.to(torch.float)
+        for _ in range(num_stages):
+            lengths = torch.floor((lengths + add_pad) / self._stride + 1.0)
+        return lengths.to(dtype=torch.int)
+
+    def _mask_time(self, x: Tensor, lengths: Tensor) -> Tensor:
+        """
+        Zero out the padded tail along the time axis (dim 2). The subsampling
+        convolutions are strided and have a receptive field wider than the
+        stride, so the padded frames of shorter samples leak into the last
+        valid frames. Left unmasked, the padding is the log-mel floor
+        (``log(1e-9) ~= -20.7``) of zero-padded audio, not zero, so a batched
+        short sample sees a different boundary than the same sample run alone
+        (where conv zero-padding applies instead). Re-zeroing after every conv
+        stage keeps the valid frames of batched inference aligned with the
+        batch-size-1 result.
+        """
+        time = torch.arange(x.size(2), device=x.device)
+        pad = time[None, :] >= lengths[:, None]  # [b, t]
+        pad = pad[:, None]  # add channel dim -> [b, 1, t]
+        if x.dim() == 4:
+            pad = pad[..., None]  # add feature dim for conv2d -> [b, 1, t, 1]
+        return x.masked_fill(pad, 0.0)
+
+    def forward(self, x: Tensor, lengths: Tensor) -> Tuple[Tensor, Tensor]:
+        if self.subsampling_type == "conv2d":
+            x = x.unsqueeze(1)
+        else:
+            x = x.transpose(1, 2)
+
+        cur_len = lengths
+        x = self._mask_time(x, cur_len)
+        for module in self.conv:
+            x = module(x)
+            if isinstance(module, (torch.nn.Conv1d, torch.nn.Conv2d)):
+                cur_len = self.calc_output_length(cur_len, 1)
+                x = self._mask_time(x, cur_len)
+
+        if self.subsampling_type == "conv2d":
+            b, _, t, _ = x.size()
+            x = self.out(x.transpose(1, 2).reshape(b, t, -1))
+        else:
+            x = x.transpose(1, 2)
+        return x, self.calc_output_length(lengths)
+
+
+class MultiHeadAttention(nn.Module, ABC):
+    """
+    Base class of Multi-Head Attention Mechanisms.
+    """
+
+    def __init__(
+        self, n_head: int, n_feat: int, flash_attn=False, torch_sdpa_attn=False
+    ):
+        super().__init__()
+        assert n_feat % n_head == 0
+        self.d_k = n_feat // n_head
+        self.h = n_head
+        self.linear_q = nn.Linear(n_feat, n_feat)
+        self.linear_k = nn.Linear(n_feat, n_feat)
+        self.linear_v = nn.Linear(n_feat, n_feat)
+        self.linear_out = nn.Linear(n_feat, n_feat)
+        self.flash_attn = flash_attn
+        self.torch_sdpa_attn = torch_sdpa_attn
+        if self.flash_attn and not IMPORT_FLASH:
+            raise RuntimeError(
+                f"flash_attn_func was imported with err {IMPORT_FLASH_ERR}. "
+                "Please install flash_attn or use --no_flash flag. "
+                "If you have already done this, "
+                "--force-reinstall flag might be useful"
+            )
+
+    def forward_qkv(
+        self, query: Tensor, key: Tensor, value: Tensor
+    ) -> Tuple[Tensor, Tensor, Tensor]:
+        """
+        Projects the inputs into queries, keys, and values for multi-head attention.
+        """
+        b = query.size(0)
+        q = self.linear_q(query).view(b, -1, self.h, self.d_k)
+        k = self.linear_k(key).view(b, -1, self.h, self.d_k)
+        v = self.linear_v(value).view(b, -1, self.h, self.d_k)
+        if self.flash_attn:
+            return q, k, v
+        return q.transpose(1, 2), k.transpose(1, 2), v.transpose(1, 2)
+
+    def forward_attention(
+        self, value: Tensor, scores: Tensor, mask: Optional[Tensor]
+    ) -> Tensor:
+        """
+        Computes the scaled dot-product attention given the projected values and scores.
+        """
+        b = value.size(0)
+        if mask is not None:
+            mask = mask.unsqueeze(1)
+            scores = scores.masked_fill(mask, -10000.0)
+            attn = torch.softmax(scores, dim=-1).masked_fill(mask, 0.0)
+        else:
+            attn = torch.softmax(scores, dim=-1)
+        x = torch.matmul(attn, value)
+        x = x.transpose(1, 2).reshape(b, -1, self.h * self.d_k)
+        return self.linear_out(x)
+
+
+class RelPositionMultiHeadAttention(MultiHeadAttention):
+    """
+    Relative Position Multi-Head Attention module.
+    """
+
+    def __init__(self, n_head: int, n_feat: int):
+        super().__init__(n_head, n_feat)
+        self.linear_pos = nn.Linear(n_feat, n_feat, bias=False)
+        self.pos_bias_u = nn.Parameter(torch.FloatTensor(self.h, self.d_k))
+        self.pos_bias_v = nn.Parameter(torch.FloatTensor(self.h, self.d_k))
+
+    def rel_shift(self, x: Tensor) -> Tensor:
+        b, h, qlen, pos_len = x.size()
+        x = torch.nn.functional.pad(x, pad=(1, 0))
+        x = x.view(b, h, -1, qlen)
+        return x[:, :, 1:].view(b, h, qlen, pos_len)
+
+    def forward(
+        self,
+        query: Tensor,
+        key: Tensor,
+        value: Tensor,
+        pos_emb: Tensor,
+        mask: Optional[Tensor] = None,
+    ) -> Tensor:
+        q, k, v = self.forward_qkv(query, key, value)
+        q = q.transpose(1, 2)
+        pos_emb = pos_emb.to(dtype=self.linear_pos.weight.dtype)
+        p = self.linear_pos(pos_emb)
+        p = p.view(pos_emb.shape[0], -1, self.h, self.d_k).transpose(1, 2)
+        q_with_bias_u = (q + self.pos_bias_u).transpose(1, 2)
+        q_with_bias_v = (q + self.pos_bias_v).transpose(1, 2)
+        matrix_bd = torch.matmul(q_with_bias_v, p.transpose(-2, -1))
+        matrix_bd = self.rel_shift(matrix_bd)
+        matrix_ac = torch.matmul(q_with_bias_u, k.transpose(-2, -1))
+        matrix_bd = matrix_bd[:, :, :, : matrix_ac.size(-1)]
+        scores = (matrix_ac + matrix_bd) / math.sqrt(self.d_k)
+        return self.forward_attention(v, scores, mask)
+
+
+class RotaryPositionMultiHeadAttention(MultiHeadAttention):
+    """
+    Rotary Position Multi-Head Attention module.
+    """
+
+    def forward(
+        self,
+        query: Tensor,
+        key: Tensor,
+        value: Tensor,
+        pos_emb: List[Tensor],
+        mask: Optional[Tensor] = None,
+    ) -> Tensor:
+        b, t, _ = value.size()
+        query = query.transpose(0, 1).view(t, b, self.h, self.d_k)
+        key = key.transpose(0, 1).view(t, b, self.h, self.d_k)
+        value = value.transpose(0, 1).view(t, b, self.h, self.d_k)
+
+        cos, sin = pos_emb
+        query, key = apply_rotary_pos_emb(query, key, cos, sin, offset=0)
+
+        q, k, v = self.forward_qkv(
+            query.view(t, b, self.h * self.d_k).transpose(0, 1),
+            key.view(t, b, self.h * self.d_k).transpose(0, 1),
+            value.view(t, b, self.h * self.d_k).transpose(0, 1),
+        )
+
+        if self.flash_attn:
+            if mask is None:
+                scores = flash_attn_func(q, k, v)
+            else:
+                scores = apply_masked_flash_attn(q, k, v, mask, self.h, self.d_k)
+            scores = scores.view(b, -1, self.h * self.d_k)
+            return self.linear_out(scores)
+        elif self.torch_sdpa_attn:
+            attn_mask = None
+            if mask is not None:
+                attn_mask = ~mask.unsqueeze(1)
+                # SDPA masks padding queries with true -inf; softmax over such a row is NaN in forward and backward.
+                # Unmask such rows entirely: their output is finite garbage that nothing reads.
+                attn_mask = attn_mask | (~attn_mask.any(dim=-1, keepdim=True))
+            attn_output = F.scaled_dot_product_attention(q, k, v, attn_mask=attn_mask)
+            attn_output = attn_output.transpose(1, 2).reshape(b, t, self.h * self.d_k)
+            return self.linear_out(attn_output)
+        else:
+            scores = torch.matmul(q, k.transpose(-2, -1) / math.sqrt(self.d_k))
+            return self.forward_attention(v, scores, mask)
+
+
+class PositionalEncoding(nn.Module, ABC):
+    """
+    Base class of Positional Encodings.
+    """
+
+    def __init__(self, dim: int, base: int):
+        super().__init__()
+        self.dim = dim
+        self.base = base
+
+    @abstractmethod
+    def create_pe(self, length: int, device: torch.device) -> Optional[Tensor]:
+        pass
+
+    def extend_pe(self, length: int, device: torch.device):
+        """
+        Extends the positional encoding buffer to process longer sequences.
+        """
+        pe = self.create_pe(length, device)
+        if pe is None:
+            return
+        if hasattr(self, "pe"):
+            self.pe = pe
+        else:
+            self.register_buffer("pe", pe, persistent=False)
+
+
+class RelPositionalEmbedding(PositionalEncoding):
+    """
+    Relative Positional Embedding module.
+    """
+
+    def create_pe(self, length: int, device: torch.device) -> Optional[Tensor]:
+        """
+        Creates the relative positional encoding matrix.
+        """
+        if hasattr(self, "pe") and self.pe.shape[1] >= 2 * length - 1:
+            return None
+        positions = torch.arange(length - 1, -length, -1, device=device).unsqueeze(1)
+        pos_length = positions.size(0)
+        pe = torch.zeros(pos_length, self.dim, device=positions.device)
+        div_term = torch.exp(
+            torch.arange(0, self.dim, 2, device=pe.device)
+            * -(math.log(10000.0) / self.dim)
+        )
+        pe[:, 0::2] = torch.sin(positions * div_term)
+        pe[:, 1::2] = torch.cos(positions * div_term)
+        return pe.unsqueeze(0)
+
+    def forward(self, x: torch.Tensor) -> Tuple[Tensor, Tensor]:
+        input_len = x.size(1)
+        center_pos = self.pe.size(1) // 2 + 1
+        start_pos = center_pos - input_len
+        end_pos = center_pos + input_len - 1
+        return x, self.pe[:, start_pos:end_pos]
+
+
+class RotaryPositionalEmbedding(PositionalEncoding):
+    """
+    Rotary Positional Embedding module.
+    """
+
+    def create_pe(self, length: int, device: torch.device) -> Optional[Tensor]:
+        """
+        Creates or extends the rotary positional encoding matrix.
+        """
+        if hasattr(self, "pe") and self.pe.size(0) >= 2 * length:
+            return None
+        positions = torch.arange(0, length, dtype=torch.float32, device=device)
+        inv_freq = 1.0 / (
+            self.base ** (torch.arange(0, self.dim, 2).float() / self.dim)
+        )
+        t = torch.arange(length, device=positions.device).type_as(inv_freq)
+        freqs = torch.einsum("i,j->ij", t, inv_freq)
+        emb = torch.cat((freqs, freqs), dim=-1).to(positions.device)
+        return torch.cat([emb.cos()[:, None, None, :], emb.sin()[:, None, None, :]])
+
+    def forward(self, x: torch.Tensor) -> Tuple[Tensor, List[Tensor]]:
+        cos_emb = self.pe[0 : x.shape[1]]
+        half_pe = self.pe.shape[0] // 2
+        sin_emb = self.pe[half_pe : half_pe + x.shape[1]]
+        return x, [cos_emb, sin_emb]
+
+
+class ConformerConvolution(nn.Module):
+    """
+    Conformer Convolution module.
+    """
+
+    def __init__(
+        self,
+        d_model: int,
+        kernel_size: int,
+        norm_type: str,
+    ):
+        super().__init__()
+        assert (kernel_size - 1) % 2 == 0
+        assert norm_type in ["batch_norm", "layer_norm"]
+        self.norm_type = norm_type
+        self.pointwise_conv1 = nn.Conv1d(d_model, d_model * 2, kernel_size=1)
+        self.depthwise_conv = nn.Conv1d(
+            in_channels=d_model,
+            out_channels=d_model,
+            kernel_size=kernel_size,
+            padding=(kernel_size - 1) // 2,
+            groups=d_model,
+            bias=True,
+        )
+        self.batch_norm = (
+            nn.BatchNorm1d(d_model)
+            if norm_type == "batch_norm"
+            else nn.LayerNorm(d_model)
+        )
+        self.activation = nn.SiLU()
+        self.pointwise_conv2 = nn.Conv1d(d_model, d_model, kernel_size=1)
+
+    def forward(self, x: Tensor, pad_mask: Optional[Tensor] = None) -> Tensor:
+        x = x.transpose(1, 2)
+        x = self.pointwise_conv1(x)
+        x = nn.functional.glu(x, dim=1)
+        if pad_mask is not None:
+            x = x.masked_fill(pad_mask.unsqueeze(1), 0.0)
+        x = self.depthwise_conv(x)
+        if self.norm_type == "batch_norm":
+            x = self.batch_norm(x)
+        else:
+            x = self.batch_norm(x.transpose(1, 2)).transpose(1, 2)
+        x = self.activation(x)
+        x = self.pointwise_conv2(x)
+        return x.transpose(1, 2)
+
+
+class ConformerFeedForward(nn.Module):
+    """
+    Conformer Feed Forward module.
+    """
+
+    def __init__(self, d_model: int, d_ff: int, use_bias=True):
+        super().__init__()
+        self.linear1 = nn.Linear(d_model, d_ff, bias=use_bias)
+        self.activation = nn.SiLU()
+        self.linear2 = nn.Linear(d_ff, d_model, bias=use_bias)
+
+    def forward(self, x: Tensor) -> Tensor:
+        return self.linear2(self.activation(self.linear1(x)))
+
+
+class ConformerLayer(nn.Module):
+    """
+    Conformer Layer module.
+    This module combines several submodules including feed forward networks,
+    depthwise separable convolution, and multi-head self-attention
+    to form a single Conformer block.
+    """
+
+    def __init__(
+        self,
+        d_model: int,
+        d_ff: int,
+        self_attention_model: str,
+        n_heads: int = 16,
+        conv_norm_type: str = "batch_norm",
+        conv_kernel_size: int = 31,
+        flash_attn: bool = False,
+    ):
+        super().__init__()
+        self.fc_factor = 0.5
+        self.norm_feed_forward1 = nn.LayerNorm(d_model)
+        self.feed_forward1 = ConformerFeedForward(d_model=d_model, d_ff=d_ff)
+        self.norm_conv = nn.LayerNorm(d_model)
+        self.conv = ConformerConvolution(
+            d_model=d_model,
+            kernel_size=conv_kernel_size,
+            norm_type=conv_norm_type,
+        )
+        self.norm_self_att = nn.LayerNorm(d_model)
+        if self_attention_model == "rotary":
+            self.self_attn: nn.Module = RotaryPositionMultiHeadAttention(
+                n_head=n_heads,
+                n_feat=d_model,
+                flash_attn=flash_attn,
+                torch_sdpa_attn=not flash_attn,
+            )
+        else:
+            assert not flash_attn, "Not supported flash_attn for rel_pos"
+            self.self_attn = RelPositionMultiHeadAttention(
+                n_head=n_heads,
+                n_feat=d_model,
+            )
+        self.norm_feed_forward2 = nn.LayerNorm(d_model)
+        self.feed_forward2 = ConformerFeedForward(d_model=d_model, d_ff=d_ff)
+        self.norm_out = nn.LayerNorm(d_model)
+
+    def forward(
+        self,
+        x: Tensor,
+        pos_emb: Union[Tensor, List[Tensor]],
+        att_mask: Optional[Tensor] = None,
+        pad_mask: Optional[Tensor] = None,
+    ) -> Tensor:
+        residual = x
+        x = self.norm_feed_forward1(x)
+        x = self.feed_forward1(x)
+        residual = residual + x * self.fc_factor
+
+        x = self.norm_self_att(residual)
+        x = self.self_attn(x, x, x, pos_emb, mask=att_mask)
+        residual = residual + x
+
+        x = self.norm_conv(residual)
+        x = self.conv(x, pad_mask=pad_mask)
+        residual = residual + x
+
+        x = self.norm_feed_forward2(residual)
+        x = self.feed_forward2(x)
+        residual = residual + x * self.fc_factor
+
+        x = self.norm_out(residual)
+        return x
+
+
+class ConformerEncoder(nn.Module):
+    """
+    Conformer Encoder module.
+    This module encapsulates the entire Conformer encoder architecture,
+    consisting of a StridingSubsampling layer, positional embeddings, and
+    a stack of Conformer Layers.
+    It serves as the main component responsible for processing speech features.
+    """
+
+    def __init__(
+        self,
+        feat_in: int = 64,
+        n_layers: int = 16,
+        d_model: int = 768,
+        subsampling: str = "conv2d",
+        subs_kernel_size: int = 3,
+        subsampling_factor: int = 4,
+        ff_expansion_factor: int = 4,
+        self_attention_model: str = "rotary",
+        n_heads: int = 16,
+        pos_emb_max_len: int = 5000,
+        conv_norm_type: str = "batch_norm",
+        conv_kernel_size: int = 31,
+        flash_attn: bool = False,
+        activation_checkpointing: bool = False,
+    ):
+        super().__init__()
+        self.feat_in = feat_in
+        self.activation_checkpointing = activation_checkpointing
+        assert self_attention_model in [
+            "rotary",
+            "rel_pos",
+        ], f"Not supported attn = {self_attention_model}"
+
+        self.pre_encode = StridingSubsampling(
+            subsampling=subsampling,
+            kernel_size=subs_kernel_size,
+            subsampling_factor=subsampling_factor,
+            feat_in=feat_in,
+            feat_out=d_model,
+            conv_channels=d_model,
+        )
+
+        self.pos_emb_max_len = pos_emb_max_len
+        if self_attention_model == "rotary":
+            self.pos_enc: PositionalEncoding = RotaryPositionalEmbedding(
+                d_model // n_heads, pos_emb_max_len
+            )
+        else:
+            self.pos_enc = RelPositionalEmbedding(d_model, pos_emb_max_len)
+
+        self.layers = nn.ModuleList()
+        for _ in range(n_layers):
+            layer = ConformerLayer(
+                d_model=d_model,
+                d_ff=d_model * ff_expansion_factor,
+                self_attention_model=self_attention_model,
+                n_heads=n_heads,
+                conv_norm_type=conv_norm_type,
+                conv_kernel_size=conv_kernel_size,
+                flash_attn=flash_attn,
+            )
+            self.layers.append(layer)
+
+    def input_example(
+        self,
+        batch_size: int = 8,
+        seqlen: int = 200,
+    ) -> Tuple[Tensor, Tensor]:
+        device = next(self.parameters()).device
+        features = torch.randn(batch_size, self.feat_in, seqlen)
+        feature_lengths = torch.randint(1, seqlen + 1, (batch_size,))
+        feature_lengths[0] = seqlen
+        return features.float().to(device), feature_lengths.to(device)
+
+    def input_names(self) -> List[str]:
+        return ["audio_signal", "length"]
+
+    def output_names(self) -> List[str]:
+        return ["encoded", "encoded_len"]
+
+    @contextmanager
+    def onnx_export_mode(self):
+        saved = []
+        for layer in self.layers:
+            attn = layer.self_attn
+            saved.append((attn.flash_attn, attn.torch_sdpa_attn))
+            attn.flash_attn = False
+            attn.torch_sdpa_attn = False
+        try:
+            yield
+        finally:
+            for layer, (fa, sdpa) in zip(self.layers, saved):
+                layer.self_attn.flash_attn = fa
+                layer.self_attn.torch_sdpa_attn = sdpa
+
+    def dynamic_axes(self) -> Dict[str, Dict[int, str]]:
+        return {
+            "audio_signal": {0: "batch_size", 2: "seq_len"},
+            "length": {0: "batch_size"},
+            "encoded": {0: "batch_size", 1: "seq_len"},
+            "encoded_len": {0: "batch_size"},
+        }
+
+    def forward(self, audio_signal: Tensor, length: Tensor) -> Tuple[Tensor, Tensor]:
+        if not hasattr(self.pos_enc, "pe"):
+            self.pos_enc.extend_pe(self.pos_emb_max_len, audio_signal.device)
+
+        audio_signal, length = self.pre_encode(
+            x=audio_signal.transpose(1, 2), lengths=length
+        )
+
+        max_len = audio_signal.size(1)
+        audio_signal, pos_emb = self.pos_enc(x=audio_signal)
+
+        pad_mask = torch.arange(0, max_len, device=audio_signal.device).expand(
+            length.size(0), -1
+        ) < length.unsqueeze(-1)
+
+        att_mask = None
+        if audio_signal.shape[0] > 1:
+            att_mask = pad_mask.unsqueeze(1).repeat([1, max_len, 1])
+            att_mask = torch.logical_and(att_mask, att_mask.transpose(1, 2))
+            att_mask = ~att_mask
+
+        pad_mask = ~pad_mask
+
+        for layer in self.layers:
+            if self.activation_checkpointing and self.training:
+                audio_signal = checkpoint(
+                    _conformer_layer_fwd,
+                    layer,
+                    audio_signal,
+                    pos_emb,
+                    att_mask,
+                    pad_mask,
+                    use_reentrant=False,
+                )
+            else:
+                audio_signal = layer(
+                    x=audio_signal,
+                    pos_emb=pos_emb,
+                    att_mask=att_mask,
+                    pad_mask=pad_mask,
+                )
+
+        return audio_signal.transpose(1, 2), length
+
+
+# ==== gigaam/model.py ====
+
+LONGFORM_THRESHOLD = 25 * SAMPLE_RATE
+
+
+class GigaAM(nn.Module):
+    """
+    Giga Acoustic Model: Self-Supervised Model for Speech Tasks
+    """
+
+    def __init__(self, cfg: omegaconf.DictConfig):
+        super().__init__()
+        self.cfg = cfg
+        self.preprocessor = hydra.utils.instantiate(self.cfg.preprocessor)
+        self.encoder = hydra.utils.instantiate(self.cfg.encoder)
+
+    def forward(
+        self, features: Tensor, feature_lengths: Tensor
+    ) -> Tuple[Tensor, Tensor]:
+        """
+        Perform forward pass through the preprocessor and encoder.
+        """
+        features, feature_lengths = self.preprocessor(features, feature_lengths)
+        if self._device.type == "cpu":
+            return self.encoder(features, feature_lengths)
+        with torch.autocast(device_type=self._device.type, dtype=torch.float16):
+            return self.encoder(features, feature_lengths)
+
+    @property
+    def _device(self) -> torch.device:
+        return next(self.parameters()).device
+
+    @property
+    def _dtype(self) -> torch.dtype:
+        return next(self.parameters()).dtype
+
+    def prepare_wav(self, wav_file: str) -> Tuple[Tensor, Tensor]:
+        """
+        Prepare an audio file for processing by loading it onto
+        the correct device and converting its format.
+        """
+        wav = load_audio(wav_file)
+        wav = wav.to(self._device).to(self._dtype).unsqueeze(0)
+        length = torch.full([1], wav.shape[-1], device=self._device)
+        return wav, length
+
+    def embed_audio(self, wav_file: str) -> Tuple[Tensor, Tensor]:
+        """
+        Extract audio representations using the GigaAM model.
+        """
+        wav, length = self.prepare_wav(wav_file)
+        encoded, encoded_len = self.forward(wav, length)
+        return encoded, encoded_len
+
+    def to_onnx(self, dir_path: str = ".", dtype: torch.dtype = torch.float32) -> None:
+        """
+        Export onnx model encoder to the specified dir.
+        """
+        with self.encoder.onnx_export_mode():
+            self._to_onnx(dir_path, dtype=dtype)
+        omegaconf.OmegaConf.save(self.cfg, f"{dir_path}/{self.cfg.model_name}.yaml")
+
+    def _to_onnx(self, dir_path: str = ".", dtype: torch.dtype = torch.float32) -> None:
+        """
+        Export onnx model encoder to the specified dir.
+        """
+        onnx_converter(
+            model_name=f"{self.cfg.model_name}_encoder",
+            out_dir=dir_path,
+            module=self.encoder,
+            dynamic_axes=self.encoder.dynamic_axes(),
+            export_dtype=dtype,
+        )
+
+
+class GigaAMASR(GigaAM):
+    """
+    Giga Acoustic Model for Speech Recognition
+    """
+
+    def __init__(self, cfg: omegaconf.DictConfig):
+        super().__init__(cfg)
+        self.head = hydra.utils.instantiate(self.cfg.head)
+        self.decoding = hydra.utils.instantiate(self.cfg.decoding)
+
+    def _decode(
+        self,
+        encoded: Tensor,
+        encoded_len: Tensor,
+        wav_lens: Tensor,
+        word_timestamps: bool = False,
+    ) -> List[Tuple[str, Optional[List[Word]]]]:
+        decoded = self.decoding.decode(self.head, encoded, encoded_len)
+        if not word_timestamps:
+            return [(t, None) for t, _, _ in decoded]
+
+        out: List[Tuple[str, Optional[List[Word]]]] = []
+        for i, (text, token_ids, token_frames) in enumerate(decoded):
+            frame_shift = compute_frame_shift(
+                int(wav_lens[i].item()), int(encoded_len[i].item())
+            )
+            out.append(
+                (
+                    text,
+                    frames_to_words(
+                        self.decoding.tokenizer,
+                        token_ids,
+                        token_frames,
+                        frame_shift,
+                    ),
+                )
+            )
+        return out
+
+    @torch.inference_mode()
+    def transcribe(
+        self, wav_file: str, word_timestamps: bool = False
+    ) -> TranscriptionResult:
+        """
+        Transcribes a short audio file into text.
+        Returns TranscriptionResult with optional word-level timestamps.
+        """
+        wav, length = self.prepare_wav(wav_file)
+        if length.item() > LONGFORM_THRESHOLD:
+            raise ValueError("Too long wav file, use 'transcribe_longform' method.")
+
+        encoded, encoded_len = self.forward(wav, length)
+        text, words = self._decode(encoded, encoded_len, length, word_timestamps)[0]
+        return TranscriptionResult(text=text, words=words)
+
+    def forward_for_export(
+        self, features: Tensor, feature_lengths: Tensor
+    ) -> Tuple[Tensor, Tensor]:
+        """
+        Encoder-decoder forward to save model entirely in onnx format.
+        """
+        encoded, encoded_len = self.encoder(features, feature_lengths)
+        return self.head(encoded), encoded_len
+
+    def _to_onnx(self, dir_path: str = ".", dtype: torch.dtype = torch.float32) -> None:
+        """
+        Export onnx ASR model.
+        `ctc`:  exported entirely in encoder-decoder format.
+        `rnnt`: exported in encoder/decoder/joint parts separately.
+        """
+        if "ctc" in self.cfg.model_name:
+            saved_forward = self.forward
+            self.forward = self.forward_for_export  # type: ignore[assignment, method-assign]
+            try:
+                onnx_converter(
+                    model_name=self.cfg.model_name,
+                    out_dir=dir_path,
+                    module=self,
+                    inputs=self.encoder.input_example(),
+                    input_names=["features", "feature_lengths"],
+                    output_names=["log_probs", "encoded_lengths"],
+                    dynamic_axes={
+                        "features": {0: "batch_size", 2: "seq_len"},
+                        "feature_lengths": {0: "batch_size"},
+                        "log_probs": {0: "batch_size", 1: "seq_len"},
+                        "encoded_lengths": {0: "batch_size"},
+                    },
+                    export_dtype=dtype,
+                )
+            finally:
+                self.forward = saved_forward  # type: ignore[assignment, method-assign]
+        else:
+            super()._to_onnx(dir_path, dtype=dtype)
+            onnx_converter(
+                model_name=f"{self.cfg.model_name}_decoder",
+                out_dir=dir_path,
+                module=self.head.decoder,
+                dynamic_axes=self.head.decoder.dynamic_axes(),
+                export_dtype=dtype,
+            )
+            onnx_converter(
+                model_name=f"{self.cfg.model_name}_joint",
+                out_dir=dir_path,
+                module=self.head.joint,
+                dynamic_axes=self.head.joint.dynamic_axes(),
+                export_dtype=dtype,
+            )
+
+    @torch.inference_mode()
+    def transcribe_longform(
+        self,
+        wav_file: str,
+        word_timestamps: bool = False,
+        fr_batch_size: int = 16,
+        fr_num_workers: int = 0,
+        **kwargs,
+    ) -> LongformTranscriptionResult:
+        """
+        Transcribes a long audio file by splitting it into segments and
+        then transcribing each segment (batched inference via AudioDataset).
+        Use fr_batch_size and fr_num_workers to control the batched inference.
+        Returns LongformTranscriptionResult with segments containing optional word-level timestamps.
+        """
+
+        segments, boundaries = segment_audio_file(
+            wav_file, SAMPLE_RATE, device=self._device, **kwargs
+        )
+
+        if not segments:
+            return LongformTranscriptionResult(segments=[])
+
+        ds = AudioDataset(segments, tokenizer=None)
+        dl = DataLoader(
+            ds,
+            batch_size=fr_batch_size,
+            shuffle=False,
+            collate_fn=AudioDataset.collate,
+            num_workers=fr_num_workers,
+        )
+
+        result_segments: List[Segment] = []
+        idx = 0
+        for wav_pad, wav_lens in dl:
+            wav_pad = wav_pad.to(self._device).to(self._dtype)
+            wav_lens = wav_lens.to(self._device)
+            encoded, encoded_len = self.forward(wav_pad, wav_lens)
+            for text, words in self._decode(
+                encoded, encoded_len, wav_lens, word_timestamps
+            ):
+                seg_start, seg_end = boundaries[idx]
+                idx += 1
+                if word_timestamps:
+                    result_segments.append(
+                        Segment(
+                            text=text,
+                            start=seg_start,
+                            end=seg_end,
+                            words=[
+                                Word(
+                                    text=w.text,
+                                    start=round(w.start + seg_start, 3),
+                                    end=round(w.end + seg_start, 3),
+                                )
+                                for w in words or []
+                            ],
+                        )
+                    )
+                else:
+                    result_segments.append(
+                        Segment(text=text, start=seg_start, end=seg_end)
+                    )
+        return LongformTranscriptionResult(segments=result_segments)
+
+
+class GigaAMEmo(GigaAM):
+    """
+    Giga Acoustic Model for Emotion Recognition
+    """
+
+    def __init__(self, cfg: omegaconf.DictConfig):
+        super().__init__(cfg)
+        self.head = hydra.utils.instantiate(self.cfg.head)
+        self.id2name = cfg.id2name
+
+    def get_probs(self, wav_file: str) -> Dict[str, float]:
+        """
+        Calculate probabilities for each emotion class based on the provided audio file.
+        """
+        wav, length = self.prepare_wav(wav_file)
+        encoded, _ = self.forward(wav, length)
+        encoded_pooled = nn.functional.avg_pool1d(
+            encoded, kernel_size=encoded.shape[-1]
+        ).squeeze(-1)
+
+        logits = self.head(encoded_pooled)[0]
+        probs = nn.functional.softmax(logits, dim=-1).detach().tolist()
+
+        return {self.id2name[i]: probs[i] for i in range(len(self.id2name))}
+
+    def forward_for_export(self, features: Tensor, feature_lengths: Tensor) -> Tensor:
+        """
+        Encoder-decoder forward to save model entirely in onnx format.
+        """
+        encoded, _ = self.encoder(features, feature_lengths)
+        enc_pooled = encoded.mean(dim=-1)
+        return nn.functional.softmax(self.head(enc_pooled), dim=-1)
+
+    def _to_onnx(self, dir_path: str = ".", dtype: torch.dtype = torch.float32) -> None:
+        """
+        Export onnx Emo model.
+        """
+        saved_forward = self.forward
+        self.forward = self.forward_for_export  # type: ignore[assignment, method-assign]
+        try:
+            onnx_converter(
+                model_name=self.cfg.model_name,
+                out_dir=dir_path,
+                module=self,
+                inputs=self.encoder.input_example(),
+                input_names=["features", "feature_lengths"],
+                output_names=["probs"],
+                dynamic_axes={
+                    "features": {0: "batch_size", 2: "seq_len"},
+                    "feature_lengths": {0: "batch_size"},
+                    "probs": {0: "batch_size", 1: "seq_len"},
+                },
+                export_dtype=dtype,
+            )
+        finally:
+            self.forward = saved_forward  # type: ignore[assignment, method-assign]
+
+
+# ==== HF glue ====
+
+
+class GigaAMConfig(PretrainedConfig):
+    model_type = "gigaam"
+
+    def __init__(self, cfg: omegaconf.DictConfig = None, **kwargs):
+        super().__init__(**kwargs)
+        self.cfg = cfg
+
+
+class GigaAMModel(PreTrainedModel):
+    config_class = GigaAMConfig
+    base_model_prefix = "gigaam"
+
+    def __init__(self, config: GigaAMConfig):
+        super().__init__(config)
+        self.config = config
+        inner = self.config.cfg["model"]["cfg"]
+        if "decoding" in inner and "model_path" in inner["decoding"]:
+            inner["decoding"]["model_path"] = cached_file(
+                config.name_or_path,
+                "tokenizer.model",
+                revision=getattr(config, "_commit_hash", None),
+                cache_dir=getattr(config, "cache_dir", None),
+                token=getattr(config, "token", None),
+            )
+        with torch.device("cpu"):  # transformers>=5 inits under meta device
+            self.model = instantiate(config.cfg["model"], _recursive_=False)
+        self.post_init()
+
+    def forward(self, features: torch.Tensor, feature_lengths: torch.Tensor):
+        return self.model(features, feature_lengths)
+
+    def embed_audio(self, wav_file: str) -> torch.Tensor:
+        return self.model.embed_audio(wav_file)
+
+    def transcribe(
+        self, wav_file: str, word_timestamps: bool = False
+    ) -> TranscriptionResult:
+        return self.model.transcribe(wav_file, word_timestamps=word_timestamps)
+
+    def transcribe_longform(
+        self, wav_file: str, **kwargs
+    ) -> LongformTranscriptionResult:
+        return self.model.transcribe_longform(wav_file, **kwargs)
+
+    def get_probs(self, wav_file: str) -> Dict[str, float]:
+        return self.model.get_probs(wav_file)
+
+    @torch.no_grad()
+    def to_onnx(self, dir_path: str = ".") -> None:
+        self.model.to_onnx(dir_path)

pytorch_model.bin

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:e1db43873ec5e296f229572e06e2470fc157ac9f8d4aacabda295630b9b91728
+size 883170115