From a755fec6f962c882bdde563e04b020e070bc052b Mon Sep 17 00:00:00 2001 From: Alexander4127 Date: Wed, 15 Apr 2026 15:47:27 +0000 Subject: [PATCH] Fine-tuning, batched rnnt, fp16 onnx --- .github/workflows/gigaam.yml | 29 +- .gitignore | 4 +- README.md | 17 +- README_ru.md | 9 +- colab_example.ipynb | 6 +- gigaam/__init__.py | 23 +- gigaam/decoder.py | 39 +- gigaam/decoding.py | 161 ++- gigaam/encoder.py | 37 +- gigaam/model.py | 159 ++- gigaam/onnx_utils.py | 371 ++++-- gigaam/types.py | 13 +- gigaam/utils.py | 261 +++- pyproject.toml | 25 +- tests/test_onnx.py | 56 +- tests/test_reading.py | 16 +- tests/test_training.py | 222 ++++ train_utils/README.md | 147 +++ train_utils/eval.py | 96 ++ train_utils/example.ipynb | 1093 +++++++++++++++++ train_utils/module.py | 271 ++++ train_utils/train.py | 212 ++++ train_utils/utils.py | 230 ++++ triton_scripts/Dockerfile | 7 +- triton_scripts/README.md | 16 +- .../repos/ctc_encoder_onnx/config.pbtxt | 2 +- .../repos/ctc_encoder_trt/config.pbtxt | 2 +- .../repos/ctc_postprocessing/1/model.py | 16 +- .../repos/gigaam_encoder_onnx/config.pbtxt | 4 +- .../repos/gigaam_encoder_trt/config.pbtxt | 4 +- triton_scripts/repos/preprocessing/1/model.py | 2 +- .../repos/preprocessing/config.pbtxt | 3 +- .../repos/rnnt_postprocessing/1/model.py | 22 +- .../repos/rnnt_postprocessing/config.pbtxt | 2 +- triton_scripts/run_convert_onnx.py | 54 +- 35 files changed, 3223 insertions(+), 408 deletions(-) create mode 100644 tests/test_training.py create mode 100644 train_utils/README.md create mode 100644 train_utils/eval.py create mode 100644 train_utils/example.ipynb create mode 100644 train_utils/module.py create mode 100644 train_utils/train.py create mode 100644 train_utils/utils.py diff --git a/.github/workflows/gigaam.yml b/.github/workflows/gigaam.yml index 966c997..9920df4 100644 --- a/.github/workflows/gigaam.yml +++ b/.github/workflows/gigaam.yml @@ -14,7 +14,7 @@ jobs: test: name: Test (Python ${{ matrix.python-version }}) runs-on: ubuntu-latest - timeout-minutes: 45 + timeout-minutes: 60 strategy: fail-fast: false @@ -56,8 +56,7 @@ jobs: - name: Install Python dependencies run: | python -m pip install --upgrade pip wheel setuptools - pip install --no-cache-dir torch==2.8.0 torchaudio==2.8.0 --index-url https://download.pytorch.org/whl/cpu - pip install --no-cache-dir -e ".[longform,tests]" + pip install --no-cache-dir -e ".[longform,tests,train]" - name: Show disk usage after install run: df -h @@ -88,21 +87,9 @@ jobs: run: | pytest -v tests/test_timestamps.py --tb=short - - name: Run all tests with coverage - if: matrix.python-version == '3.10' - env: - HF_TOKEN: ${{ secrets.HF_TOKEN }} + - name: Run training tests run: | - pytest --cov=gigaam --cov-report=xml --cov-report=term-missing tests/ - - - name: Upload coverage to Codecov - if: matrix.python-version == '3.10' - uses: codecov/codecov-action@v4 - with: - file: ./coverage.xml - flags: unittests - name: codecov-umbrella - fail_ci_if_error: false + pytest -v tests/test_training.py --tb=short lint: name: Lint and Format Check @@ -134,16 +121,16 @@ jobs: - name: Check code formatting with black run: | - black --check --diff gigaam/ tests/ triton_scripts/ + black --check --diff gigaam/ tests/ triton_scripts/ train_utils/*.py - name: Check imports with isort run: | - isort --check-only --diff gigaam/ tests/ triton_scripts/ + isort --check-only --diff gigaam/ tests/ triton_scripts/ train_utils/*.py - name: Lint with flake8 run: | - flake8 --ignore=E203,W503,W504 --max-line-length=120 --statistics gigaam/ tests/ triton_scripts/ + flake8 --ignore=E203,W503,W504 --max-line-length=120 --statistics gigaam/ tests/ triton_scripts/ train_utils/*.py - name: Type check with mypy run: | - mypy gigaam/ --ignore-missing-imports --no-strict-optional + mypy gigaam/ train_utils/*.py --ignore-missing-imports --no-strict-optional diff --git a/.gitignore b/.gitignore index 4775fdd..fa2205f 100644 --- a/.gitignore +++ b/.gitignore @@ -4,4 +4,6 @@ build *.wav .DS_Store *tmp* -onnx \ No newline at end of file +onnx +train_utils/data +train_utils/checkpoints \ No newline at end of file diff --git a/README.md b/README.md index 2cdf132..3c9a9cf 100644 --- a/README.md +++ b/README.md @@ -15,6 +15,7 @@ ![plot](./assets/gigaam_scheme.svg) ## Latest News +* 2026/04 — [model fine-tuning](#model-fine-tuning) (CTC / RNNT), word-level timestamps, [Triton Inference Server](#triton-inference-server-and-tensorrt) * 2025/11 — GigaAM-v3: **30%** WER reduction on new data domains; GigaAM-v3-e2e: end-to-end transcription support (**70:30** win in Side-by-Side vs Whisper-large-v3) * 2025/06 — Our [research paper on GigaAM](https://arxiv.org/abs/2506.01192) was accepted to InterSpeech 2025! * 2024/12 — [MIT License](./LICENSE), GigaAM-v2 (**-15%** and **-12%** WER Reduction for CTC and RNN-T models, respectively), [ONNX export support](#onnx-export-and-inference) @@ -126,6 +127,10 @@ emotion2prob = model.get_probs(audio_path) print(", ".join([f"{emotion}: {prob:.3f}" for emotion, prob in emotion2prob.items()])) ``` +### Model Fine-tuning + +Both CTC and RNNT models can be fine-tuned on custom data using PyTorch Lightning. For a detailed description of all training arguments, see [`train_utils/README.md`](./train_utils/README.md). End-to-end examples with different VRAM constraints are available in [`train_utils/example.ipynb`](./train_utils/example.ipynb). + ### Loading from Hugging Face > **Note:** Install requirements from the [example](./colab_example.ipynb). @@ -138,7 +143,7 @@ model = AutoModel.from_pretrained("ai-sage/GigaAM-v3", revision="e2e_rnnt", trus ### ONNX Export and Inference -> **Note:** GPU support can be enabled with `pip install onnxruntime-gpu==1.23.*` if applicable. +> **Note:** `to_onnx` exports in **fp32** by default. Pass `dtype=torch.float16` for GPU deployment — it is faster and uses less VRAM. GPU support can be enabled with uninstalling onnxruntime and running `pip install onnxruntime-gpu==1.22.*`. 1. Export the model to ONNX using the `model.to_onnx` method: ```python @@ -146,7 +151,7 @@ model = AutoModel.from_pretrained("ai-sage/GigaAM-v3", revision="e2e_rnnt", trus model_version = "v3_ctc" # Options: any version model = gigaam.load_model(model_version) - model.to_onnx(dir_path=onnx_dir) + model.to_onnx(dir_path=onnx_dir, dtype=torch.float32) # or fp16 (recommended for GPU) ``` 2. Run ONNX inference: @@ -154,8 +159,12 @@ model = AutoModel.from_pretrained("ai-sage/GigaAM-v3", revision="e2e_rnnt", trus from gigaam.onnx_utils import load_onnx, infer_onnx sessions, model_cfg = load_onnx(onnx_dir, model_version) - result = infer_onnx(audio_path, model_cfg, sessions) - print(result) # string for ctc / rnnt, np.ndarray for ssl / emo + result = infer_onnx([audio_path], model_cfg, sessions) + print(result[0]) + + # or use the whole dataset + texts = infer_onnx("/path/to/eval/manifest.tsv", model_cfg, sessions) + print(texts[0]) ``` These and more advanced (e.g. custom audio loading, batching) examples can be found in the [Colab notebook](https://colab.research.google.com/github/salute-developers/GigaAM/blob/main/colab_example.ipynb). diff --git a/README_ru.md b/README_ru.md index 846651e..86906e1 100644 --- a/README_ru.md +++ b/README_ru.md @@ -15,6 +15,7 @@ ![plot](./assets/gigaam_scheme.svg) ## Последние обновления +* **2026/04** — [дообучение моделей](#дообучение-моделей) (CTC / RNNT), таймстемпы на уровне слов, [Triton Inference Server](#triton-inference-server-и-tensorrt) * **2025/11** — GigaAM-v3: снижение WER на **30%** на новых доменах данных; GigaAM-v3-e2e: end-to-end распознавание речи (**70:30** в side-by-side сравнении против Whisper-large-v3) * **2025/06** — Наша [научная статья о GigaAM](https://arxiv.org/abs/2506.01192) принята на InterSpeech 2025! * **2024/12** — [MIT-лицензия](./LICENSE), GigaAM-v2 (**снижение WER на 15% и 12%** для CTC и RNN-T моделей), [поддержка экспорта в ONNX](#конвертация-в-onnx-и-использование-графа) @@ -125,6 +126,10 @@ emotion2prob = model.get_probs(audio_path) print(", ".join([f"{emotion}: {prob:.3f}" for emotion, prob in emotion2prob.items()])) ``` +### Дообучение моделей + +CTC и RNNT модели можно дообучать на собственных данных с помощью PyTorch Lightning. Подробное описание всех аргументов обучения — в [`train_utils/README.md`](./train_utils/README.md). Примеры с разными ограничениями VRAM доступны в [`train_utils/example.ipynb`](./train_utils/example.ipynb). + ### Загрузка из Hugging Face > Используйте установку зависимостей из [примера](./colab_example.ipynb). @@ -137,7 +142,7 @@ model = AutoModel.from_pretrained("ai-sage/GigaAM-v3", revision="e2e_rnnt", trus ### Конвертация в ONNX и использование графа -> GPU будет использоваться после установки `pip install onnxruntime-gpu==1.23.*` (если доступно). +> **Примечание:** `to_onnx` по умолчанию экспортирует в **fp32**. Для GPU рекомендуется передать `dtype=torch.float16` — это ускоряет инференс и снижает потребление VRAM. GPU будет использоваться после удаления onnxruntime и установки `pip install onnxruntime-gpu==1.22.*`. 1. Экспорт модели в ONNX с помощью метода `model.to_onnx`: ```python @@ -145,7 +150,7 @@ model = AutoModel.from_pretrained("ai-sage/GigaAM-v3", revision="e2e_rnnt", trus model_version = "v3_ctc" # Варианты: любая версия модели model = gigaam.load_model(model_version) - model.to_onnx(dir_path=onnx_dir) + model.to_onnx(dir_path=onnx_dir, dtype=torch.float32) # или fp16 (рекомендовано для GPU) ``` 2. Запуск с использованием ONNX: diff --git a/colab_example.ipynb b/colab_example.ipynb index b567e36..ead522a 100644 --- a/colab_example.ipynb +++ b/colab_example.ipynb @@ -426,8 +426,8 @@ " wav_tns.to(model._device).to(model._dtype), lengths.to(model._device)\n", " )\n", " results = model.decoding.decode(model.head, encoded, encoded_len)\n", - " for token_ids, _ in results:\n", - " print(model.decoding.tokenizer.decode(token_ids))\n", + " for text, _, __ in results:\n", + " print(text)\n", "\n", "# outputs expected to be equal" ] @@ -514,7 +514,7 @@ " with torch.no_grad():\n", " encoded, encoded_len = model(wav_tns, lengths)\n", " results = model.decoding.decode(model.head, encoded, encoded_len)\n", - " pred_texts.extend(model.decoding.tokenizer.decode(ids) for ids, _ in results)\n", + " pred_texts.extend(text for text, _, __ in results)\n", "\n", "for (start, end), text in zip(boundaries, pred_texts):\n", " print(f\"[{gigaam.format_time(start)} - {gigaam.format_time(end)}]: {text}\")" diff --git a/gigaam/__init__.py b/gigaam/__init__.py index d43a8fc..89fc069 100644 --- a/gigaam/__init__.py +++ b/gigaam/__init__.py @@ -115,12 +115,12 @@ def load_model( download_root: Optional[str] = None, ) -> Union[GigaAM, GigaAMEmo, GigaAMASR]: """ - Load the GigaAM model by name. + Load the GigaAM model by name, or a local ``.ckpt`` from fine-tuning with ``train_utils/train.py``. Parameters ---------- model_name : str - The name of the model to load. + Model name or a path to a ``.ckpt`` file. fp16_encoder: Whether to convert encoder weights to FP16 precision. use_flash : Optional[bool] @@ -136,6 +136,25 @@ def load_model( if download_root is None: download_root = _CACHE_DIR + local_path = os.path.expanduser(model_name) + if os.path.isfile(local_path): + finetuned = torch.load(local_path, map_location="cpu", weights_only=False) + base_name = finetuned["hyper_parameters"]["model_name"] + model = load_model( + base_name, + fp16_encoder=fp16_encoder, + use_flash=use_flash, + device=device_obj, + download_root=download_root, + ) + sd = { + k: v + for k, v in finetuned["state_dict"].items() + if k.startswith(("preprocessor.", "encoder.", "head.")) + } + model.load_state_dict(sd) + return model + model_name, model_path = _download_model(model_name, download_root) tokenizer_path = _download_tokenizer(model_name, download_root) diff --git a/gigaam/decoder.py b/gigaam/decoder.py index 066d397..ebcb49c 100644 --- a/gigaam/decoder.py +++ b/gigaam/decoder.py @@ -46,10 +46,10 @@ def joint(self, encoder_out: Tensor, decoder_out: Tensor) -> Tensor: pred = self.pred(decoder_out).unsqueeze(1) return self.joint_net(enc + pred).log_softmax(-1) - def input_example(self) -> Tuple[Tensor, Tensor]: + def input_example(self, batch_size: int = 8) -> Tuple[Tensor, Tensor]: device = next(self.parameters()).device - enc = torch.zeros(1, self.enc_hidden, 1) - dec = torch.zeros(1, self.pred_hidden, 1) + 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]: @@ -58,6 +58,13 @@ def input_names(self) -> List[str]: 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)) @@ -94,18 +101,32 @@ def predict( g, hid = self.lstm(emb.transpose(0, 1), state) return g.transpose(0, 1), hid - def input_example(self) -> Tuple[Tensor, Tensor, Tensor]: + def input_example(self, batch_size: int = 8) -> Tuple[Tensor, Tensor, Tensor]: device = next(self.parameters()).device - label = torch.tensor([[0]]).to(device) - hidden_h = torch.zeros(1, 1, self.pred_hidden).to(device) - hidden_c = torch.zeros(1, 1, self.pred_hidden).to(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", "h", "c"] + return ["x", "hi", "ci"] def output_names(self) -> List[str]: - return ["dec", "h", "c"] + 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]: """ diff --git a/gigaam/decoding.py b/gigaam/decoding.py index aae4b9f..e703646 100644 --- a/gigaam/decoding.py +++ b/gigaam/decoding.py @@ -59,20 +59,17 @@ def decode( head: "CTCHead", encoded: Tensor, lengths: Tensor, - labels: Optional[Tensor] = None, - ) -> List[Tuple[List[int], List[int]]]: + ) -> List[Tuple[str, List[int], List[int]]]: """ - CTC greedy decode: returns (token_ids, token_frames) per sample. + CTC greedy decode: returns (text, token_ids, token_frames) per sample. Token frames are time indices (0..T-1) where a token is emitted. - If labels are provided, encoded and head are not used. """ - if labels is None: - 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) + 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 @@ -93,10 +90,17 @@ def decode( ids_splits = token_ids_flat.cpu().split(counts) fr_splits = token_frames_flat.cpu().split(counts) - return [(ids.tolist(), fr.tolist()) for ids, fr in zip(ids_splits, fr_splits)] + 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], @@ -107,64 +111,97 @@ def __init__( self.blank_id = len(self.tokenizer) self.max_symbols = max_symbols_per_step - def _greedy_decode( + @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", - x: Tensor, - seqlen: Tensor, - ) -> Tuple[List[int], List[int]]: + encoded: Tensor, + enc_len: Tensor, + ) -> List[Tuple[str, List[int], List[int]]]: """ - Greedy decode a single sequence. - Returns (token_ids, token_frames). - Token frames are encoder time indices t where a token is emitted. + RNN-T greedy decode: returns (text, token_ids, token_frames) per sample. + Token frames are encoder time indices where tokens are emitted. """ - T = int(seqlen.item()) if torch.is_tensor(seqlen) else int(seqlen) - - hyp: List[int] = [] - token_frames: List[int] = [] - dec_state: Optional[Tensor] = None - - last_label: Optional[Tensor] = None - - last_label_buf = torch.empty((1, 1), device=x.device, dtype=torch.long) - + 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): - f = x[t, :, :].unsqueeze(1) - new_symbols = 0 - - while new_symbols < self.max_symbols: - g, hidden = head.decoder.predict(last_label, dec_state) - k = int(head.joint.joint(f, g)[0, 0, 0, :].argmax(0).item()) + active = (t < enc_len).nonzero(as_tuple=False).squeeze(1).tolist() + if not active: + break - if k == self.blank_id: + for _ in range(self.max_symbols): + if not active: break - hyp.append(k) - token_frames.append(t) + 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] - dec_state = hidden - last_label_buf.fill_(k) - last_label = last_label_buf - new_symbols += 1 + next_active = [] + if fresh: + next_active.extend(emit_batch(fresh, t, fresh=True)) + if stateful: + next_active.extend(emit_batch(stateful, t, fresh=False)) - return hyp, token_frames + if not next_active: + break - @torch.inference_mode() - def decode( - self, - head: "RNNTHead", - encoded: Tensor, - enc_len: Tensor, - ) -> List[Tuple[List[int], List[int]]]: - """ - Decode RNN-T outputs for a batch. - Returns (token_ids, token_frames) per sample. - """ - B = encoded.shape[0] - encoded = encoded.transpose(1, 2) - - results: List[Tuple[List[int], List[int]]] = [] - for i in range(B): - inseq = encoded[i, :, :].unsqueeze(1) - results.append(self._greedy_decode(head, inseq, enc_len[i])) - return results + active = next_active + + return [(self.tokenizer.decode(h), h, tf) for h, tf in zip(hyps, token_frames)] diff --git a/gigaam/encoder.py b/gigaam/encoder.py index 55f7b45..3ae60df 100644 --- a/gigaam/encoder.py +++ b/gigaam/encoder.py @@ -6,6 +6,7 @@ import torch import torch.nn.functional as F from torch import Tensor, nn +from torch.utils.checkpoint import checkpoint try: from flash_attn import flash_attn_func @@ -18,6 +19,16 @@ from .utils import apply_masked_flash_attn, apply_rotary_pos_emb +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. @@ -169,6 +180,7 @@ def forward( ) -> 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) @@ -475,9 +487,11 @@ def __init__( 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", @@ -577,11 +591,22 @@ def forward(self, audio_signal: Tensor, length: Tensor) -> Tuple[Tensor, Tensor] pad_mask = ~pad_mask for layer in self.layers: - audio_signal = layer( - x=audio_signal, - pos_emb=pos_emb, - att_mask=att_mask, - pad_mask=pad_mask, - ) + 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 diff --git a/gigaam/model.py b/gigaam/model.py index 0ac91d3..fbcd07e 100644 --- a/gigaam/model.py +++ b/gigaam/model.py @@ -4,10 +4,11 @@ import omegaconf import torch from torch import Tensor, nn +from torch.utils.data import DataLoader from .preprocess import SAMPLE_RATE, load_audio from .types import LongformTranscriptionResult, Segment, TranscriptionResult, Word -from .utils import onnx_converter +from .utils import AudioDataset, onnx_converter LONGFORM_THRESHOLD = 25 * SAMPLE_RATE @@ -61,15 +62,15 @@ def embed_audio(self, wav_file: str) -> Tuple[Tensor, Tensor]: encoded, encoded_len = self.forward(wav, length) return encoded, encoded_len - def to_onnx(self, dir_path: str = ".") -> None: + 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) + 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 = ".") -> None: + def _to_onnx(self, dir_path: str = ".", dtype: torch.dtype = torch.float32) -> None: """ Export onnx model encoder to the specified dir. """ @@ -78,6 +79,7 @@ def _to_onnx(self, dir_path: str = ".") -> None: out_dir=dir_path, module=self.encoder, dynamic_axes=self.encoder.dynamic_axes(), + export_dtype=dtype, ) @@ -95,37 +97,31 @@ def _decode( self, encoded: Tensor, encoded_len: Tensor, - audio_length: int, + wav_lens: Tensor, word_timestamps: bool = False, - ) -> Tuple[str, Optional[List[Word]]]: - """ - Decode encoder output to text with optional word-level timestamps. - - Args: - encoded: Encoder output tensor - encoded_len: Length of encoded sequence - audio_length: Original audio length in samples - word_timestamps: Whether to compute word-level timestamps - - Returns: - Tuple of (text, words) where words is None if word_timestamps=False - """ - token_ids, token_frames = self.decoding.decode(self.head, encoded, encoded_len)[ - 0 - ] - - text = self.decoding.tokenizer.decode(token_ids) - + ) -> List[Tuple[str, Optional[List[Word]]]]: + decoded = self.decoding.decode(self.head, encoded, encoded_len) if not word_timestamps: - return text, None - + return [(t, None) for t, _, _ in decoded] from .timestamps_utils import compute_frame_shift, frames_to_words - frame_shift = compute_frame_shift(audio_length, int(encoded_len[0].item())) - words = frames_to_words( - self.decoding.tokenizer, token_ids, token_frames, frame_shift - ) - return text, words + 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( @@ -140,18 +136,19 @@ def transcribe( raise ValueError("Too long wav file, use 'transcribe_longform' method.") encoded, encoded_len = self.forward(wav, length) - text, words = self._decode( - encoded, encoded_len, int(length[0].item()), word_timestamps - ) + 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) -> Tensor: + def forward_for_export( + self, features: Tensor, feature_lengths: Tensor + ) -> Tuple[Tensor, Tensor]: """ Encoder-decoder forward to save model entirely in onnx format. """ - return self.head(self.encoder(features, feature_lengths)[0]) + encoded, encoded_len = self.encoder(features, feature_lengths) + return self.head(encoded), encoded_len - def _to_onnx(self, dir_path: str = ".") -> None: + def _to_onnx(self, dir_path: str = ".", dtype: torch.dtype = torch.float32) -> None: """ Export onnx ASR model. `ctc`: exported entirely in encoder-decoder format. @@ -167,35 +164,47 @@ def _to_onnx(self, dir_path: str = ".") -> None: module=self, inputs=self.encoder.input_example(), input_names=["features", "feature_lengths"], - output_names=["log_probs"], + 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) # export encoder + 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, **kwargs + 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. + 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. """ from .vad_utils import segment_audio_file @@ -204,36 +213,49 @@ def transcribe_longform( wav_file, SAMPLE_RATE, device=self._device, **kwargs ) - result_segments: List[Segment] = [] - for segment, segment_boundaries in zip(segments, boundaries): - wav = segment.to(self._device).unsqueeze(0).to(self._dtype) - length = torch.full([1], wav.shape[-1], device=self._device) - encoded, encoded_len = self.forward(wav, length) + if not segments: + return LongformTranscriptionResult(segments=[]) - seg_start = segment_boundaries[0] - seg_end = segment_boundaries[1] - - text, words = self._decode( - encoded, encoded_len, int(length[0].item()), word_timestamps - ) + ds = AudioDataset(segments, tokenizer=None) + dl = DataLoader( + ds, + batch_size=fr_batch_size, + shuffle=False, + collate_fn=AudioDataset.collate, + num_workers=fr_num_workers, + ) - if word_timestamps: - # Adjust word timestamps to absolute time positions - adjusted_words = [ - Word( - text=w.text, - start=round(w.start + seg_start, 3), - end=round(w.end + seg_start, 3), + 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 [] + ], + ) ) - for w in words - ] - result_segments.append( - Segment( - text=text, start=seg_start, end=seg_end, words=adjusted_words + else: + result_segments.append( + Segment(text=text, start=seg_start, end=seg_end) ) - ) - else: - result_segments.append(Segment(text=text, start=seg_start, end=seg_end)) return LongformTranscriptionResult(segments=result_segments) @@ -270,7 +292,7 @@ def forward_for_export(self, features: Tensor, feature_lengths: Tensor) -> Tenso enc_pooled = encoded.mean(dim=-1) return nn.functional.softmax(self.head(enc_pooled), dim=-1) - def _to_onnx(self, dir_path: str = ".") -> None: + def _to_onnx(self, dir_path: str = ".", dtype: torch.dtype = torch.float32) -> None: """ Export onnx Emo model. """ @@ -289,6 +311,7 @@ def _to_onnx(self, dir_path: str = ".") -> None: "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] diff --git a/gigaam/onnx_utils.py b/gigaam/onnx_utils.py index 1ae8c59..e01fc1c 100644 --- a/gigaam/onnx_utils.py +++ b/gigaam/onnx_utils.py @@ -1,116 +1,293 @@ +import logging import warnings -from typing import List, Optional, Tuple, Union +from pathlib import Path +from typing import Dict, List, Optional, Sequence, Tuple, Union import hydra import numpy as np import omegaconf import onnxruntime as rt import torch +from tqdm.auto import tqdm from .decoding import Tokenizer -from .preprocess import FeatureExtractor, load_audio +from .preprocess import FeatureExtractor +from .utils import AudioDataset warnings.simplefilter("ignore", category=UserWarning) - -DTYPE = np.float32 MAX_LETTERS_PER_FRAME = 3 +def _session_float_dtype(session: rt.InferenceSession) -> np.dtype: + """Infer numpy float dtype from the first float input of an ONNX session.""" + _type_map: Dict[str, np.dtype] = { + "tensor(float16)": np.dtype(np.float16), + "tensor(float)": np.dtype(np.float32), + "tensor(double)": np.dtype(np.float64), + } + for inp in session.get_inputs(): + if inp.type in _type_map: + return _type_map[inp.type] + return np.dtype(np.float32) + + +def _build_inputs(session: rt.InferenceSession, values: List[np.ndarray]) -> dict: + return {node.name: data for node, data in zip(session.get_inputs(), values)} + + +def _decode_ctc_batch( + labels: np.ndarray, + lengths: np.ndarray, + tokenizer: Tokenizer, +) -> List[str]: + blank_id = len(tokenizer) + b, t = labels.shape + lengths = np.clip(np.asarray(lengths, dtype=np.int64).reshape(-1), 0, t) + + skip_mask = labels != blank_id + skip_mask[:, 1:] &= labels[:, 1:] != labels[:, :-1] + + time = np.arange(t, dtype=np.int64)[None, :] + skip_mask &= time < lengths[:, None] + + return [tokenizer.decode(labels[i][skip_mask[i]].tolist()) for i in range(b)] + + +def _cat_states( + states: List[Tuple[np.ndarray, np.ndarray]], +) -> Tuple[np.ndarray, np.ndarray]: + hs = [s[0] for s in states] + cs = [s[1] for s in states] + return np.concatenate(hs, axis=1), np.concatenate(cs, axis=1) + + +def _split_state( + state: Tuple[np.ndarray, np.ndarray], +) -> List[Tuple[np.ndarray, np.ndarray]]: + h, c = state + b = h.shape[1] + return [(h[:, i : i + 1], c[:, i : i + 1]) for i in range(b)] + + +def _decode_rnnt_batch( + enc_features: np.ndarray, + enc_len: np.ndarray, + model_cfg: omegaconf.DictConfig, + sessions: List[Optional[rt.InferenceSession]], + tokenizer: Tokenizer, +) -> List[str]: + pred_sess, joint_sess = sessions[1:] + dtype = _session_float_dtype(pred_sess) + + enc_features = np.asarray(enc_features, dtype=dtype, order="C") + blank_idx = len(tokenizer) + pred_hidden = model_cfg.head.decoder.pred_hidden + pred_rnn_layers = model_cfg.head.decoder.pred_rnn_layers + B, _, T = enc_features.shape + + hyps: List[List[int]] = [[] for _ in range(B)] + last_label: List[Optional[np.ndarray]] = [None] * B + dec_state: List[Optional[Tuple[np.ndarray, np.ndarray]]] = [None] * B + + def emit_batch(batch_idx: List[int], t: int, fresh: bool) -> List[int]: + idx = np.asarray(batch_idx, dtype=np.int64) + f = enc_features[idx, :, t : t + 1] + + if fresh: + labels = np.full((len(batch_idx), 1), blank_idx, dtype=np.int64) + h = np.zeros((pred_rnn_layers, len(batch_idx), pred_hidden), dtype=dtype) + c = np.zeros((pred_rnn_layers, len(batch_idx), pred_hidden), dtype=dtype) + else: + labels = np.concatenate([last_label[i] for i in batch_idx], axis=0) + h, c = _cat_states([dec_state[i] for i in batch_idx]) + + pred_outputs = pred_sess.run( + [node.name for node in pred_sess.get_outputs()], + _build_inputs(pred_sess, [labels, h, c]), + ) + + joint_outputs = joint_sess.run( + [node.name for node in joint_sess.get_outputs()], + _build_inputs( + joint_sess, + [f, pred_outputs[0].swapaxes(1, 2)], + ), + ) + + k = joint_outputs[0][:, 0, 0, :].argmax(axis=-1) + emit_pos = np.nonzero(k != blank_idx)[0] + if emit_pos.size == 0: + return [] + + hidden_parts = _split_state((pred_outputs[1], pred_outputs[2])) + out = [] + + for p in emit_pos.tolist(): + bi = batch_idx[p] + tok = int(k[p]) + + hyps[bi].append(tok) + last_label[bi] = np.array([[tok]], dtype=np.int64) + dec_state[bi] = hidden_parts[p] + out.append(bi) + + return out + + enc_len = np.asarray(enc_len, dtype=np.int64).reshape(-1) + for t in range(T): + active = np.nonzero(t < enc_len)[0].tolist() + if not active: + break + + for _ in range(MAX_LETTERS_PER_FRAME): + 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 [tokenizer.decode(h) for h in hyps] + + def infer_onnx( - wav_file: Optional[str], + data: Union[str, Sequence[Union[str, np.ndarray, torch.Tensor]]], model_cfg: omegaconf.DictConfig, sessions: List[Optional[rt.InferenceSession]], - enc_features: Optional[np.ndarray] = None, preprocessor: Optional[FeatureExtractor] = None, tokenizer: Optional[Tokenizer] = None, -) -> Union[str, np.ndarray]: + batch_size: int = 16, + num_workers: int = 0, + progress: bool = True, +) -> Union[List[str], np.ndarray, List[np.ndarray]]: """ - Run ONNX sessions for the model, requires preprocessor instantiating. - The first session (the encoder one) and wav_file can be None if enc_features is provided. + Perform inference of GigaAM model with ONNX Runtime. + + Parameters + ---------- + data : Path to a manifest file or an iterable of audio paths / waveforms. + model_cfg : Model configuration. + sessions : List of ONNX Runtime inference sessions. + preprocessor : Optional[FeatureExtractor]. + tokenizer : Optional[Tokenizer]. + batch_size : Inference batch size. + num_workers : Number of workers for data loading (use for large datasets). + progress : Whether to show progress bar. + + Returns + ------- + Union[List[str], np.ndarray, List[np.ndarray]] + List of texts (ASR) / probs (Emo) / arrays (SSL) per sample. """ model_name = model_cfg.model_name - assert ( - enc_features is not None or sessions[0] is not None - ), "At least one of encoder session or enc_features is required" + if any(s in model_name for s in ["v1", "v2", "emo"]) and batch_size > 32: + logging.warning( + f"Batch size {batch_size} can be too large for v1/v2-family models. " + "This value can cause CUDA/cuDNN errors in Conv2d subsampling. " + "Forcing batch size to 32." + ) + batch_size = 32 - if preprocessor is None and enc_features is None: + if preprocessor is None: preprocessor = hydra.utils.instantiate(model_cfg.preprocessor) if tokenizer is None and ("ctc" in model_name or "rnnt" in model_name): tokenizer = hydra.utils.instantiate(model_cfg.decoding).tokenizer - if enc_features is None: - sgn = load_audio(wav_file) - input_signal = ( - preprocessor(sgn.unsqueeze(0), torch.tensor([sgn.shape[-1]]))[0] - .detach() - .numpy() - ) - enc_sess = sessions[0] - enc_inputs = { - node.name: data - for (node, data) in zip( - enc_sess.get_inputs(), - [input_signal.astype(DTYPE), [input_signal.shape[-1]]], - ) - } - enc_features = enc_sess.run( - [node.name for node in enc_sess.get_outputs()], enc_inputs - )[0] + if isinstance(data, str) and Path(data).suffix != ".tsv": + data = [data] + + dataset = AudioDataset(data) + + loader = torch.utils.data.DataLoader( + dataset, + batch_size=batch_size, + shuffle=False, + collate_fn=AudioDataset.collate, + num_workers=num_workers, + ) + loader_iter = ( + tqdm(loader, desc="Inference") + if progress and ("emo" in model_name or "ssl" in model_name) + else loader + ) + + enc_sess = sessions[0] + dtype = _session_float_dtype(enc_sess) if "emo" in model_name or "ssl" in model_name: - return enc_features + outputs = [] + for wavs, wav_lens in loader_iter: + input_signal, input_lengths = preprocessor(wavs.float(), wav_lens) + batch_outputs = enc_sess.run( + [node.name for node in enc_sess.get_outputs()], + _build_inputs( + enc_sess, + [ + input_signal.contiguous().numpy().astype(dtype), + input_lengths.numpy().astype(np.int64), + ], + ), + ) + outputs.extend(list(batch_outputs[0])) - blank_idx = len(tokenizer) - token_ids = [] - prev_token = blank_idx - if "ctc" in model_name: - prev_tok = blank_idx - for tok in enc_features.argmax(-1).squeeze().tolist(): - if (tok != prev_tok or prev_tok == blank_idx) and tok != blank_idx: - token_ids.append(tok) - prev_tok = tok - else: - pred_states = [ - np.zeros(shape=(1, 1, model_cfg.head.decoder.pred_hidden), dtype=DTYPE), - np.zeros(shape=(1, 1, model_cfg.head.decoder.pred_hidden), dtype=DTYPE), - ] - pred_sess, joint_sess = sessions[1:] - for j in range(enc_features.shape[-1]): - emitted_letters = 0 - while emitted_letters < MAX_LETTERS_PER_FRAME: - pred_inputs = { - node.name: data - for (node, data) in zip( - pred_sess.get_inputs(), [np.array([[prev_token]])] + pred_states - ) - } - pred_outputs = pred_sess.run( - [node.name for node in pred_sess.get_outputs()], pred_inputs - ) + return outputs - joint_inputs = { - node.name: data - for node, data in zip( - joint_sess.get_inputs(), - [enc_features[:, :, [j]], pred_outputs[0].swapaxes(1, 2)], - ) - } - log_probs = joint_sess.run( - [node.name for node in joint_sess.get_outputs()], joint_inputs + texts = [] + asr_iter = tqdm(loader, desc="ASR inference") if progress else loader + for wavs, wav_lens in asr_iter: + input_signal, input_lengths = preprocessor(wavs.float(), wav_lens) + batch_outputs = enc_sess.run( + [node.name for node in enc_sess.get_outputs()], + _build_inputs( + enc_sess, + [ + input_signal.contiguous().numpy().astype(dtype), + input_lengths.numpy().astype(np.int64), + ], + ), + ) + + batch_features = batch_outputs[0] + assert ( + len(batch_outputs) > 1 + ), "encoder must return enc_lengths for batched decoding" + batch_lengths = np.asarray(batch_outputs[1], dtype=np.int64).reshape(-1) + + if "ctc" in model_name: + texts.extend( + _decode_ctc_batch(batch_features.argmax(-1), batch_lengths, tokenizer) + ) + else: + texts.extend( + _decode_rnnt_batch( + batch_features, batch_lengths, model_cfg, sessions, tokenizer ) - token = log_probs[0].argmax(-1)[0][0] + ) - if token != blank_idx: - prev_token = token.item() - pred_states = pred_outputs[1:] - token_ids.append(token.item()) - emitted_letters += 1 - else: - break + return texts - return tokenizer.decode(token_ids) + +def _providers_list(provider: Optional[str]) -> List[Union[str, Tuple[str, dict]]]: + if provider == "CPUExecutionProvider": + return ["CPUExecutionProvider"] + if provider is not None and provider != "CUDAExecutionProvider": + return [provider] + cuda_opts = {"cudnn_conv_algo_search": "HEURISTIC"} + if "CUDAExecutionProvider" in rt.get_available_providers(): + return [("CUDAExecutionProvider", cuda_opts), "CPUExecutionProvider"] + return ["CPUExecutionProvider"] def load_onnx( @@ -120,41 +297,35 @@ def load_onnx( ) -> Tuple[ List[rt.InferenceSession], Union[omegaconf.DictConfig, omegaconf.ListConfig] ]: - """Load ONNX sessions for the given versions and cpu / cuda provider""" - if provider is None and "CUDAExecutionProvider" in rt.get_available_providers(): - provider = "CUDAExecutionProvider" - elif provider is None: - provider = "CPUExecutionProvider" + """ + Load a GigaAM model from ONNX Runtime given a model version. + Supports any family of models (ASR, Emo, SSL). + """ + providers = _providers_list(provider) opts = rt.SessionOptions() - opts.intra_op_num_threads = 16 + opts.graph_optimization_level = rt.GraphOptimizationLevel.ORT_ENABLE_ALL + opts.intra_op_num_threads = 8 if providers == ["CPUExecutionProvider"] else 1 opts.execution_mode = rt.ExecutionMode.ORT_SEQUENTIAL opts.log_severity_level = 3 model_cfg = omegaconf.OmegaConf.load(f"{onnx_dir}/{model_version}.yaml") + def _sess(path: str) -> rt.InferenceSession: + return rt.InferenceSession(path, providers=providers, sess_options=opts) + if "rnnt" not in model_version and "ssl" not in model_version: model_path = f"{onnx_dir}/{model_version}.onnx" - sessions = [ - rt.InferenceSession(model_path, providers=[provider], sess_options=opts) - ] + sessions = [_sess(model_path)] elif "ssl" in model_version: pth = f"{onnx_dir}/{model_version}" - enc_sess = rt.InferenceSession( - f"{pth}_encoder.onnx", providers=[provider], sess_options=opts - ) - sessions = [enc_sess] + sessions = [_sess(f"{pth}_encoder.onnx")] else: pth = f"{onnx_dir}/{model_version}" - enc_sess = rt.InferenceSession( - f"{pth}_encoder.onnx", providers=[provider], sess_options=opts - ) - pred_sess = rt.InferenceSession( - f"{pth}_decoder.onnx", providers=[provider], sess_options=opts - ) - joint_sess = rt.InferenceSession( - f"{pth}_joint.onnx", providers=[provider], sess_options=opts - ) - sessions = [enc_sess, pred_sess, joint_sess] + sessions = [ + _sess(f"{pth}_encoder.onnx"), + _sess(f"{pth}_decoder.onnx"), + _sess(f"{pth}_joint.onnx"), + ] return sessions, model_cfg diff --git a/gigaam/types.py b/gigaam/types.py index d66b0e9..92cb57f 100644 --- a/gigaam/types.py +++ b/gigaam/types.py @@ -1,5 +1,16 @@ from dataclasses import dataclass -from typing import List, Optional +from typing import List, Optional, Union + +import numpy as np +from torch import Tensor + + +@dataclass +class AudioDatasetSample: + item: Union[str, np.ndarray, Tensor] + duration: float + text: Optional[str] = None + tokens: Optional[List[int]] = None @dataclass diff --git a/gigaam/utils.py b/gigaam/utils.py index 2c1ec35..f1b2810 100644 --- a/gigaam/utils.py +++ b/gigaam/utils.py @@ -1,15 +1,20 @@ +import csv import os import warnings +from collections.abc import Iterable from pathlib import Path -from typing import Dict, List, Optional, Tuple, Union +from typing import Dict, List, Optional, Tuple, Union, cast import numpy as np +import soundfile as sf import torch import torch.nn.functional as F +import torchaudio from torch import Tensor from torch.jit import TracerWarning -from .preprocess import load_audio +from .preprocess import SAMPLE_RATE +from .types import AudioDatasetSample def onnx_converter( @@ -23,7 +28,12 @@ def onnx_converter( 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: @@ -31,14 +41,17 @@ def onnx_converter( 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") - saved_dtype = next(module.parameters()).dtype with warnings.catch_warnings(), torch.no_grad(): warnings.simplefilter("ignore", category=UserWarning) warnings.simplefilter("ignore", category=TracerWarning) torch.onnx.export( - module.to(torch.float32), + module.to(export_dtype), inputs, out_path, input_names=input_names, @@ -48,7 +61,8 @@ def onnx_converter( dynamo=False, ) print(f"Successfully ported onnx {model_name} to {out_path}.") - module.to(saved_dtype) + # We force the whole module to float32 to avoid fp16 preprocessing issues + module.float() def format_time(seconds: float) -> str: @@ -81,6 +95,8 @@ def apply_rotary_pos_emb( 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) @@ -139,7 +155,7 @@ def apply_masked_flash_attn( return scores -def download_short_audio(): +def download_short_audio() -> str: """Download test audio file if not exists""" audio_file = "example.wav" if not os.path.exists(audio_file): @@ -150,7 +166,7 @@ def download_short_audio(): return audio_file -def download_long_audio(): +def download_long_audio() -> str: """Download test audio file if not exists""" audio_file = "long_example.wav" if not os.path.exists(audio_file): @@ -163,37 +179,214 @@ def download_long_audio(): class AudioDataset(torch.utils.data.Dataset): """ - Helper class for creating batched inputs + 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, lst: List[Union[str, np.ndarray, torch.Tensor]]): - if len(lst) == 0: - raise ValueError("AudioDataset cannot be initialized with an empty list") - assert isinstance( - lst[0], (str, np.ndarray, torch.Tensor) - ), f"Unexpected dtype: {type(lst[0])}" - self.lst = lst + 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 - def __len__(self): - return len(self.lst) + 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] - def __getitem__(self, idx): - item = self.lst[idx] + return tokenizer.model.encode + + def normalize_text(self, text: str) -> str: + if not self.raw_text: + return text + + text = text.replace("ё", "е").replace("Ё", "Е") + text = " ".join(text.split()) + + if self.tokenizer is not None and getattr(self.tokenizer, "charwise", False): + vocab = set(self.tokenizer.vocab) + return "".join(c for c in text.lower() if c in vocab) + + return text.lower() + + @staticmethod + def _get_duration(item: Union[str, np.ndarray, Tensor]) -> float: if isinstance(item, str): - wav_tns = load_audio(item) - elif isinstance(item, np.ndarray): - wav_tns = torch.from_numpy(item) - elif isinstance(item, torch.Tensor): - wav_tns = item - else: - raise RuntimeError(f"Unexpected sample type: {type(item)} at idx={idx}") - return wav_tns + 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 collate(wavs): - lengths = torch.tensor([len(wav) for wav in wavs]) - max_len = lengths.max().item() - wav_tns = torch.zeros(len(wavs), max_len, dtype=wavs[0].dtype) - for idx, wav in enumerate(wavs): - wav_tns[idx, : wav.shape[-1]] = wav.squeeze() - return wav_tns, lengths + 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 diff --git a/pyproject.toml b/pyproject.toml index 1cf1912..e4aa8e1 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -19,28 +19,37 @@ dependencies = [ "omegaconf==2.3.*", "onnx==1.19.*", "onnxruntime==1.23.*", + "soundfile", "sentencepiece", "tqdm", ] [project.optional-dependencies] +# Why torch is optional: often preinstalled and reinstall is slow. +# Older builds also work, pinned extras are recommended. torch = [ - "torch>=2.5,<2.9", - "torchaudio>=2.5,<2.9", + "torch>=2.6", + "torchaudio>=2.6", ] longform = [ - "torch==2.8.*", - "torchaudio==2.8.*", - "pyannote.audio==4.0", - "torchcodec==0.7", "numba>=0.62", + "pyannote.audio==4.0.*", + "pyarrow>=23", + "torch==2.10.*", + "torchaudio==2.10.*", + "torchcodec==0.10", + "transformers==5.*", +] +train = [ + "editdistance", + "lightning>=2.6", + "tensorboard>=2.20", + "wandb>=0.25", ] tests = [ "pytest", "pytest-cov", "scipy", - "soundfile", - "librosa", ] lint = [ "black==26.1.0", diff --git a/tests/test_onnx.py b/tests/test_onnx.py index 948df1e..d6d7b76 100644 --- a/tests/test_onnx.py +++ b/tests/test_onnx.py @@ -3,6 +3,7 @@ import numpy as np import pytest +import torch import gigaam from gigaam.onnx_utils import infer_onnx, load_onnx @@ -18,31 +19,60 @@ def test_audio(): return download_short_audio() -@pytest.mark.parametrize("revision", ["emo", "v2_ssl", "v3_ctc", "v3_e2e_rnnt"]) -def test_onnx_converting(revision, test_audio): - """Test specific model revision loads and processes audio (partial models enabled)""" +@pytest.mark.parametrize( + "revision, export_dtype", + [ + ("emo", torch.float32), + ("v2_ssl", torch.float16), + ("v3_ctc", torch.float16), + ("v3_e2e_rnnt", torch.float32), + ], +) +def test_onnx_converting(revision, export_dtype, test_audio): + """Test model revision converts to ONNX and produces correct batched output.""" onnx_dir = "test_onnx_tmp" - model = gigaam.load_model(revision, fp16_encoder=False) - model.to_onnx(dir_path=onnx_dir) + model = gigaam.load_model(revision) + model.to_onnx(dir_path=onnx_dir, dtype=export_dtype) sessions, model_cfg = load_onnx(onnx_dir, revision) - result = infer_onnx(test_audio, model_cfg, sessions) + + data = [test_audio, test_audio] + result = infer_onnx(data, model_cfg, sessions, batch_size=2) shutil.rmtree(onnx_dir) + assert isinstance(result, list), f"{revision}: expected list, got {type(result)}" + assert len(result) == 2, f"{revision}: expected 2 results, got {len(result)}" + if "ssl" in revision: + for i, r in enumerate(result): + assert isinstance(r, np.ndarray), f"{revision}[{i}]: expected ndarray" + assert r.ndim == 2, f"{revision}[{i}]: expected 2D array, got {r.ndim}D" + orig_embed = model.embed_audio(test_audio)[0].detach().cpu().numpy() - diff = np.abs(orig_embed - result).max() - assert diff < 0.01, f"{revision}: ONNX embed failed with diff {diff}" + tol = 0.01 if export_dtype == torch.float16 else 0.001 + for i in range(2): + diff = np.abs(orig_embed - result[i]).mean() + assert diff < tol, f"{revision}[{i}]: ONNX embed diff {diff}" elif "emo" in revision: orig_probs = model.get_probs(test_audio) - pred_probs = {model.id2name[i]: result[0, i] for i in range(len(model.id2name))} - assert all( - abs(orig_probs[em] - pred_probs[em]) < 1e-3 for em in orig_probs - ), f"{revision}: ONNX emotions probs failed: {pred_probs}" + tol = 1e-3 if export_dtype == torch.float16 else 1e-4 + for i in range(2): + r = result[i] + assert isinstance(r, np.ndarray), f"{revision}[{i}]: expected ndarray" + pred_probs = { + model.id2name[j]: float(r[j]) for j in range(len(model.id2name)) + } + assert all( + abs(orig_probs[em] - pred_probs[em]) < tol for em in orig_probs + ), f"{revision}[{i}]: ONNX emo probs failed: {pred_probs}" else: orig_text = model.transcribe(test_audio).text - assert orig_text == result, f"{revision}: ONNX transcribe failed: {result}" + for i in range(2): + assert isinstance(result[i], str), f"{revision}[{i}]: expected str" + assert ( + orig_text == result[i] + ), f"{revision}[{i}]: ONNX transcribe failed: {result[i]}" if __name__ == "__main__": diff --git a/tests/test_reading.py b/tests/test_reading.py index 1179ef0..d4bffb1 100644 --- a/tests/test_reading.py +++ b/tests/test_reading.py @@ -1,13 +1,11 @@ import logging -import librosa import pytest import torch from torch.nn.functional import softmax import gigaam -from gigaam.preprocess import SAMPLE_RATE -from gigaam.utils import download_short_audio +from gigaam.utils import AudioDataset, download_short_audio logging.basicConfig(level=logging.INFO) logger = logging.getLogger(__name__) @@ -20,13 +18,13 @@ def test_audio(): @pytest.mark.parametrize("revision", ["emo"]) -def test_librosa_loading(revision, test_audio): - """Test the outputs with librosa.load are close to load_audio""" +def test_torchaudio_loading(revision, test_audio): + """Torchaudio-loaded waveform should match get_probs(path) (ffmpeg load_audio).""" model = gigaam.load_model(revision) - wav_tns = torch.from_numpy(librosa.load(test_audio, sr=SAMPLE_RATE)[0]) + wav_tns = AudioDataset([test_audio])[0] lengths = torch.full([1], wav_tns.shape[-1], device=model._device) with torch.no_grad(): - encoded, encoded_len = model( + encoded, _ = model( wav_tns.unsqueeze(0).to(model._device).to(model._dtype), lengths ) orig_probs = model.get_probs(test_audio) @@ -37,7 +35,9 @@ def test_librosa_loading(revision, test_audio): model.id2name[i]: pred_probs[i] for i in range(len(model.id2name)) } are_close = max(abs(pred_probs[k] - orig_probs[k]) for k in orig_probs) < 1e-3 - assert are_close, f"Emotions with librosa failed: {orig_probs} != {pred_probs}" + assert ( + are_close + ), f"Emotions with torchaudio failed: {orig_probs} != {pred_probs}" if __name__ == "__main__": diff --git a/tests/test_training.py b/tests/test_training.py new file mode 100644 index 0000000..bef02b4 --- /dev/null +++ b/tests/test_training.py @@ -0,0 +1,222 @@ +import gc +import os +import re +import shutil +import subprocess +import sys +from pathlib import Path + +import pytest +import soundfile as sf + +import gigaam +from gigaam.preprocess import SAMPLE_RATE +from gigaam.utils import download_long_audio + +TRAIN_UTILS_DIR = Path(__file__).resolve().parents[1] / "train_utils" +MAX_SAMPLES = 2 + + +def _write_manifest( + manifest_path: Path, audio_paths: list[Path], texts: list[str] +) -> None: + rows = ["path\tduration\ttranscription"] + for audio_path, text in zip(audio_paths, texts): + duration = sf.info(audio_path).duration + rows.append(f"{audio_path}\t{duration:.3f}\t{text}") + manifest_path.write_text("\n".join(rows) + "\n", encoding="utf-8") + + +def _run_python(args: list[str], cwd: Path) -> subprocess.CompletedProcess[str]: + env = os.environ.copy() + env["CUDA_VISIBLE_DEVICES"] = "" + return subprocess.run( + [sys.executable, *args], + cwd=cwd, + env=env, + check=True, + capture_output=True, + text=True, + ) + + +def _extract_e2e_wer(output: str) -> float: + match = re.search(r"WER e2e:\s+([0-9.]+)%", output) + assert match, f"WER e2e not found in output:\n{output}" + return float(match.group(1)) + + +def _clear_gigaam_cache_checkpoints() -> None: + cache_dir = Path.home() / ".cache" / "gigaam" + if not cache_dir.is_dir(): + return + for ckpt in cache_dir.glob("*.ckpt"): + ckpt.unlink(missing_ok=True) + + +@pytest.fixture(scope="session") +def pseudo_labeled_dataset(tmp_path_factory: pytest.TempPathFactory) -> Path: + root = tmp_path_factory.mktemp("pseudo_labels") + audio_dir = root / "audio" + audio_dir.mkdir() + + audio_path = download_long_audio() + model = gigaam.load_model("v3_ctc", device="cpu", fp16_encoder=False) + labels = model.transcribe_longform(audio_path, fr_batch_size=8, fr_num_workers=0) + del model + gc.collect() + audio, sr = sf.read(audio_path, dtype="float32") + assert sr == SAMPLE_RATE + + assert labels.segments + + segment_paths = [] + texts = [] + for idx, labeled in enumerate(labels.segments[:MAX_SAMPLES]): + segment_path = audio_dir / f"{idx:02d}.wav" + start_idx = int(labeled.start * SAMPLE_RATE) + end_idx = int(labeled.end * SAMPLE_RATE) + sf.write(segment_path, audio[start_idx:end_idx], SAMPLE_RATE) + segment_paths.append(segment_path) + texts.append(str(labeled.text).lower()) + + manifest_path = root / "manifest.tsv" + _write_manifest(manifest_path, segment_paths, texts) + return manifest_path + + +@pytest.mark.parametrize( + ("model_name", "extra_args", "max_steps", "e2e_threshold", "min_e2e_gain"), + [ + ("v3_e2e_ctc", ["--lr", "5e-4", "--activation_checkpointing"], 3, 20.0, 20.0), + ( + "v3_e2e_rnnt", + ["--rnnt_subbatch_size", "1", "--lr", "5e-4", "--activation_checkpointing"], + 6, + 20.0, + 20.0, + ), + ], +) +def test_training_and_eval_on_cpu( + tmp_path: Path, + pseudo_labeled_dataset: Path, + model_name: str, + extra_args: list[str], + max_steps: int, + e2e_threshold: float, + min_e2e_gain: float, +) -> None: + exp_name = f"pytest_{model_name}" + output_dir = tmp_path / "artifacts" + + try: + _run_training_case( + model_name, + extra_args, + max_steps, + e2e_threshold, + min_e2e_gain, + pseudo_labeled_dataset, + output_dir, + exp_name, + ) + finally: + if output_dir.exists(): + shutil.rmtree(output_dir, ignore_errors=True) + + +def _run_training_case( + model_name: str, + extra_args: list[str], + max_steps: int, + e2e_threshold: float, + min_e2e_gain: float, + pseudo_labeled_dataset: Path, + output_dir: Path, + exp_name: str, +) -> None: + baseline_eval = _run_python( + [ + "eval.py", + "--model_name", + model_name, + "--eval_manifest", + str(pseudo_labeled_dataset), + "--batch_size", + str(MAX_SAMPLES), + "--num_workers", + "0", + "--device", + "cpu", + "--disable_tqdm", + ], + cwd=TRAIN_UTILS_DIR, + ) + baseline_e2e = _extract_e2e_wer(baseline_eval.stdout) + + # Сlear gigaam cache checkpoints to save disk space in github actions + _clear_gigaam_cache_checkpoints() + + train_cmd = [ + "train.py", + "--model_name", + model_name, + "--train_manifest", + str(pseudo_labeled_dataset), + "--val_manifest", + str(pseudo_labeled_dataset), + "--output_dir", + str(output_dir), + "--exp_name", + exp_name, + "--batch_size", + "1", + "--eval_batch_size", + str(MAX_SAMPLES), + "--num_workers", + "0", + "--precision", + "32", + "--max_steps", + str(max_steps), + "--val_check_steps", + str(max_steps), + "--disable_tqdm", + "--log_every_n_steps", + "1", + "--skip_initial_validation", + "--save_top_k", + "1", + ] + _run_python(train_cmd + extra_args, cwd=TRAIN_UTILS_DIR) + + ckpt_dir = output_dir / "models" / exp_name + checkpoints = sorted(ckpt_dir.glob("*.ckpt")) + assert checkpoints, f"No checkpoints found in {ckpt_dir}" + + eval_run = _run_python( + [ + "eval.py", + "--checkpoint", + str(checkpoints[0]), + "--eval_manifest", + str(pseudo_labeled_dataset), + "--batch_size", + str(MAX_SAMPLES), + "--num_workers", + "0", + "--device", + "cpu", + "--disable_tqdm", + ], + cwd=TRAIN_UTILS_DIR, + ) + + e2e_wer = _extract_e2e_wer(eval_run.stdout) + assert e2e_wer <= e2e_threshold + assert baseline_e2e - e2e_wer >= min_e2e_gain + + +if __name__ == "__main__": + pytest.main([__file__, "-v"]) diff --git a/train_utils/README.md b/train_utils/README.md new file mode 100644 index 0000000..ddd2da5 --- /dev/null +++ b/train_utils/README.md @@ -0,0 +1,147 @@ +# GigaAM Fine-tuning (CTC / RNNT) + +For training and evaluation examples and a more detailed description of manifests, see [`example.ipynb`](./example.ipynb). + +## Setup + +From the repository root, install the training dependencies: + +```bash +pip install -e ".[train]" + +cd train_utils +``` + +## Data format + +TSV manifest (tab-separated columns): `path`, `duration`, and optionally `transcription`. Paths may be absolute or relative to the manifest directory. `transcription` may be omitted for audio-only manifests. + +``` +path duration transcription +audio/0001.wav 3.21 привет как дела +``` + +## Training + +```bash +python train.py \ + --model_name v3_e2e_ctc \ + --train_manifest /path/to/manifest_train.tsv \ + --val_manifest /path/to/manifest_val.tsv \ + --max_epochs 3 \ + --val_check_interval 0.5 \ + --batch_size 64 \ + --eval_batch_size 64 \ + --lr 8e-5 \ + --activation_checkpointing + +python train.py \ + --model_name v3_rnnt \ + --train_manifest /path/to/manifest_train.tsv \ + --val_manifest /path/to/manifest_val.tsv \ + --raw_text \ + --max_epochs 2 \ + --val_check_interval 0.5 \ + --batch_size 8 \ + --accumulate_grad_batches 2 \ + --rnnt_subbatch_size 2 \ + --eval_batch_size 64 \ + --lr 2e-5 \ + --val_first_batches 50 +``` + +### Arguments + +#### Model and data + +| Argument | Default | Description | +|---|---|---| +| `--model_name` | required | Pretrained GigaAM model name | +| `--train_manifest` | required | TSV manifest for training | +| `--val_manifest` | required | TSV manifest for validation | +| `--raw_text` | off | For non-E2E setups: lowercase text, drop punctuation, restrict to the character vocabulary | +| `--max_duration` | `20.0` | Maximum audio length in seconds (dataset filter) | +| `--min_duration` | `0.1` | Minimum audio length in seconds (dataset filter) | + +#### Scheduling (epochs vs steps) + +| Argument | Default | Description | +|---|---|---| +| `--max_epochs` | `None` | Train for this many epochs (omit if using `--max_steps`) | +| `--val_check_interval` | `1.0` | Run validation every N-th fraction of an epoch | +| `--max_steps` | `None` | Train for this many steps (requires `--val_check_steps`) | +| `--val_check_steps` | `None` | With `--max_steps`: validate every N training steps | +| `--val_first_batches` | `None` | If set, run validation on only the first N batches | + +#### Batching and memory + +| Argument | Default | Description | +|---|---|---| +| `--batch_size` | `8` | Per-device training batch size | +| `--eval_batch_size` | `32` | Validation batch size | +| `--rnnt_subbatch_size` | `0` | RNNT loss sub-batches (`0` disables) | +| `--num_workers` | `4` | Number of `DataLoader` workers | +| `--precision` | `32` | Lightning precision (`16`, `bf16`, `32`, ...) | +| `--accumulate_grad_batches` | `1` | Gradient accumulation steps | +| `--accelerator` | `auto` | Lightning accelerator (`auto`, `cpu`, `gpu`, ...) | +| `--devices` | `1` | Number of devices (DDP when `> 1`) | +| `--activation_checkpointing` | off | Activation checkpointing for each Conformer layer | +| `--freeze_encoder` | off | Freeze encoder weights | + +#### Optimizer + +| Argument | Default | Description | +|---|---|---| +| `--lr` | `2e-5` | Peak learning rate (AdamW) | +| `--weight_decay` | `0.01` | AdamW weight decay | +| `--warmup_ratio` | `0.1` | Linear warmup fraction before cosine decay | +| `--gradient_clip_val` | `1.0` | Global gradient norm clipping | +| `--seed` | `42` | Passed to `pl.seed_everything` | + +#### SpecAugment + +| Argument | Default | Description | +|---|---|---| +| `--freq_masks` | `2` | Number of frequency masks | +| `--freq_width` | `27` | Maximum width of each frequency mask (bins) | +| `--time_masks` | `2` | Number of time masks | +| `--time_width` | `20` | Maximum width of each time mask (frames) | +| `--disable_spec_augment` | off | Disable SpecAugment (enabled by default) | + +#### Outputs, logging and resuming + +| Argument | Default | Description | +|---|---|---| +| `--output_dir` | `./checkpoints` | Checkpoints under `models//`, TensorBoard under `tb_logs/` | +| `--exp_name` | auto | Run directory name; if omitted, derived from hyperparameters | +| `--log_every_n_steps` | `25` | Logging interval in steps | +| `--save_top_k` | `2` | Keep this many best checkpoints by `val_wer` | +| `--disable_tqdm` | off | Disable progress bars | +| `--skip_initial_validation` | off | Skip `trainer.validate` before `fit` | +| `--resume_from_checkpoint` | `None` | Path to a Lightning `.ckpt` file to resume training from | + +## Evaluation + +Evaluate a fine-tuned checkpoint: + +```bash +python eval.py \ + --checkpoint ./checkpoints/models//gigaam-*.ckpt \ + --eval_manifest /path/to/manifest.tsv +``` + +Evaluate a pretrained GigaAM model: + +```bash +python eval.py --model_name v3_e2e_ctc --eval_manifest /path/to/manifest.tsv +``` + +This writes `preds.jsonl` and prints WER. Predictions are saved under `predictions///step_/preds.jsonl` (`step_` is omitted for pretrained models) next to the manifest. WER is reported on the original transcripts (end-to-end) and on raw texts. + +## Loading fine-tuned checkpoints + +`gigaam.load_model` accepts a path to a Lightning `.ckpt` file, so fine-tuned models can be loaded the same way as pretrained ones: + +```python +model = gigaam.load_model("./checkpoints/models//gigaam-*.ckpt") +``` diff --git a/train_utils/eval.py b/train_utils/eval.py new file mode 100644 index 0000000..aeba55b --- /dev/null +++ b/train_utils/eval.py @@ -0,0 +1,96 @@ +"""Evaluate a GigaAM checkpoint (pretrained or fine-tuned).""" + +import argparse +import json +import os +import re +from pathlib import Path + +import torch +from torch.utils.data import DataLoader +from tqdm import tqdm +from utils import compute_wer + +import gigaam +from gigaam.utils import AudioDataset + + +def main(): + p = argparse.ArgumentParser() + p.add_argument("--eval_manifest", required=True) + p.add_argument("--checkpoint", default=None) + p.add_argument("--model_name", default=None) + p.add_argument("--batch_size", type=int, default=64) + p.add_argument("--num_workers", type=int, default=4) + p.add_argument("--max_duration", type=float, default=None) + p.add_argument("--min_duration", type=float, default=0.0) + p.add_argument("--device", default="cuda" if torch.cuda.is_available() else "cpu") + p.add_argument("--disable_tqdm", action="store_true", default=False) + args = p.parse_args() + + src = args.checkpoint or args.model_name + assert src, "Pass --checkpoint or --model_name" + model = gigaam.load_model(src, device=args.device) + + ds = AudioDataset( + args.eval_manifest, + tokenizer=model.decoding.tokenizer, + max_duration=args.max_duration, + min_duration=args.min_duration, + raw_text=False, + return_tokens=False, + ) + samples = ds.samples + print(f"Loaded {len(samples)} samples") + + dl = DataLoader( + ds, + batch_size=args.batch_size, + shuffle=False, + collate_fn=AudioDataset.collate, + num_workers=args.num_workers, + pin_memory=args.device != "cpu", + ) + + preds, idx = [], 0 + with torch.inference_mode(): + for wav_pad, wav_lens in tqdm(dl, desc="Inference", disable=args.disable_tqdm): + enc, enc_len = model(wav_pad.to(args.device), wav_lens.to(args.device)) + for txt, _, _ in model.decoding.decode(model.head, enc, enc_len): + s = samples[idx] + preds.append( + { + "audio_filepath": s.item, + "text": s.text or "", + "pred_text": txt, + "duration": s.duration, + } + ) + idx += 1 + + manifest_path = Path(args.eval_manifest) + if src and os.path.isfile(os.path.expanduser(src)): + ckpt_path = Path(src) + experiment = ckpt_path.parent.name + step_match = re.search(r"step=(\d+)", ckpt_path.stem) + step_tag = f"step_{step_match.group(1)}" if step_match else ckpt_path.stem + ckpt_name = f"{experiment}/{step_tag}" + else: + ckpt_name = src + out = manifest_path.parent / "predictions" / manifest_path.stem / ckpt_name + out.mkdir(parents=True, exist_ok=True) + with open(out / "preds.jsonl", "w", encoding="utf-8") as f: + for r in preds: + f.write(json.dumps(r, ensure_ascii=False) + "\n") + + print(f"Saved predictions to {out / 'preds.jsonl'}") + + wer_e2e, wer_raw, e2e_err, e2e_w, raw_err, raw_w = compute_wer(preds) + print( + f"WER e2e: {wer_e2e:.2f}% ({e2e_err}/{e2e_w} words)\n" + f"WER raw: {wer_raw:.2f}% ({raw_err}/{raw_w} words)" + ) + + +if __name__ == "__main__": + main() diff --git a/train_utils/example.ipynb b/train_utils/example.ipynb new file mode 100644 index 0000000..87b4682 --- /dev/null +++ b/train_utils/example.ipynb @@ -0,0 +1,1093 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "id": "dd4d3e32", + "metadata": {}, + "source": [ + "## GigaAM CTC / RNNT Fine-tuning\n", + "\n", + "This notebook describes the dataset format and provides fine-tuning examples for CTC and RNNT models in end-to-end and raw setups, using activation checkpointing, RNNT loss sub-batching, gradient accumulation, DDP, and encoder freezing for faster, more memory-efficient training." + ] + }, + { + "cell_type": "markdown", + "id": "0a8a1a96", + "metadata": {}, + "source": [ + "### Dataset" + ] + }, + { + "cell_type": "markdown", + "id": "abe1a8da", + "metadata": {}, + "source": [ + "#### Format\n", + "\n", + "The training pipeline supports a `.tsv` manifest format like the example below.\n", + "\n", + "Each row describes one sample:\n", + "\n", + "* `path` - relative or absolute path to the audio file\n", + "* `duration` - audio length in seconds\n", + "* `transcription` - reference text for this audio sample\n", + "\n", + "Example:\n", + "\n", + "```tsv\n", + "path duration transcription\n", + "audio/train/000000.wav 5.265 Вот только они совсем не радовали, а, напротив, потрясали и ужасали.\n", + "audio/train/000001.wav 4.268 Убедившись, что никого нет, он приблизился ко мне и понизил голос.\n", + "```\n", + "\n", + "Usage:\n", + "\n", + "```bash\n", + "python train.py \\\n", + " --train_manifest /path/to/train/manifest.tsv \\\n", + " --val_manifest /path/to/val/manifest.tsv \\\n", + " ...\n", + "\n", + "python eval.py \\\n", + " --eval_manifest /path/to/eval/manifest.tsv \\\n", + " ...\n", + "```\n", + "\n", + "**Note:** By default, samples in the training and validation sets are filtered by duration to the `[0.1s, 20s]` range. You can change these limits with `min_duration` and `max_duration`." + ] + }, + { + "cell_type": "markdown", + "id": "4ab9dff4", + "metadata": {}, + "source": [ + "#### Loading data" + ] + }, + { + "cell_type": "code", + "execution_count": 1, + "id": "ef723e4c", + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "Loading Vikhrmodels/ToneBooks...\n" + ] + }, + { + "name": "stderr", + "output_type": "stream", + "text": [ + "Warning: You are sending unauthenticated requests to the HF Hub. Please set a HF_TOKEN to enable higher rate limits and faster downloads.\n" + ] + }, + { + "data": { + "application/vnd.jupyter.widget-view+json": { + "model_id": "f73eb34fa6054af7b3bf125e2f407416", + "version_major": 2, + "version_minor": 0 + }, + "text/plain": [ + "Loading dataset shards: 0%| | 0/21 [00:00 "GigaAMFineTuner": + super().train(mode) + # We keep it verbose for the freezed preprocessor + self.preprocessor.eval() + if self._freeze_encoder: + self.encoder.eval() + return self + + def on_train_epoch_start(self): + self.train() + + def _ctc_loss( + self, + log_probs: Tensor, + targets: Tensor, + input_lens: Tensor, + target_lens: Tensor, + ) -> Tensor: + return self._ctc( + log_probs.transpose(0, 1), + targets.long(), + input_lens.long(), + target_lens.long(), + ).mean() + + def _rnnt_loss( + self, logits: Tensor, targets: Tensor, logit_lens: Tensor, target_lens: Tensor + ) -> Tensor: + return _ta_rnnt_loss( + logits=logits.float(), + targets=targets.int(), + logit_lengths=logit_lens.int(), + target_lengths=target_lens.int(), + blank=self._blank_id, + reduction="mean", + fused_log_softmax=True, + ) + + def _encode(self, wavs: Tensor, wav_lens: Tensor) -> Tuple[Tensor, Tensor]: + amp_dev = next(self.parameters()).device.type + with torch.amp.autocast(amp_dev, enabled=False): + features, feat_lens = self.preprocessor(wavs.float(), wav_lens) + if self.training and self._spec_augment: + for aug in self._freq_aug: + features = aug(features) + for aug in self._time_aug: + features = aug(features) + return self.encoder(features, feat_lens) + + def _rnnt_joint(self, encoded: Tensor, tokens: Tensor) -> Tensor: + head = self.head + assert isinstance(head, RNNTHead) + enc_t = encoded.float().transpose(1, 2) + dec, jnt = head.decoder, head.joint + bos = torch.zeros( + enc_t.size(0), 1, dec.pred_hidden, device=enc_t.device, dtype=enc_t.dtype + ) + pred_out, _ = dec.lstm( + torch.cat([bos, dec.embed(tokens)], dim=1).transpose(0, 1) + ) + return jnt.joint_net( + jnt.enc(enc_t).unsqueeze(2) + + jnt.pred(pred_out.transpose(0, 1)).unsqueeze(1) + ) + + def _rnnt_forward( + self, encoded: Tensor, enc_lens: Tensor, tokens: Tensor, tok_lens: Tensor + ) -> Tensor: + B = encoded.size(0) + subb = self._rnnt_subbatch_size or B + + # Ensure int32-safe indexing in _ta_rnnt_loss + max_t = int(enc_lens.max().item()) + max_u1 = int(tok_lens.max().item()) + 1 + vocab_size = self._rnnt_vocab_size + while subb > 1 and subb * max_t * max_u1 * vocab_size >= (1 << 31): + subb = max(1, subb // 2) + + losses: List[torch.Tensor] = [] + weights: List[int] = [] + for i in range(0, B, subb): + s = slice(i, min(i + subb, B)) + enc_i, enc_lens_i = encoded[s], enc_lens[s] + tok_i, tok_lens_i = tokens[s], tok_lens[s] + enc_i = enc_i[:, :, : int(enc_lens_i.max().item())].contiguous() + tok_i = tok_i[:, : int(tok_lens_i.max().item())].contiguous() + logits_i = self._rnnt_joint(enc_i, tok_i) # [b, T, U+1, V] + T, U1 = logits_i.shape[1:3] + enc_lens_i = enc_lens_i.clamp(min=1, max=T) + tok_lens_i = tok_lens_i.clamp(min=1, max=U1 - 1) + loss_i = self._rnnt_loss(logits_i, tok_i, enc_lens_i, tok_lens_i) + losses.append(loss_i * (enc_i.size(0))) + weights.append(enc_i.size(0)) + del logits_i + + return torch.stack(losses).sum() / sum(weights) + + def _batch_wer( + self, hyps: List[str], tokens: Tensor, tok_lens: Tensor + ) -> Tuple[int, int]: + errors = words = 0 + for i, hyp in enumerate(hyps): + ref = self._tokenizer.decode(tokens[i, : tok_lens[i]].tolist()).split() + hyp_w = hyp.split() + errors += editdistance.eval(ref, hyp_w) + words += max(len(ref), 1) + return errors, words + + def training_step(self, batch: Tuple[Tensor, ...], batch_idx: int) -> Tensor: + wavs, wav_lens, tokens, tok_lens = batch + encoded, enc_lens = self._encode(wavs, wav_lens) + + if self.mode == "ctc": + log_probs = self.head(encoded) + loss = self._ctc_loss(log_probs, tokens, enc_lens, tok_lens) + else: + loss = self._rnnt_forward(encoded, enc_lens, tokens, tok_lens) + self.log("train/loss", loss, prog_bar=True) + + if ( + self.global_step % self._log_every == 0 + and self.global_step != self._last_log_step + ): + self._last_log_step = self.global_step + lr = self.trainer.optimizers[0].param_groups[0]["lr"] + self.log("train/lr", lr) + with torch.no_grad(): + res = self._decoding.decode( + self.head, encoded.detach().float(), enc_lens + ) + errs, wds = self._batch_wer([h[0] for h in res], tokens, tok_lens) + train_wer = errs / max(wds, 1) + self.log("train/wer", train_wer) + return loss + + def validation_step(self, batch: Tuple[Tensor, ...], batch_idx: int): + wavs, wav_lens, tokens, tok_lens = batch + encoded, enc_lens = self._encode(wavs, wav_lens) + + if self.mode == "ctc": + log_probs = self.head(encoded) + loss = self._ctc_loss(log_probs, tokens, enc_lens, tok_lens) + else: + loss = self._rnnt_forward(encoded, enc_lens, tokens, tok_lens) + + res = self._decoding.decode(self.head, encoded, enc_lens) + errs, wds = self._batch_wer([h[0] for h in res], tokens, tok_lens) + self._val_errors += errs + self._val_words += wds + self.log("val/loss", loss, sync_dist=True) + + def on_validation_epoch_end(self): + errs_t = torch.tensor(self._val_errors, device=self.device, dtype=torch.float64) + wds_t = torch.tensor(self._val_words, device=self.device, dtype=torch.float64) + if self.trainer.world_size > 1 and dist.is_initialized(): + dist.all_reduce(errs_t, op=dist.ReduceOp.SUM) + dist.all_reduce(wds_t, op=dist.ReduceOp.SUM) + self._val_errors = self._val_words = 0 + if wds_t.item() <= 0: + return + wer = (errs_t / wds_t).item() + # All ranks must log val_wer so ModelCheckpoint sees monitor on every process. + self.log("val/wer", wer, sync_dist=False) + self.log("val_wer", wer, logger=False, sync_dist=False) + if self.trainer.is_global_zero: + print( + f" [val] step={self.global_step} epoch={self.current_epoch} " + f"val/loss={self.trainer.callback_metrics.get('val/loss', 0):.6f} " + f"val/wer={wer:.4f}" + ) + + def configure_optimizers( + self, + ) -> Tuple[List[torch.optim.Optimizer], List[Dict[str, Any]]]: + opt = torch.optim.AdamW( + [p for p in self.parameters() if p.requires_grad], + lr=self._lr, + weight_decay=self._wd, + ) + total = self.trainer.estimated_stepping_batches + warmup = max(1, int(self._warmup_ratio * total)) + decay = max(1, total - warmup) + print(f" LR: {warmup} warmup + {decay} cosine = {total} steps") + + def lr_lambda(step): + if step < warmup: + return step / warmup + return max(0.0, 0.5 * (1 + math.cos(math.pi * (step - warmup) / decay))) + + sch = torch.optim.lr_scheduler.LambdaLR(opt, lr_lambda) + return [opt], [{"scheduler": sch, "interval": "step"}] diff --git a/train_utils/train.py b/train_utils/train.py new file mode 100644 index 0000000..dd3006b --- /dev/null +++ b/train_utils/train.py @@ -0,0 +1,212 @@ +"""Fine-tune a GigaAM pretrained model.""" + +import argparse +import warnings + +import pytorch_lightning as pl +import torch +from module import GigaAMFineTuner +from pytorch_lightning.callbacks import ModelCheckpoint +from pytorch_lightning.loggers import TensorBoardLogger +from torch.utils.data import DataLoader +from utils import ( + EpochTimeLogger, + StepProgressBar, + build_exp_name, + prepare_experiment_dirs, +) + +import gigaam +from gigaam.utils import AudioDataset + + +def parse_args() -> argparse.Namespace: + p = argparse.ArgumentParser() + p.add_argument("--model_name", required=True) + p.add_argument("--train_manifest", required=True) + p.add_argument("--val_manifest", required=True) + p.add_argument("--output_dir", default="./checkpoints") + p.add_argument("--exp_name", default=None) + p.add_argument("--batch_size", type=int, default=8) + p.add_argument("--eval_batch_size", type=int, default=32) + p.add_argument("--rnnt_subbatch_size", type=int, default=0) + p.add_argument("--num_workers", type=int, default=4) + p.add_argument("--lr", type=float, default=2e-5) + p.add_argument("--weight_decay", type=float, default=1e-2) + p.add_argument("--max_duration", type=float, default=20.0) + p.add_argument("--min_duration", type=float, default=0.1) + p.add_argument("--accumulate_grad_batches", type=int, default=1) + p.add_argument("--gradient_clip_val", type=float, default=1.0) + p.add_argument("--precision", default="32") + p.add_argument("--accelerator", default="auto") + p.add_argument("--devices", type=int, default=1) + p.add_argument("--seed", type=int, default=42) + p.add_argument("--activation_checkpointing", action="store_true") + p.add_argument("--freeze_encoder", action="store_true") + p.add_argument("--raw_text", action="store_true") + p.add_argument("--warmup_ratio", type=float, default=0.1) + p.add_argument("--max_epochs", type=int, default=None) + p.add_argument("--val_check_interval", type=float, default=1.0) + p.add_argument("--max_steps", type=int, default=None) + p.add_argument("--val_check_steps", type=int, default=None) + p.add_argument("--val_first_batches", type=int, default=None) + p.add_argument("--log_every_n_steps", type=int, default=25) + p.add_argument("--disable_tqdm", action="store_true", default=False) + p.add_argument("--skip_initial_validation", action="store_true", default=False) + p.add_argument("--save_top_k", type=int, default=2) + p.add_argument("--disable_spec_augment", action="store_true") + p.add_argument("--freq_masks", type=int, default=2) + p.add_argument("--freq_width", type=int, default=27) + p.add_argument("--time_masks", type=int, default=2) + p.add_argument("--time_width", type=int, default=20) + p.add_argument("--resume_from_checkpoint", type=str, default=None) + + args = p.parse_args() + assert (args.max_steps is not None) ^ ( + args.max_epochs is not None + ), "Either --max_steps or --max_epochs must be provided, but not both" + step_mode = args.max_steps is not None + if step_mode: + assert args.val_check_steps, "--max_steps requires --val_check_steps" + assert not ( + args.max_epochs is not None and args.val_check_steps is not None + ), "Use --val_check_interval instead of --val_check_steps with epoch mode" + return args + + +def main(): + args = parse_args() + pl.seed_everything(args.seed) + torch.set_float32_matmul_precision("high") + step_mode = args.max_steps is not None + exp_name = build_exp_name(args) + print(f"Experiment: {exp_name}") + model_dir, tb_dir = prepare_experiment_dirs(args.output_dir, exp_name) + + print(f"Loading pretrained {args.model_name} ...") + model = gigaam.load_model(args.model_name, fp16_encoder=False, device="cpu") + assert isinstance( + model, gigaam.GigaAMASR + ), "Fine-tuning expects an ASR model (GigaAMASR)" + + if args.activation_checkpointing: + model.encoder.activation_checkpointing = True + print("Encoder: activation checkpointing on (per Conformer layer)") + tokenizer = model.decoding.tokenizer + blank_id = model.decoding.blank_id + orig_model_name = model.cfg.model_name + is_e2e = "e2e" in orig_model_name + print( + f"Mode: {'rnnt' if 'rnnt' in orig_model_name else 'ctc'} | vocab={len(tokenizer)}, blank={blank_id}" + ) + + if args.raw_text and is_e2e: + raise ValueError("--raw_text is only for non-e2e models (charwise vocab)") + if not is_e2e and not args.raw_text: + warnings.warn( + "Non-e2e model without --raw_text: text won't be normalized. " + "Consider --raw_text to strip punctuation to vocab chars.", + stacklevel=1, + ) + + ds_kw = dict( + tokenizer=tokenizer, + max_duration=args.max_duration, + min_duration=args.min_duration, + raw_text=args.raw_text, + return_tokens=True, + ) + train_ds = AudioDataset(args.train_manifest, **ds_kw) + val_ds = AudioDataset(args.val_manifest, **ds_kw) + print(f"Train: {len(train_ds)} Val: {len(val_ds)}") + + dl_kw = dict(num_workers=args.num_workers, pin_memory=True) + train_dl = DataLoader( + train_ds, + batch_size=args.batch_size, + shuffle=True, + drop_last=True, + collate_fn=train_ds.collate_fn, + **dl_kw, + ) + val_dl = DataLoader( + val_ds, + batch_size=args.eval_batch_size, + shuffle=False, + collate_fn=val_ds.collate_fn, + **dl_kw, + ) + + lit = GigaAMFineTuner( + model=model, + blank_id=blank_id, + lr=args.lr, + freeze_encoder=args.freeze_encoder, + rnnt_subbatch_size=args.rnnt_subbatch_size, + weight_decay=args.weight_decay, + warmup_ratio=args.warmup_ratio, + log_every_n_steps=args.log_every_n_steps, + spec_augment=not args.disable_spec_augment, + freq_masks=args.freq_masks, + freq_width=args.freq_width, + time_masks=args.time_masks, + time_width=args.time_width, + cli_args=vars(args), + ) + + ckpt_cb = ModelCheckpoint( + dirpath=model_dir, + filename=f"gigaam-{args.model_name}-" + "{epoch:02d}-{step:06d}-{val_wer:.4f}", + monitor="val_wer", + mode="min", + save_top_k=max(1, args.save_top_k), + ) + + trainer_kw = dict( + accelerator=args.accelerator, + devices=args.devices, + strategy="ddp" if args.devices > 1 else "auto", + precision=args.precision, + accumulate_grad_batches=args.accumulate_grad_batches, + gradient_clip_val=args.gradient_clip_val, + callbacks=( + [ckpt_cb, EpochTimeLogger()] + + ( + [StepProgressBar(args.val_check_steps if step_mode else None)] + if not args.disable_tqdm + else [] + ) + ), + logger=TensorBoardLogger(save_dir=tb_dir, name=exp_name), + log_every_n_steps=args.log_every_n_steps, + num_sanity_val_steps=0, + default_root_dir=args.output_dir, + limit_val_batches=( + args.val_first_batches if args.val_first_batches is not None else 1.0 + ), + enable_progress_bar=not args.disable_tqdm, + ) + if step_mode: + trainer_kw.update( + max_steps=args.max_steps, + max_epochs=-1, + limit_train_batches=args.val_check_steps * args.accumulate_grad_batches, + ) + else: + trainer_kw.update( + max_epochs=args.max_epochs, val_check_interval=args.val_check_interval + ) + + trainer = pl.Trainer(**trainer_kw) + if not args.skip_initial_validation: + print("Running initial validation...") + trainer.validate(lit, val_dl) + trainer.fit(lit, train_dl, val_dl, ckpt_path=args.resume_from_checkpoint) + print(f"Best: {ckpt_cb.best_model_path}") + + if torch.distributed.is_available() and torch.distributed.is_initialized(): + torch.distributed.destroy_process_group() + + +if __name__ == "__main__": + main() diff --git a/train_utils/utils.py b/train_utils/utils.py new file mode 100644 index 0000000..43713e6 --- /dev/null +++ b/train_utils/utils.py @@ -0,0 +1,230 @@ +import os +import re +import time +import warnings +from concurrent.futures import ThreadPoolExecutor +from pathlib import Path +from typing import Any, Dict, List, Optional, Tuple + +import editdistance +import soundfile as sf +from pytorch_lightning.callbacks import Callback, TQDMProgressBar +from tqdm.auto import tqdm + +from gigaam.preprocess import SAMPLE_RATE + + +def normalize_raw_text(text: str) -> str: + text = text.replace("ё", "е").replace("Ё", "Е") + text = " ".join(text.split()) + return "".join( + c for c in text.lower() if ord("а") <= ord(c) <= ord("я") or c == " " + ) + + +def compute_wer(preds: List[Dict[str, Any]]) -> Tuple[float, float, int, int, int, int]: + """ + Word error rates for prediction rows with ``text`` (reference) and ``pred_text`` (hypothesis). + + Returns (wer_e2e_pct, wer_raw_pct, e2e_err, e2e_words, raw_err, raw_words). + """ + e2e_err = e2e_w = raw_err = raw_w = 0 + for r in preds: + ref_s, hyp_s = r["text"].strip(), r["pred_text"].strip() + rw, hw = ref_s.split(), hyp_s.split() + e2e_err += editdistance.eval(rw, hw) + e2e_w += len(rw) + nr, nh = normalize_raw_text(ref_s), normalize_raw_text(hyp_s) + rr, hh = nr.split(), nh.split() + raw_err += editdistance.eval(rr, hh) + raw_w += len(rr) + return ( + e2e_err / max(e2e_w, 1) * 100, + raw_err / max(raw_w, 1) * 100, + e2e_err, + e2e_w, + raw_err, + raw_w, + ) + + +def save_split(ds, split: str, out_dir: str, max_dur: float, workers: int) -> List[str]: + audio_dir = Path(out_dir) / "audio" / split + audio_dir.mkdir(parents=True, exist_ok=True) + n = len(ds) + w = max(1, min(workers, n)) + + def process_one(i: int) -> Optional[str]: + sample = ds[i] + text = sample["text"].strip() + arr, sr = sample["audio"]["array"], sample["audio"]["sampling_rate"] + if len(arr) / sr > max_dur: + return None + rel_p = f"audio/{split}/{i:06d}.wav" + p = audio_dir / f"{i:06d}.wav" + if not p.exists(): + sf.write(str(p), arr, sr) + return f"{rel_p}\t{len(arr) / sr:.3f}\t{text}" + + with ThreadPoolExecutor(max_workers=w) as ex: + lines = list( + tqdm( + ex.map(process_one, range(n)), + total=n, + desc=f"{split} ({n})", + ) + ) + return [ln for ln in lines if ln is not None] + + +def load_tonebooks(out_dir: str, max_duration: float = 30.0, workers: int = 8) -> Path: + """ + Load ToneBooks dataset and save manifests to the output directory. + Creates train / val .tsv files with the following columns: path, duration, transcription. + """ + + from datasets import Audio, load_dataset + + out = Path(out_dir) + out.mkdir(parents=True, exist_ok=True) + + print("Loading Vikhrmodels/ToneBooks...") + ds = load_dataset("Vikhrmodels/ToneBooks") + train, val = ds["train"], ds.get("validation") or ds.get("test") + train = train.cast_column("audio", Audio(sampling_rate=SAMPLE_RATE)) + val = val.cast_column("audio", Audio(sampling_rate=SAMPLE_RATE)) + print(f"Splits: {list(ds.keys())}, train={len(train)}, val={len(val)}") + + for fname, rows in ( + ( + "manifest_train.tsv", + save_split(train, "train", out_dir, max_duration, workers), + ), + ("manifest_val.tsv", save_split(val, "val", out_dir, max_duration, workers)), + ): + path = out / fname + path.write_text( + "path\tduration\ttranscription\n" + "\n".join(rows) + "\n", + encoding="utf-8", + ) + print(f" {path} ({len(rows)} samples)") + + print(f"\nDone! Manifests at {out}") + return out + + +class StepProgressBar(TQDMProgressBar): + + def __init__(self, steps_per_epoch: Optional[int] = None): + super().__init__() + self._steps_per_epoch = steps_per_epoch + + def on_train_epoch_start(self, trainer: Any, pl_module: Any): + super().on_train_epoch_start(trainer, pl_module) + acc = trainer.accumulate_grad_batches + if acc <= 1: + return + if self._steps_per_epoch: + remaining = ( + trainer.max_steps - trainer.global_step + if trainer.max_steps > 0 + else self._steps_per_epoch + ) + total = min(self._steps_per_epoch, remaining) + else: + total = (self.train_progress_bar.total or 0) // acc + self.train_progress_bar.reset(total=total) + + def on_train_batch_end( + self, trainer: Any, pl_module: Any, outputs: Any, batch: Any, batch_idx: int + ): + acc = trainer.accumulate_grad_batches + if acc <= 1: + super().on_train_batch_end(trainer, pl_module, outputs, batch, batch_idx) + return + if (batch_idx + 1) % acc == 0: + self.train_progress_bar.n = (batch_idx + 1) // acc + self.train_progress_bar.set_postfix(self.get_metrics(trainer, pl_module)) + self.train_progress_bar.refresh() + + +class EpochTimeLogger(Callback): + def on_train_epoch_start(self, trainer: Any, pl_module: Any): + self.start_time = time.time() + + def on_train_epoch_end(self, trainer: Any, pl_module: Any): + duration = time.time() - self.start_time + if trainer.is_global_zero: + print(f"[epoch {trainer.current_epoch}] time: {duration:.2f} sec") + + +def _fmt_float(v: float) -> str: + return f"{v:g}".replace("+0", "+").replace("-0", "-") + + +def _sanitize_name(name: str) -> str: + return re.sub(r"[^a-zA-Z0-9._-]+", "_", name).strip("._-") or "exp" + + +def build_exp_name(args) -> str: + """ + Create a unique experiment name based on the command line arguments. + Ignores arguments which do not affect training dynamics. + """ + if args.exp_name: + return _sanitize_name(args.exp_name) + parts = [f"{args.model_name.replace('_', '')}"] + parts += [f"lr{_fmt_float(args.lr)}", f"wd{_fmt_float(args.weight_decay)}"] + parts.append(f"b{args.batch_size}") + if args.accumulate_grad_batches > 1: + parts.append(f"agb{args.accumulate_grad_batches}") + if args.devices > 1: + parts.append(f"{args.devices}gpu") + if args.max_steps is not None: + parts.append(f"{args.max_steps}steps") + parts.append(f"vcs{args.val_check_steps}") + else: + parts.append(f"{args.max_epochs}ep") + if args.val_check_interval != 1.0: + parts.append(f"vci{_fmt_float(args.val_check_interval)}") + if args.warmup_ratio != 0.1: + parts.append(f"wmp{_fmt_float(args.warmup_ratio)}") + if args.freeze_encoder: + parts.append("frenc") + if args.activation_checkpointing: + parts.append("acckpt") + if args.val_first_batches is not None: + parts.append(f"vfb{args.val_first_batches}") + if args.raw_text: + parts.append("raw") + parts.append(f"dur{_fmt_float(args.min_duration)}-{_fmt_float(args.max_duration)}s") + if args.gradient_clip_val != 1.0: + parts.append(f"gc{_fmt_float(args.gradient_clip_val)}") + if args.precision != "16": + parts.append(f"pr-{str(args.precision).replace('-', '')}") + if args.seed != 42: + parts.append(f"seed{args.seed}") + if args.disable_spec_augment: + parts.append("nospecaug") + if not args.disable_spec_augment: + if args.freq_masks != 2: + parts.append(f"fm{args.freq_masks}") + if args.freq_width != 27: + parts.append(f"fw{args.freq_width}") + if args.time_masks != 2: + parts.append(f"tm{args.time_masks}") + if args.time_width != 20: + parts.append(f"tw{args.time_width}") + return _sanitize_name("_".join(parts)) + + +def prepare_experiment_dirs(output_dir: str, exp_name: str) -> Tuple[str, str]: + model_dir = os.path.join(output_dir, "models", exp_name) + tb_dir = os.path.join(output_dir, "tb_logs") + if os.path.isdir(model_dir) and os.listdir(model_dir): + warnings.warn( + f"Checkpoint dir is not empty: {model_dir}. Checkpoints may be overwritten.", + stacklevel=1, + ) + os.makedirs(model_dir, exist_ok=True) + return model_dir, tb_dir diff --git a/triton_scripts/Dockerfile b/triton_scripts/Dockerfile index 0f6d143..a278ec0 100644 --- a/triton_scripts/Dockerfile +++ b/triton_scripts/Dockerfile @@ -3,8 +3,9 @@ FROM nvcr.io/nvidia/tritonserver:24.10-py3 RUN pip install --no-cache-dir \ "torch>=2.6,<2.11" \ "torchaudio>=2.6,<2.11" \ + hydra-core==1.3.* \ + soundfile==0.13.* \ + omegaconf==2.3.* \ sentencepiece \ - omegaconf \ onnxruntime-gpu \ - tqdm \ - hydra-core + tqdm diff --git a/triton_scripts/README.md b/triton_scripts/README.md index b643d65..643fa39 100644 --- a/triton_scripts/README.md +++ b/triton_scripts/README.md @@ -72,11 +72,11 @@ python run_client.py ctc trt audio1.wav audio2.wav audio3.wav Forward pass time in seconds on CUDA for the first 4 segments from `long_example.wav` (VAD-segmented, ~65s total audio). For torch/onnx — both single-sample and batched inference are shown. -| Backend | v3_ctc | v3_e2e_rnnt | -|:--------------|:----------------|:----------------| -| triton/trt | 0.034 ± 0.000 | 0.403 ± 0.008 | -| triton/onnx | 0.046 ± 0.001 | 0.413 ± 0.005 | -| onnx (batch) | 0.037 ± 0.004 | 0.949 ± 0.017 | -| onnx | 0.047 ± 0.001 | 1.093 ± 0.045 | -| torch (batch) | 0.036 ± 0.002 | 0.919 ± 0.002 | -| torch | 0.112 ± 0.003 | 1.008 ± 0.001 | +| Backend | v3_ctc | v3_e2e_rnnt | +|:----------------|:----------------|:----------------| +| triton/trt | 0.034 ± 0.000 | 0.403 ± 0.008 | +| triton/onnx | 0.046 ± 0.001 | 0.434 ± 0.005 | +| onnx (bs=4) | 0.037 ± 0.004 | 0.470 ± 0.017 | +| onnx (b=1) | 0.047 ± 0.001 | 1.561 ± 0.014 | +| torch (bs=4) | 0.036 ± 0.002 | 0.539 ± 0.012 | +| torch (bs=1) | 0.112 ± 0.003 | 1.517 ± 0.028 | diff --git a/triton_scripts/repos/ctc_encoder_onnx/config.pbtxt b/triton_scripts/repos/ctc_encoder_onnx/config.pbtxt index 864d274..19b680a 100644 --- a/triton_scripts/repos/ctc_encoder_onnx/config.pbtxt +++ b/triton_scripts/repos/ctc_encoder_onnx/config.pbtxt @@ -5,7 +5,7 @@ max_batch_size: 0 input [ { name: "features" - data_type: TYPE_FP32 + data_type: TYPE_FP16 dims: [-1, 64, -1] }, { diff --git a/triton_scripts/repos/ctc_encoder_trt/config.pbtxt b/triton_scripts/repos/ctc_encoder_trt/config.pbtxt index 49f2a64..1753fdf 100644 --- a/triton_scripts/repos/ctc_encoder_trt/config.pbtxt +++ b/triton_scripts/repos/ctc_encoder_trt/config.pbtxt @@ -5,7 +5,7 @@ max_batch_size: 0 input [ { name: "features" - data_type: TYPE_FP32 + data_type: TYPE_FP16 dims: [-1, 64, -1] }, { diff --git a/triton_scripts/repos/ctc_postprocessing/1/model.py b/triton_scripts/repos/ctc_postprocessing/1/model.py index a0a3269..cb765de 100644 --- a/triton_scripts/repos/ctc_postprocessing/1/model.py +++ b/triton_scripts/repos/ctc_postprocessing/1/model.py @@ -3,9 +3,9 @@ import numpy as np import omegaconf -import torch -from gigaam.decoding import CTCGreedyDecoding +from gigaam.decoding import Tokenizer +from gigaam.onnx_utils import _decode_ctc_batch class TritonPythonModel: @@ -30,7 +30,7 @@ def initialize(self, args: Dict[str, Any]) -> None: else: tokenizer_path = None - self.decoding = CTCGreedyDecoding(vocabulary=vocab, model_path=tokenizer_path) + self.tokenizer = Tokenizer(vocab=vocab, model_path=tokenizer_path) def execute(self, requests: Any) -> List[Any]: import triton_python_backend_utils as pb_utils # type: ignore @@ -46,13 +46,11 @@ def execute(self, requests: Any) -> List[Any]: token_ids_np = token_ids.as_numpy() token_ids_lengths_np = token_ids_lengths.as_numpy() - results = self.decoding.decode( - head=None, - encoded=None, - lengths=torch.from_numpy(token_ids_lengths_np), - labels=torch.from_numpy(token_ids_np), + texts = _decode_ctc_batch( + token_ids_np, + token_ids_lengths_np, + self.tokenizer, ) - texts = [self.decoding.tokenizer.decode(result[0]) for result in results] texts_bytes = [text.encode("utf-8") for text in texts] texts_array = np.array(texts_bytes, dtype=object) diff --git a/triton_scripts/repos/gigaam_encoder_onnx/config.pbtxt b/triton_scripts/repos/gigaam_encoder_onnx/config.pbtxt index 47645f4..9ffc552 100644 --- a/triton_scripts/repos/gigaam_encoder_onnx/config.pbtxt +++ b/triton_scripts/repos/gigaam_encoder_onnx/config.pbtxt @@ -5,7 +5,7 @@ max_batch_size: 0 input [ { name: "audio_signal" - data_type: TYPE_FP32 + data_type: TYPE_FP16 dims: [-1, 64, -1] }, { @@ -18,7 +18,7 @@ input [ output [ { name: "encoded" - data_type: TYPE_FP32 + data_type: TYPE_FP16 dims: [-1, 768, -1] }, { diff --git a/triton_scripts/repos/gigaam_encoder_trt/config.pbtxt b/triton_scripts/repos/gigaam_encoder_trt/config.pbtxt index 2a746cb..73a0db9 100644 --- a/triton_scripts/repos/gigaam_encoder_trt/config.pbtxt +++ b/triton_scripts/repos/gigaam_encoder_trt/config.pbtxt @@ -5,7 +5,7 @@ max_batch_size: 0 input [ { name: "audio_signal" - data_type: TYPE_FP32 + data_type: TYPE_FP16 dims: [-1, 64, -1] }, { @@ -18,7 +18,7 @@ input [ output [ { name: "encoded" - data_type: TYPE_FP32 + data_type: TYPE_FP16 dims: [-1, 768, -1] }, { diff --git a/triton_scripts/repos/preprocessing/1/model.py b/triton_scripts/repos/preprocessing/1/model.py index c3ead81..86274c2 100644 --- a/triton_scripts/repos/preprocessing/1/model.py +++ b/triton_scripts/repos/preprocessing/1/model.py @@ -74,7 +74,7 @@ def execute(self, requests: Any) -> List[Any]: batch_audio, batch_lengths ) - features_np = features.detach().cpu().numpy().astype(np.float32) + features_np = features.detach().cpu().numpy().astype(np.float16) feature_lengths_np = feature_lengths.detach().cpu().numpy().astype(np.int64) output_tensors = [ diff --git a/triton_scripts/repos/preprocessing/config.pbtxt b/triton_scripts/repos/preprocessing/config.pbtxt index 01a5de2..2d5c542 100644 --- a/triton_scripts/repos/preprocessing/config.pbtxt +++ b/triton_scripts/repos/preprocessing/config.pbtxt @@ -18,7 +18,7 @@ input [ output [ { name: "features" - data_type: TYPE_FP32 + data_type: TYPE_FP16 dims: [-1, 64, -1] }, { @@ -29,4 +29,3 @@ output [ ] instance_group [{ kind: KIND_CPU }] - diff --git a/triton_scripts/repos/rnnt_postprocessing/1/model.py b/triton_scripts/repos/rnnt_postprocessing/1/model.py index cabaca4..cae1fcf 100644 --- a/triton_scripts/repos/rnnt_postprocessing/1/model.py +++ b/triton_scripts/repos/rnnt_postprocessing/1/model.py @@ -6,7 +6,7 @@ import onnxruntime as rt from gigaam.decoding import Tokenizer -from gigaam.onnx_utils import infer_onnx +from gigaam.onnx_utils import _decode_rnnt_batch class TritonPythonModel: @@ -70,8 +70,6 @@ def initialize(self, args: Dict[str, Any]) -> None: sess_options=opts, ) - self.pred_hidden = cfg.head.decoder.pred_hidden - def execute(self, requests: Any) -> List[Any]: import triton_python_backend_utils as pb_utils # type: ignore @@ -86,17 +84,13 @@ def execute(self, requests: Any) -> List[Any]: encoded_np = encoded.as_numpy() encoded_lengths_np = encoded_lengths.as_numpy() - texts = [ - infer_onnx( - wav_file=None, - model_cfg=self.cfg, - sessions=[None, self.pred_sess, self.joint_sess], - enc_features=encoded_np[i : i + 1, :, : encoded_lengths_np[i]], - tokenizer=self.tokenizer, - ) - for i in range(encoded_np.shape[0]) - ] - + texts = _decode_rnnt_batch( + encoded_np, + encoded_lengths_np, + self.cfg, + [None, self.pred_sess, self.joint_sess], + self.tokenizer, + ) texts_bytes = [text.encode("utf-8") for text in texts] texts_array = np.array(texts_bytes, dtype=object) diff --git a/triton_scripts/repos/rnnt_postprocessing/config.pbtxt b/triton_scripts/repos/rnnt_postprocessing/config.pbtxt index 01efa2f..ab99b17 100644 --- a/triton_scripts/repos/rnnt_postprocessing/config.pbtxt +++ b/triton_scripts/repos/rnnt_postprocessing/config.pbtxt @@ -5,7 +5,7 @@ max_batch_size: 0 input [ { name: "encoded" - data_type: TYPE_FP32 + data_type: TYPE_FP16 dims: [-1, 768, -1] }, { diff --git a/triton_scripts/run_convert_onnx.py b/triton_scripts/run_convert_onnx.py index 620b0a9..ebf35b9 100644 --- a/triton_scripts/run_convert_onnx.py +++ b/triton_scripts/run_convert_onnx.py @@ -6,6 +6,7 @@ from typing import Any, Tuple import omegaconf +import torch from torch import Tensor sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) @@ -13,6 +14,8 @@ import gigaam # noqa: E402 from gigaam.utils import onnx_converter # noqa: E402 +EXPORT_DTYPE = torch.float16 + # We need to override the forward method to return the argmax of the logits # to avoid variable output shapes in the ONNX / TRT and keep the unified config. @@ -25,30 +28,34 @@ def forward_for_export_with_argmax( return token_ids, encoded_len.long() -def _to_onnx_with_token_ids(self: Any, dir_path: str = ".") -> None: +def _to_onnx_with_token_ids( + self: Any, dir_path: str = ".", dtype: torch.dtype = torch.float32 +) -> None: """Convert to ONNX with token ids instead of logits.""" saved_forward = self.forward self.forward = self.forward_for_export try: - onnx_converter( - model_name="model", - out_dir=dir_path, - module=self, - inputs=self.encoder.input_example(), - input_names=["features", "feature_lengths"], - output_names=["token_ids", "token_ids_lengths"], - dynamic_axes={ - "features": {0: "batch_size", 2: "seq_len"}, - "feature_lengths": {0: "batch_size"}, - "token_ids": {0: "batch_size", 1: "seq_len"}, - "token_ids_lengths": {0: "batch_size"}, - }, - ) + with self.encoder.onnx_export_mode(): + onnx_converter( + model_name="model", + out_dir=dir_path, + module=self, + inputs=self.encoder.input_example(), + input_names=["features", "feature_lengths"], + output_names=["token_ids", "token_ids_lengths"], + dynamic_axes={ + "features": {0: "batch_size", 2: "seq_len"}, + "feature_lengths": {0: "batch_size"}, + "token_ids": {0: "batch_size", 1: "seq_len"}, + "token_ids_lengths": {0: "batch_size"}, + }, + export_dtype=dtype, + ) finally: self.forward = saved_forward -def convert_ctc(model: Any) -> tuple[str, str]: +def convert_ctc(model: Any, dtype: torch.dtype = torch.float16) -> tuple[str, str]: save_path = "repos/ctc_encoder_onnx/1" postprocessing_dir = "repos/ctc_postprocessing/1" @@ -59,7 +66,7 @@ def convert_ctc(model: Any) -> tuple[str, str]: model._to_onnx = types.MethodType(_to_onnx_with_token_ids, model) try: - model.to_onnx(save_path) + model.to_onnx(save_path, dtype=dtype) finally: model.forward_for_export = original_forward model._to_onnx = original_to_onnx @@ -67,12 +74,11 @@ def convert_ctc(model: Any) -> tuple[str, str]: return save_path, postprocessing_dir -def convert_rnnt(model: Any) -> tuple[str, str]: +def convert_rnnt(model: Any, dtype: torch.dtype = torch.float16) -> tuple[str, str]: save_path = "repos/gigaam_encoder_onnx/1" postprocessing_dir = "repos/rnnt_postprocessing/1" - # Save encoder, decoder and joint parts to onnx - model.to_onnx(save_path) + model.to_onnx(save_path, dtype=dtype) rename_onnx( f"{save_path}/{model.cfg.model_name}_encoder.onnx", @@ -132,10 +138,14 @@ def main() -> None: model_version = sys.argv[1] model = gigaam.load_model(model_version) + dtype = EXPORT_DTYPE + dtype_name = {torch.float16: "fp16", torch.float32: "fp32"}[dtype] + print(f"ONNX export dtype: {dtype_name}") + if "ctc" in model_version: - save_path, postprocessing_dir = convert_ctc(model) + save_path, postprocessing_dir = convert_ctc(model, dtype=dtype) else: - save_path, postprocessing_dir = convert_rnnt(model) + save_path, postprocessing_dir = convert_rnnt(model, dtype=dtype) # Save config and tokenizer for the postprocessing save_and_distribute_config(model, save_path, postprocessing_dir)