Fine-tuning, batched rnnt, fp16 onnx

.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

.gitignore

@@ -4,4 +4,6 @@ build
 *.wav
 .DS_Store
 *tmp*
-onnx
+onnx
+train_utils/data
+train_utils/checkpoints

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,24 +143,28 @@ 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
    onnx_dir = "onnx"
    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:
    ```python
    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).

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,15 +142,15 @@ 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
    onnx_dir = "onnx"
    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:

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}\")"

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)
 

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]:
         """