Word-level timestamps

[mirix]

Is it possible to obtain word-level timestamps?

[georgygospodinov]

Hello!

Right now the GigaAM API does not provide word-level timestamps directly.

However, both CTC and RNN-T models learn an alignment between the input acoustic features and the output symbol sequence. This alignment can be used to reconstruct timestamps from the encoder frames on which tokens are emitted, and then aggregate those tokens into words.

Below is a minimal example of how you can do this for a char-level RNN-T model (v3_rnnt). It performs greedy RNN-T decoding while keeping track of the encoder frame index for every emitted token, and then groups characters into word segments with start/end times:

from typing import Dict, List, Tuple, Union

import gigaam
import torch
from gigaam.preprocess import SAMPLE_RATE


def _decode_with_alignment(
    head,
    encoded_seq: torch.Tensor,
    seq_len: int,
    blank_id: int,
    max_symbols: int,
) -> Tuple[List[int], List[int]]:
    """
    Greedy RNNT decoding that also keeps the encoder frame index for every token.
    """
    hyp: List[int] = []
    token_frames: List[int] = []
    dec_state = None
    last_label = None

    for t in range(seq_len):
        encoder_step = encoded_seq[t, :, :].unsqueeze(1)
        not_blank = True
        emitted = 0
        while not_blank and emitted < max_symbols:
            decoder_step, hidden = head.decoder.predict(last_label, dec_state)
            joint_logp = head.joint.joint(encoder_step, decoder_step)[0, 0, 0, :]
            k = int(torch.argmax(joint_logp).item())
            if k == blank_id:
                not_blank = False
                continue
            hyp.append(k)
            token_frames.append(t)
            dec_state = hidden
            last_label = torch.tensor(
                [[k]], dtype=torch.long, device=encoded_seq.device
            )
            emitted += 1

    return hyp, token_frames


def _chars_to_words(
    chars: List[str],
    frames: List[int],
    frame_shift: float,
) -> List[Dict[str, Union[str, float]]]:
    words: List[Dict[str, Union[str, float]]] = []
    current_chars: List[str] = []
    current_frames: List[int] = []

    def collect_word():
        if not current_chars:
            return
        text = "".join(current_chars).strip()
        if not text:
            current_chars.clear()
            current_frames.clear()
            return
        start = current_frames[0] * frame_shift
        end = (current_frames[-1] + 1) * frame_shift
        words.append({"word": text, "start": start, "end": end})
        current_chars.clear()
        current_frames.clear()

    for char, frame in zip(chars, frames):
        if char == " ":
            collect_word()
            continue
        current_chars.append(char)
        current_frames.append(frame)

    collect_word()
    return words


@torch.inference_mode()
def extract_word_timestamps(
    model,
    audio_path: str,
    max_symbols_per_step: int = 3,
) -> Dict[str, object]:
    wav, length = model.prepare_wav(audio_path)
    encoded, encoded_len = model.forward(wav, length)

    encoded = encoded.transpose(1, 2)
    seq = encoded[0, :, :].unsqueeze(1)
    seq_len = int(encoded_len[0].item())
    frame_shift = int(length[0].item()) / SAMPLE_RATE / seq_len

    blank_id = model.decoding.blank_id

    token_ids, token_frames = _decode_with_alignment(
        model.head, seq, seq_len, blank_id, max_symbols_per_step
    )

    tokenizer = model.decoding.tokenizer
    transcript = tokenizer.decode(token_ids)
    chars = [tokenizer.vocab[idx] for idx in token_ids]
    word_segments = _chars_to_words(chars, token_frames, frame_shift)

    return {
        "transcript": transcript.strip(),
        "frame_shift": frame_shift,
        "words": word_segments,
    }


if __name__ == "__main__":
    model = gigaam.load_model("v3_rnnt", device="cpu")
    result = extract_word_timestamps(model, audio_path="example.wav")

    for word in result["words"]:
        print(f"{word['start']:.2f} - {word['end']:.2f}\t{word['word']}")

[Maksim-Goncharovskiy]

Right now the GigaAM API does not provide word-level timestamps directly.

However, both CTC and RNN-T models learn an alignment between the input acoustic features and the output symbol sequence. This alignment can be used to reconstruct timestamps from the encoder frames on which tokens are emitted, and then aggregate those tokens into words.

Below is a minimal example of how you can do this for a char-level RNN-T model (v3_rnnt). It performs greedy RNN-T decoding while keeping track of the encoder frame index for every emitted token, and then groups characters into word segments with start/end times:

Do you have a similar code for the ctc model? And is it possible for the new e2e models? :)

[Maksim-Goncharovskiy]

For the word timestamps of the ctc models (including v3_e2e_ctc model), the code below can be used (but transcription is a bit different from the transcribe_longform() method):

from typing import Dict, List, Tuple, Union

import gigaam
import torch
from gigaam.preprocess import SAMPLE_RATE


def _decode_with_alignment_ctc(
    head,
    encoded_seq: torch.Tensor,
    seq_len: int,
    blank_id: int,
) -> Tuple[List[int], List[int]]:
    """
    Greedy CTC decoding
    """
    log_probs = head(encoder_output=encoded_seq)  # [B, T, C]
    labels = log_probs.argmax(dim=-1, keepdim=False)  # [B, T]
    
    frame_labels = labels[0, :seq_len].cpu().tolist()  # [T]
    
    hyp: List[int] = []
    token_frames: List[int] = []
    
    prev_label = blank_id
    for t, label in enumerate(frame_labels):
        # skip blanks and consecutive duplicates
        if label != blank_id and (label != prev_label or prev_label == blank_id):
            hyp.append(int(label))
            token_frames.append(t)
        prev_label = label
    
    return hyp, token_frames


def _get_token_string(tokenizer, token_id: int) -> str:
    if tokenizer.charwise:
        return tokenizer.vocab[token_id]
    else:
        return tokenizer.model.IdToPiece(token_id)


def _chars_to_words(
    chars: List[str],
    frames: List[int],
    frame_shift: float,
) -> List[Dict[str, Union[str, float]]]:
    """
    Convert character-level tokens with frame indices to word-level timestamps.
    For SentencePiece models, chars may contain subword pieces.
    """
    words: List[Dict[str, Union[str, float]]] = []
    current_chars: List[str] = []
    current_frames: List[int] = []

    def collect_word():
        if not current_chars:
            return
        text = "".join(current_chars).strip()
        if not text:
            current_chars.clear()
            current_frames.clear()
            return
        start = current_frames[0] * frame_shift
        end = (current_frames[-1] + 1) * frame_shift
        words.append({"word": text, "start": start, "end": end})
        current_chars.clear()
        current_frames.clear()

    for char, frame in zip(chars, frames):
        if char == " " or (char.startswith("▁") and len(current_chars) > 0):
            collect_word()
            if char.startswith("▁"):
                char = char[1:]
            if char:  
                current_chars.append(char)
                current_frames.append(frame)
            continue

        if char.startswith("▁"):
            collect_word()
            char = char[1:]
        current_chars.append(char)
        current_frames.append(frame)

    collect_word()
    return words


@torch.inference_mode()
def extract_word_timestamps(
    model,
    audio_path: str,
    max_symbols_per_step: int = 3,  
) -> Dict[str, object]:
    wav, length = model.prepare_wav(audio_path)
    encoded, encoded_len = model.forward(wav, length)

    seq_len = int(encoded_len[0].item())
    frame_shift = int(length[0].item()) / SAMPLE_RATE / seq_len

    blank_id = model.decoding.blank_id

    token_ids, token_frames = _decode_with_alignment_ctc(
        model.head, encoded, seq_len, blank_id
    )

    tokenizer = model.decoding.tokenizer
    transcript = tokenizer.decode(token_ids)
    
    chars = [_get_token_string(tokenizer, idx) for idx in token_ids]
    word_segments = _chars_to_words(chars, token_frames, frame_shift)

    return {
        "transcript": transcript.strip(),
        "frame_shift": frame_shift,
        "words": word_segments,
    }

Note, that it does not work for long audios, so in these cases you need to implement your own transcribe_longform method.

[mirix]

Thanks. Could you also post the code for v3_e2e_rnnt with word-level timestamps?

[Maksim-Goncharovskiy]

Update:
Now you can use my fork (https://github.com/Maksim-Goncharovskiy/GigaAM) with word level timestamps functionality for all asr gigaam models via transcribe_longform method. All you need is to pass word_timestamps=True argument:

import gigaam
e2e_ctc = gigaam.load_model("v3_e2e_ctc", device="cpu")

transcription: list[dict] = e2e_ctc.transcribe_longform('example.wav', word_timestamps=True)

"""
transcription format:
[{
    "text": str
    "start": float
    "end": float
}, ...]
"""

Transcriptions with word timestamps same as transcriptions without them.

[mirix]

Is this OK?

from typing import Dict, List, Tuple, Union
import torch
import gigaam
from gigaam.preprocess import SAMPLE_RATE

################################################################################
# ------------------------- R N N T   D E C O D E R ---------------------------
################################################################################

def _decode_with_alignment_rnnt(
    head,
    encoded_seq: torch.Tensor,
    seq_len: int,
    blank_id: int,
    max_symbols: int,
) -> Tuple[List[int], List[int]]:
    """
    Greedy RNNT decoding that also keeps the encoder frame index for every token.
    Compatible with GIGAAM v3_e2e_rnnt models.
    """
    hyp: List[int] = []
    token_frames: List[int] = []
    dec_state = None
    last_label = None

    # encoded_seq must be [T, 1, D]
    for t in range(seq_len):
        encoder_step = encoded_seq[t, :, :].unsqueeze(1)
        not_blank = True
        emitted = 0

        while not_blank and emitted < max_symbols:
            decoder_step, hidden = head.decoder.predict(last_label, dec_state)
            joint = head.joint.joint(encoder_step, decoder_step)[0, 0, 0, :]
            k = int(torch.argmax(joint).item())

            if k == blank_id:
                not_blank = False
                continue

            hyp.append(k)
            token_frames.append(t)

            last_label = torch.tensor([[k]], dtype=torch.long, device=encoded_seq.device)
            dec_state = hidden
            emitted += 1

    return hyp, token_frames


################################################################################
# ---------------------------- C T C   D E C O D E R ---------------------------
################################################################################

def _decode_with_alignment_ctc(
    head,
    encoded_seq: torch.Tensor,
    seq_len: int,
    blank_id: int,
):
    """
    Greedy CTC decoding with frame alignment.
    Works with v3_e2e_ctc.
    """
    log_probs = head(encoder_output=encoded_seq)
    labels = log_probs.argmax(dim=-1)  # [B, T]

    frame_labels = labels[0, :seq_len].cpu().tolist()

    hyp, token_frames = [], []
    prev = blank_id

    for t, lbl in enumerate(frame_labels):
        if lbl != blank_id and (lbl != prev or prev == blank_id):
            hyp.append(lbl)
            token_frames.append(t)
        prev = lbl

    return hyp, token_frames


################################################################################
# -------------------------- T O K E N   H E L P E R ---------------------------
################################################################################

def _get_token_str(tokenizer, token_id: int) -> str:
    """Works for charwise and sentencepiece vocabularies."""
    if tokenizer.charwise:
        return tokenizer.vocab[token_id]
    else:
        return tokenizer.model.IdToPiece(token_id)


################################################################################
# --------------------------- W O R D   G R O U P I N G ------------------------
################################################################################

def _chars_to_words(chars, frames, frame_shift):
    words = []
    current_chars, current_frames = [], []

    def push():
        if not current_chars:
            return
        text = "".join(current_chars).strip()
        if text:
            words.append({
                "word": text,
                "start": current_frames[0] * frame_shift,
                "end": (current_frames[-1] + 1) * frame_shift,
            })
        current_chars.clear()
        current_frames.clear()

    for c, f in zip(chars, frames):
        # SentencePiece "▁" marks word boundaries
        if c == " " or c.startswith("▁"):
            push()
            c = c.lstrip("▁")
            if c:
                current_chars.append(c)
                current_frames.append(f)
            continue

        current_chars.append(c)
        current_frames.append(f)

    push()
    return words


################################################################################
# ------------------------------ M A I N   A P I -------------------------------
################################################################################

@torch.inference_mode()
def extract_word_timestamps(
    model,
    audio_path: str,
    max_symbols_per_step: int = 3,
):
    wav, length = model.prepare_wav(audio_path)
    encoded, encoded_len = model.forward(wav, length)

    seq_len = int(encoded_len[0].item())
    frame_shift = float(length[0].item()) / SAMPLE_RATE / seq_len

    blank_id = model.decoding.blank_id
    tokenizer = model.decoding.tokenizer
    head = model.head

    # ----------------------------------------------------------
    #   FIXED: Detect RNNT vs CTC for all GigaAM models
    # ----------------------------------------------------------
    is_rnnt = hasattr(head, "decoder") and hasattr(head, "joint")
    is_ctc  = not is_rnnt
    # ----------------------------------------------------------

    if is_rnnt:
        # e2e RNNT requires [T, B, D]
        encoded_seq = encoded.transpose(1, 2)[0, :, :].unsqueeze(1)

        token_ids, token_frames = _decode_with_alignment_rnnt(
            head,
            encoded_seq,
            seq_len,
            blank_id,
            max_symbols_per_step
        )

    else:
        # CTC uses [B, T, D]
        token_ids, token_frames = _decode_with_alignment_ctc(
            head,
            encoded,
            seq_len,
            blank_id
        )

    chars = [_get_token_str(tokenizer, i) for i in token_ids]
    word_segments = _chars_to_words(chars, token_frames, frame_shift)

    return {
        "transcript": tokenizer.decode(token_ids).strip(),
        "frame_shift": frame_shift,
        "words": word_segments,
    }


################################################################################
# ------------------------------ E X A M P L E --------------------------------
################################################################################

if __name__ == "__main__":
    model = gigaam.load_model("v3_e2e_rnnt", device="cuda")
    res = extract_word_timestamps(model, "example.wav")

    for w in res["words"]:
        print(f"{w['start']:.2f} – {w['end']:.2f}\t{w['word']}")

[mirix]

I have finished migrating the automated Russian transcription and translation pipeline to GigaAM.

A few examples can be found here:

https://rumble.com/c/c-7816716?e9s=src_v1_cbl

The code is here (there is a README now):

https://github.com/mirix/autosubs/tree/main/rus

Word precision has improved substantially.

However, GigaAM does not seem to have been trained on songs and still struggles with overlapping music and voice, overlapping voices and polyphony.

Thanks for the great work and the help!

[hadzhykhanov-dev]

Are word-level timestamps supported when using huggingface models?

[macsunmood]

Are word-level timestamps supported when using huggingface models?

Same question here. Doesn't work for me now..