GigaAM

repository·main·Indexed 20 days ago

https://github.com/salute-developers/gigaam

A package for audio modeling and ASR featuring a family of open-source Conformer-based acoustic models. GigaAM supports Automatic Speech Recognition (ASR), emotion recognition, and audio embedding extraction, with specialized model lines for Russian (v1-v3) and multilingual support for over 70 languages. It provides tools for long-form transcription via pyannote.audio, model export to ONNX and TensorRT, and deployment via Triton Inference Server.

Tokens
23.9K
Snippets
78
Records
85
Agent score
72%

What's inside gigaam

  1. Overview of GigaAM model families

    main

    GigaAM is a family of acoustic models based on the Conformer architecture (220M–600M parameters). The models are pre-trained on diverse speech data and can be used for Automatic Speech Recognition (ASR) or emotion recognition.

    There are four main model lineages:

    • v1: Pre-trained using Wav2vec 2.0 (50k hours). Available versions: v1_ssl, emo, v1_ctc, v1_rnnt.
    • v2: Pre-trained using HuBERT–CTC (50k hours). Available versions: v2_ssl, v2_ctc, v2_rnnt.
    • v3: Pre-trained using HuBERT–CTC (700k hours). Available versions: v3_ssl, v3_ctc, v3_rnnt, v3_e2e_ctc, v3_e2e_rnnt. Note that v3_e2e_ctc and v3_e2e_rnnt support punctuation and text normalization.
    • multilingual: Pre-trained using HuBERT-style methods (2M hours). Available versions: multilingual_ssl, multilingual_large_ssl, multilingual_ctc, multilingual_large_ctc. These provide SSL encoders (220M/600M parameters) and character-level CTC ASR models supporting 70+ languages.
  2. Optimize Attention mechanism for CUDA

    main

    When running GigaAM on CUDA, you can choose between different attention implementations.

    • SDPA (Scaled Dot Product Attention): Enabled by default. It offers better robustness across varying input dimensions and is generally efficient for standard batch sizes and sequence lengths.
    • Flash Attention (flash_attn): Only provides significant benefits for long sequences when the batch_size is much greater than 1 (batch_size >> 1).
    • Custom Implementation: Uses naive tensor multiplication.

    Performance Summary:

    • For small batches (e.g., bs=1), SDPA and Custom are comparable, while Flash Attention may be slightly slower.
    • For large batches (e.g., bs=128), Flash Attention is significantly faster than both Custom and SDPA implementations for both Attention Mechanism and Full Encoder inference.
    * Only Attention Mechanism (Time (ms) ± std) |
    | bs, seq_len |        Custom |        SDPA |        Flash |
    |-------------|---------------|----------|-------------|
    | 1, 10s      | 0.03 ± 0.00   | 0.03 ± 0.00 | 0.05 ± 0.03  |
    | 8, 20s      | 0.15 ± 0.01   | 0.14 ± 0.01 | 0.66 ± 0.14  |
    | 128, 30s    | 3.60 ± 0.10   | 3.59 ± 0.04 | 1.40 ± 0.06  |
    
    * Full Encoder Inference (Time (ms) ± std) |
    | bs, seq_len |        Custom |        SDPA |        Flash |
    |-------------|---------------|----------|-------------|
    | 1, 10s      | 10.14 ± 0.17  | 10.06 ± 0.12 | 11.57 ± 0.25 |
    | 8, 20s      | 15.84 ± 0.07  | 15.90 ± 0.02 | 25.26 ± 0.26 |
    | 128, 30s    | 324.53 ± 0.17 | 324.48 ± 0.09| 293.80 ± 0.89|
  3. Prepare TSV manifests for training and evaluation

    main

    GigaAM uses TSV (tab-separated values) manifests. Each row must contain a path and duration. Optionally, you can include a transcription column. Paths can be absolute or relative to the manifest directory.

    path	audio/0001.wav	duration	3.21	transcription	привет как дела
  4. Understand End-to-End ASR performance and label quality

    main

    GigaAM's end-to-end models (v3_e2e_ctc and v3_e2e_rnnt) are evaluated against Whisper using an LLM-as-a-Judge.

    A critical factor in model performance is the quality of training labels (punctuation and text normalization).

    • GigaChat Max Audio: Provides full normalization and precise punctuation, leading to high F1 scores for punctuation and lower WER/CER.
    • Whisper Punctuator: Only adds punctuation without normalization, which can lead to poor comma accuracy.
    • Whisper labels: Using Whisper in forced decoding mode for training results in lower F1 scores for punctuation and higher error rates compared to using GigaChat Max Audio labels.

    Metrics used:

    • F1(,), F1(.), F1(?): F1 scores for commas, periods, and question marks.
    • WER: Word Error Rate.
    • CER: Character Error Rate.
    |                                   | Normalization     | F1(,) ↑ |   F1(.) ↑ |   F1(?) ↑ |   WER ↓ |   CER ↓ |
    |:----------------------------------|------------------|----------:|----------:|----------:|--------:|--------:|
    | GigaChat Max Audio                | Full              |      84.2 |      85.6 |      74.9 |    18.4 |    10.9 |
    | Whisper Punctuator                | punctuation only  |      62.2 |      85.0 |      77.7 |     0.0 |     0.0 |
    | GigaAM from Whisper labels        | punctuation only  |      50.3 |      84.1 |      77.7 |    12.0 |     7.8 |
    | GigaAM-e2e-ctc                    | Full              |      83.7 |      86.7 |      78.6 |    16.0 |     8.7 |
    | GigaAM-e2e-rnnt                   | Full              |      84.5 |      86.7 |      79.8 |    14.2 |     8.8 |
  5. Set up Triton Inference Server for GigaAM

    main

    GigaAM models can be deployed using a Triton Inference Server ensemble. The setup process involves building a custom Docker image, converting models to ONNX and TensorRT formats, and running the Triton server.

    Inference Logic:

    • CTC models: Converted entirely to ONNX/TensorRT.
    • RNNT models: Split into an encoder (ONNX/TRT) and decoder/joint components that run in Python using onnxruntime.

    Workflow Summary:

    1. Build the gigaam-triton Docker image.
    2. Convert models to ONNX using run_convert_onnx.py.
    3. Convert ONNX to TensorRT using run_convert_trt.sh inside an NVIDIA TensorRT container.
    4. Start the Triton server using the built image.
    5. Run the client to perform inference.
    # 0. Build Docker Image
    docker build -t gigaam-triton .
  6. Evaluate a model with `eval.py`

    main

    Use eval.py to evaluate either a fine-tuned checkpoint or a pretrained GigaAM model. The script outputs preds.jsonl and prints the Word Error Rate (WER).

    Predictions are saved in predictions/<manifest_stem>/<exp_name>/step_<step>/preds.jsonl (the step_<step> part is omitted for pretrained models) in the same directory as your manifest. WER is reported on both original transcripts and raw texts.

    # Evaluate a fine-tuned checkpoint
    python eval.py \
        --checkpoint ./checkpoints/models/<exp_name>/gigaam-*.ckpt \
        --eval_manifest /path/to/manifest.tsv
    
    # Evaluate a pretrained GigaAM model
    python eval.py --model_name v3_e2e_ctc --eval_manifest /path/to/manifest.tsv
  7. Export and use GigaAM models with ONNX

    main

    GigaAM supports exporting models to ONNX format for optimized inference.

    Exporting: Use model.to_onnx(dir_path=..., dtype=...). While torch.float32 is the default, using torch.float16 is recommended for GPU inference to increase speed and reduce VRAM usage.

    Running Inference: Use gigaam.onnx_utils to load and run the exported models. You can run inference on single audio files or entire datasets (via a .tsv manifest).

    GPU Acceleration: To use the GPU with ONNX, uninstall onnxruntime and install onnxruntime-gpu==1.22.*.

    # 1. Export
    import torch
    import gigaam
    
    model_version = "v3_ctc"
    model = gigaam.load_model(model_version)
    model.to_onnx(dir_path="onnx", dtype=torch.float16)
    
    # 2. Inference
    from gigaam.onnx_utils import load_onnx, infer_onnx
    
    sessions, model_cfg = load_onnx("onnx", model_version)
    result = infer_onnx([audio_path], model_cfg, sessions)
    print(result[0])
  8. Transcribe long-form audio

    main

    For audio longer than 25 seconds, use the transcribe_longform method.

    Setup Requirements:

    1. Install longform dependencies: pip install -e ".[longform]".
    2. Generate a Hugging Face API token.
    3. Accept the terms for pyannote/segmentation-3.0 on Hugging Face.
    4. Set the HF_TOKEN environment variable.

    The method returns an iterable of segments, where each segment has start, end, and text attributes.

    import os
    import gigaam
    
    # Ensure HF_TOKEN is set for pyannote access
    os.environ["HF_TOKEN"] = "<YOUR_HF_TOKEN>"
    
    result = model.transcribe_longform(long_audio_path)
    for segment in result:
       print(f"[{gigaam.format_time(segment.start)} - {gigaam.format_time(segment.end)}]: {segment.text}")
  9. Convert GigaAM models to ONNX

    main

    Use the run_convert_onnx.py script to create .onnx checkpoints and configuration files.

    Important Note on Preprocessing: The script saves model configs to the preprocessing directory. For v3 family models, preprocessing differs from earlier versions. Because Triton uses a shared preprocessing model, you can only use models with the same preprocessing simultaneously (either all v3 models or all earlier models). The preprocessing is determined by the last model converted to ONNX.

    python run_convert_onnx.py <model_version>  # e.g., v3_ctc, v3_e2e_rnnt
  10. Start the Triton Inference Server

    main

    Run the Triton Inference Server using the gigaam-triton Docker image. The command mounts the model repository and the GigaAM repository to ensure Python backend models (like rnnt_postprocessing) can import the gigaam package.

    Configuration Note: By default, ONNX models use instance_group [{ kind: KIND_GPU }]. To switch to CPU execution, you must manually update the instance_group to KIND_CPU in the following config files:

    • repos/ctc/config.pbtxt
    • repos/gigaam_encoder_onnx/config.pbtxt
    docker run --gpus all --ipc=host -p 8000:8000 -p 8001:8001 -p 8002:8002 \
      -v "$(pwd)/repos:/models" \
      -v "$(pwd)/..:/opt/gigaam_repo" \
      -e PYTHONPATH=/opt/gigaam_repo \
      gigaam-triton \
      tritonserver --model-repository=/models --exit-on-error=false
  11. Fine-tune a new language from an SSL backbone

    main

    To adapt a self-supervised learning (SSL) backbone (e.g., multilingual_ssl) to a new language, use train.py with the SSL model name and specify a head using --head ctc (default) or --head rnnt.

    There are two ways to handle vocabulary:

    1. Derived from manifest: Use --build_vocab_from_manifest and --raw_text (recommended). This normalizes text (lowercase, no punctuation) and builds the alphabet from your training data. You can save this via --save_vocab.
    2. Explicit vocabulary file: Provide a .json list via --vocab.

    Using --raw_text is recommended as it normalizes labels to lowercase and strips punctuation.

    # Example: CTC head with vocabulary derived from manifest
    python train.py \
        --model_name multilingual_ssl \
        --head ctc \
        --train_manifest /path/to/ka_train.tsv \
        --val_manifest /path/to/ka_val.tsv \
        --build_vocab_from_manifest \
        --save_vocab ./ka_vocab.json \
        --raw_text \
        --max_epochs 5 \
        --val_check_interval 0.5 \
        --batch_size 16 \
        --lr 1e-4 \
        --activation_checkpointing
  12. Convert ONNX models to TensorRT

    main

    To convert the models to TensorRT format, you must run the conversion script inside an NVIDIA TensorRT Docker container. This process converts the version of the corresponding CTC or RNNT model that was last converted to ONNX.

    # Run the TensorRT container
    docker run --gpus all -it --rm -v $(pwd):/workspace nvcr.io/nvidia/tensorrt:24.10-py3
    
    # Inside the container, run the conversion script
    bash run_convert_trt.sh <ctc | rnnt>