It is possible with all ASR model types (CTC and LLM), although CTC is far easier to implement.

Here's a quick shim for you, though: save the file to your local directory, and run the code snippet. If you're using a CTC model, run pip install ctc-segmentation.

Note: This is just a proof-of-concept, and it only works properly for CTC models right now. The word-level timestamps are segment-level for non-whitespace languages, and the character-level timestamps do not associate with real morphological boundaries. For word segmentation, you can use pyicu or another library, but omnilingual tokenization is hard to come by, programmatically. The monkey-patch also includes automatic chunking.

omni_ts_patch.py

# if you're in Colab, uncomment the following line and run as a cell:
#%%writefile omni_ts_patch.py

# omni_ts_patch.py
# unified monkey-patch for omnilingual asr ctc & llm timestamping.
# - ctc models: use argmax-ctc token transition frames.
# - llm models: use energy-cdf alignment over audio as a robust heuristic.
# whitespace vs. non-whitespace segmentation is determined automatically by
# analyzing unicode character blocks in the transcribed text.
#
# import this once before creating asrinferencepipeline.

from __future__ import annotations

import re
import types
import warnings
from dataclasses import dataclass
from typing import Any, Dict, List, Optional, Sequence, Tuple, Union

import torch
from torch import tensor, Tensor

# suppress torchaudio deprecation warnings from the backend
warnings.filterwarnings("ignore", category=UserWarning, module="torchaudio")
warnings.filterwarnings(
    "ignore", category=UserWarning, module="torchaudio.backend.common"
)
warnings.filterwarnings(
    "ignore", category=UserWarning, module="torchaudio.backend.utils"
)

try:
    from fairseq2.nn import BatchLayout

    HAS_FAIRSEQ2_BATCHLAYOUT = True
except Exception:
    HAS_FAIRSEQ2_BATCHLAYOUT = False

try:
    from omnilingual_asr.models.inference.pipeline import ASRInferencePipeline
except Exception as e:
    raise RuntimeError(
        "failed to import asrinferencepipeline from "
        "omnilingual_asr.models.inference.pipeline. make sure the "
        "omnilingual-asr package is installed."
    ) from e


@dataclass
class WordTimestamp:
    word: str
    start: float
    end: float
    score: Optional[float] = None


def chunk_waveform(
    waveform: Tensor, sample_rate: int, max_duration_sec: float
) -> List[Tuple[Tensor, float]]:
    if waveform.ndim != 1:
        raise ValueError("waveform must be 1d")
    chunk_samples = int(max_duration_sec * sample_rate)
    T = waveform.shape[0]
    chunks: List[Tuple[Tensor, float]] = []
    for s in range(0, T, chunk_samples):
        e = min(T, s + chunk_samples)
        chunks.append((waveform[s:e], s / float(sample_rate)))
    return chunks


def _make_seqs_layout(x_batched: Tensor, seq_lens: Sequence[int]) -> Any:
    if HAS_FAIRSEQ2_BATCHLAYOUT:
        return BatchLayout.of(batch=x_batched, seq_lens=list(seq_lens))

    class _FakeBatchLayout:
        def __init__(self, seq_lens: Sequence[int]):
            self.seq_lens = [int(x) for x in seq_lens]
            self.batch_size = len(self.seq_lens)
            self.max_len = int(max(self.seq_lens)) if self.seq_lens else 0
            self.total_len = int(sum(self.seq_lens))
            self.packed = False
            self.padded = True

        def to(self, device):
            return self

    warnings.warn(
        "fairseq2.nn.batchlayout not available; using minimal fakebatchlayout."
    )
    return _FakeBatchLayout(seq_lens)


def _ensure_feats_and_layout(
    maybe_ret: Any, default_layout: Any
) -> Tuple[Tensor, Any]:
    if isinstance(maybe_ret, (tuple, list)):
        if len(maybe_ret) >= 2:
            return maybe_ret[0], maybe_ret[1]
        else:
            return maybe_ret[0], default_layout
    else:
        return maybe_ret, default_layout


def _ensure_tensor(maybe_ret: Any) -> Tensor:
    return maybe_ret[0] if isinstance(maybe_ret, (tuple, list)) else maybe_ret


def _get_frontend_and_dtype(model: Any):
    if hasattr(model, "encoder_frontend"):
        frontend = getattr(model, "encoder_frontend")
    elif hasattr(model, "frontend"):
        frontend = getattr(model, "frontend")
    else:
        raise RuntimeError("model has no encoder_frontend or frontend.")
    try:
        f_dtype = next(frontend.parameters()).dtype
    except Exception:
        f_dtype = next(model.parameters()).dtype
    return frontend, f_dtype


def get_encoder_stride_seconds(
    model: Any, waveform: Tensor, sample_rate: int
) -> Tuple[int, float]:
    device = next(model.parameters()).device
    frontend, f_dtype = _get_frontend_and_dtype(model)
    with torch.no_grad():
        x = waveform.to(device=device, dtype=f_dtype).unsqueeze(0)
        seqs_layout = _make_seqs_layout(x, [waveform.shape[0]])
        ret = frontend.extract_features(x, seqs_layout)
        feats, _ = _ensure_feats_and_layout(ret, seqs_layout)
    s_frames = int(feats.shape[1])
    t_samples = int(waveform.shape[0])
    stride_samples = max(1, round(t_samples / max(1, s_frames)))
    stride_seconds = stride_samples / float(sample_rate)
    return stride_samples, float(stride_seconds)


def get_ctc_logits(
    asr_model: Any, waveform: Tensor, sample_rate: int
) -> Tuple[Tensor, float]:
    device = next(asr_model.parameters()).device
    frontend, f_dtype = _get_frontend_and_dtype(asr_model)
    with torch.no_grad():
        wav = waveform.to(device=device, dtype=f_dtype)
        x = wav.unsqueeze(0)
        seqs_layout = _make_seqs_layout(x, [wav.shape[0]])
        ret = frontend.extract_features(x, seqs_layout)
        feats, layout = _ensure_feats_and_layout(ret, seqs_layout)
        try:
            feats = _ensure_tensor(frontend.process_features(feats, layout, None))
        except Exception:
            pass
        if hasattr(asr_model, "encoder"):
            enc = _ensure_tensor(asr_model.encoder(feats, layout))
        elif hasattr(asr_model, "w2v2"):
            enc = _ensure_tensor(asr_model.w2v2(feats, layout))
        else:
            raise RuntimeError("ctc model missing encoder block")
        proj = None
        for name in ("final_proj", "ctc_head", "output_proj", "proj", "decoder"):
            if hasattr(asr_model, name):
                proj = getattr(asr_model, name)
                break
        if proj is None:
            raise RuntimeError("ctc model missing projection head.")
        logits = proj(enc).squeeze(0)
    stride_seconds = get_encoder_stride_seconds(asr_model, waveform, sample_rate)[1]
    return logits, stride_seconds


def _get_ctc_frame_boundaries(logits: Tensor) -> List[int]:
    # return frame indices where non-blank token changes occur, including 0 and t
    path = torch.argmax(logits, dim=-1).tolist()
    boundaries = [0]
    prev_token = -1
    for i, tok in enumerate(path):
        if tok != 0 and tok != prev_token:
            boundaries.append(i)
            prev_token = tok
        elif tok != 0:
            prev_token = tok
    boundaries.append(len(path))
    # deduplicate & monotonic
    out = [boundaries[0]]
    for b in boundaries[1:]:
        if b > out[-1]:
            out.append(b)
    return out


# unicode character ranges for scripts that typically do not use spaces.
_NON_SPACE_RANGES = [
    (0x4E00, 0x9FFF),  # cjk unified ideographs
    (0x3400, 0x4DBF),  # cjk unified ideographs extension a
    (0xAC00, 0xD7AF),  # hangul syllables
    (0x3040, 0x309F),  # hiragana
    (0x30A0, 0x30FF),  # katakana
    (0x0E00, 0x0E7F),  # thai
    (0x0E80, 0x0EFF),  # lao
    (0x1780, 0x17FF),  # khmer
    (0x1000, 0x109F),  # myanmar
]


def is_non_space_script(text: str) -> bool:
    """
    determines if a text is likely from a non-whitespace script.
    it checks if the majority (>50%) of non-whitespace characters
    fall within unicode blocks for cjk, thai, lao, etc.
    """
    stripped_text = "".join(text.split())
    if not stripped_text:
        return False

    non_space_char_count = 0
    for char in stripped_text:
        code_point = ord(char)
        for start, end in _NON_SPACE_RANGES:
            if start <= code_point <= end:
                non_space_char_count += 1
                break

    # if more than half the characters are from these scripts, treat as non-space.
    return (non_space_char_count / len(stripped_text)) > 0.5


def split_text_units(transcript: str, mode: str) -> List[Tuple[str, int, int]]:
    units: List[Tuple[str, int, int]] = []
    if mode == "word":
        for m in re.finditer(r"\s+", transcript):
            units.append((m.group(0), m.start(), m.end()))
    elif mode == "char":
        for idx, ch in enumerate(transcript):
            if not ch.isspace():
                units.append((ch, idx, idx + 1))
    else:
        raise ValueError(f"unknown mode: {mode}")
    return units


def align_units_to_ctc_frames(
    units: List[Tuple[str, int, int]],
    frame_boundaries: List[int],
    stride_seconds: float,
) -> List[WordTimestamp]:
    if not frame_boundaries or len(frame_boundaries) < 2:
        return []
    if not units:
        return []
    num_frames = len(frame_boundaries) - 1
    U = len(units)
    out: List[WordTimestamp] = []
    for u_idx, (tok, _s, _e) in enumerate(units):
        start_bin = int(round((u_idx) * num_frames / U))
        end_bin = int(round((u_idx + 1) * num_frames / U))
        start_bin = min(max(start_bin, 0), num_frames)
        end_bin = min(max(end_bin, start_bin + 1), num_frames)
        start_frame = frame_boundaries[start_bin]
        end_frame = frame_boundaries[end_bin]
        out.append(
            WordTimestamp(
                word=tok,
                start=float(start_frame) * stride_seconds,
                end=float(end_frame) * stride_seconds,
            )
        )
    return out


def _make_energy_cdf_time_fn(
    waveform: Tensor, sample_rate: int
) -> Tuple[float, callable]:
    # return (duration_seconds, cdf_to_time) mapping fraction in [0,1]->seconds.
    wav = waveform.detach().float().abs()
    if wav.ndim != 1:
        wav = wav.view(-1)
    energy = (wav**2).to(torch.float32)
    csum = torch.cumsum(energy, dim=0)
    total = float(csum[-1]) if csum.numel() > 0 else 0.0
    duration = float(wav.shape[0]) / float(sample_rate)

    if total <= 0.0:

        def cdf_to_time(frac: float) -> float:
            f = min(max(float(frac), 0.0), 1.0)
            return f * duration

        return duration, cdf_to_time

    def cdf_to_time(frac: float) -> float:
        f = min(max(float(frac), 0.0), 1.0)
        target = f * total
        idx = int(torch.searchsorted(csum, tensor(target)))
        idx = min(max(idx, 0), int(wav.shape[0]))
        return float(idx) / float(sample_rate)

    return duration, cdf_to_time


def align_units_by_energy_cdf(
    units: List[Tuple[str, int, int]],
    waveform: Tensor,
    sample_rate: int,
) -> List[WordTimestamp]:
    if not units:
        return []
    duration, cdf_to_time = _make_energy_cdf_time_fn(waveform, sample_rate)
    U = len(units)
    out: List[WordTimestamp] = []
    last_t = 0.0
    for u_idx, (tok, _s, _e) in enumerate(units):
        start_f = u_idx / U
        end_f = (u_idx + 1) / U
        t0 = cdf_to_time(start_f)
        t1 = cdf_to_time(end_f)
        # ensure strict monotonicity in degenerate cases
        if t1 <= t0:
            t1 = min(duration, t0 + 1e-3)
        if t0 < last_t:
            t0 = last_t
            t1 = max(t1, t0 + 1e-3)
        out.append(WordTimestamp(word=tok, start=t0, end=t1))
        last_t = t1
    return out


def _detect_model_family(model: Any) -> str:
    name = type(model).__name__.lower()
    if hasattr(model, "ctc_head"):
        return "ctc"
    if hasattr(model, "decoder") or "llama" in name or "gpt" in name or "llm" in name:
        return "llm"
    if any(hasattr(model, n) for n in ("final_proj", "output_proj", "proj")):
        if not hasattr(model, "decoder"):
            return "ctc"
    return "llm"


def _align_chunk_any_model(
    pipeline: ASRInferencePipeline,
    chunk_wav: Tensor,
    sr: int,
    chunk_text: str,
) -> Tuple[List[WordTimestamp], List[WordTimestamp]]:
    model_type = _detect_model_family(pipeline.model)

    # both word and char units are always generated.
    # the final output stage will decide which one to use.
    word_units = split_text_units(chunk_text, mode="word")
    char_units = split_text_units(chunk_text, mode="char")

    if model_type == "ctc":
        logits, stride_seconds = get_ctc_logits(pipeline.model, chunk_wav, sr)
        frame_bounds = _get_ctc_frame_boundaries(logits)
        words_ts = align_units_to_ctc_frames(
            word_units, frame_bounds, stride_seconds
        )
        chars_ts = align_units_to_ctc_frames(
            char_units, frame_bounds, stride_seconds
        )
    else:  # llm path
        words_ts = align_units_by_energy_cdf(word_units, chunk_wav, sr)
        chars_ts = align_units_by_energy_cdf(char_units, chunk_wav, sr)

    return words_ts, chars_ts


def transcribe_with_timestamps(
    self: ASRInferencePipeline,
    inp: Union[str, Dict[str, Any], Tensor, List[Any]],
    batch_size: int = 1,
    max_chunk_sec: float = 40.0,
) -> List[Dict[str, Any]]:
    import torchaudio

    items: List[Any] = inp if isinstance(inp, list) else [inp]
    results: List[Dict[str, Any]] = []

    for idx, item in enumerate(items):
        # load/prepare waveform to mono 16k
        if isinstance(item, str):
            waveform, sr = torchaudio.load(item)
        elif isinstance(item, dict):
            waveform, sr = item["waveform"], int(item["sample_rate"])
        elif isinstance(item, Tensor):
            waveform, sr = item, 16000  # assume 16k for raw tensor
        else:
            raise ValueError(f"unsupported input type: {type(item)}")

        if not isinstance(waveform, Tensor):
            waveform = torch.tensor(waveform, dtype=torch.float32)

        if waveform.ndim > 1:
            waveform = waveform.mean(dim=0)

        target_sr = 16000
        if sr != target_sr:
            resampler = torchaudio.transforms.Resample(sr, target_sr)
            waveform = resampler(waveform)
            sr = target_sr

        chunks = chunk_waveform(waveform, sr, max_chunk_sec)
        beam_texts, all_words, all_chars = [], [], []

        for cidx, (cwav, off) in enumerate(chunks):
            text_list = self.transcribe(
                [{"waveform": cwav, "sample_rate": sr}],
                lang=None,  # ensure no lang is passed to native transcribe
                batch_size=batch_size,
            )
            chunk_text = text_list[0] if text_list else ""
            beam_texts.append(chunk_text)

            if chunk_text.strip():
                try:
                    words_ts, chars_ts = _align_chunk_any_model(
                        self, cwav, sr, chunk_text
                    )
                    for w in words_ts:
                        w.start += off
                        w.end += off
                    for c in chars_ts:
                        c.start += off
                        c.end += off
                    all_words.extend(words_ts)
                    all_chars.extend(chars_ts)
                except Exception as e:
                    print(f"warning: alignment failed for a chunk: {e}")

        full_text = " ".join(t for t in beam_texts if t.strip())
        results.append(
            {"text": full_text, "words": all_words, "chars": all_chars}
        )

    return results


# store a reference to the original init
_original_asr_pipeline_init = ASRInferencePipeline.__init__


def _patched_asr_pipeline_init(self, *args, **kwargs):
    # handle `model_card` being passed as either a positional or keyword arg
    effective_model_card, other_args = None, list(args)
    if args:
        effective_model_card, other_args = args[0], list(args[1:])
    elif 'model_card' in kwargs:
        effective_model_card = kwargs.pop('model_card')

    # for omnilingual-asr 0.1.0, 'model_card' must be positional if used.
    if effective_model_card is not None:
        _original_asr_pipeline_init(self, effective_model_card, *other_args, **kwargs)
    else:
        _original_asr_pipeline_init(self, *args, **kwargs)

    # attach our timestamping method to every new instance
    self.transcribe_with_timestamps = types.MethodType(
        transcribe_with_timestamps, self
    )

# apply the patch to the class
ASRInferencePipeline.__init__ = _patched_asr_pipeline_init

import warnings
import json
import torch
import torchaudio

# -- user configuration --
FILENAME = "ado.wav"
MODEL_TYPE = "LLM"  # options: "CTC" or "LLM"
PARAM_COUNT = "300M" # e.g., "300M", "600M", "1B"
# -- end configuration --

# this import applies the patch to the asr pipeline
import omni_ts_patch
from omnilingual_asr.models.inference.pipeline import ASRInferencePipeline


# 1. construct model card from configuration
model_card = f"omniASR_{MODEL_TYPE}_{PARAM_COUNT}"
print(f"--- Configuration ---")
print(f"File: {FILENAME}")
print(f"Model Card: {model_card}\n")


# 2. load and prepare audio
print("--- Loading Audio ---")
try:
    waveform, sr = torchaudio.load(FILENAME)
    # resample to 16khz mono
    if sr != 16000:
        resampler = torchaudio.transforms.Resample(sr, 16000)
        waveform = resampler(waveform)
    if waveform.shape[0] > 1:
        waveform = torch.mean(waveform, dim=0, keepdim=True)
    print(f"Loaded and preprocessed '{FILENAME}'. Duration: {waveform.shape[1]/16000:.2f}s\n")
except Exception as e:
    raise RuntimeError(f"failed to load audio file '{FILENAME}'. error: {e}")


# 3. load model and transcribe
print("--- Transcription ---")
try:
    print(f"Loading pipeline for {model_card}...")
    pipeline = ASRInferencePipeline(model_card)
    print("Pipeline loaded. Starting transcription...")
    results = pipeline.transcribe_with_timestamps(
        [{"waveform": waveform.squeeze(0), "sample_rate": 16000}],
        max_chunk_sec=30.0,
    )
    result = results[0]
    print("Transcription complete.\n")
except Exception as e:
    raise RuntimeError(f"failed during model loading or transcription. error: {e}")


# 4. process and save results
print("--- Results ---")
full_text = result["text"]
print("Full Text:", full_text)

# decide whether to use word or character segments based on the output text
use_char_segments = omni_ts_patch.is_non_space_script(full_text)
source_segments = result["chars"] if use_char_segments else result["words"]

print(f"\nSegmentation strategy: {'CHARACTER' if use_char_segments else 'WORD'}-level")

# create simplified, clean output
output_segments = [
    {"text": seg.word, "start": round(seg.start, 2), "end": round(seg.end, 2)}
    for seg in source_segments
]

final_output = {
    "model": model_card,
    "text": full_text,
    "segments": output_segments,
}

# print first 10 segments for preview
print("\nFirst 10 Segments:")
for seg in final_output["segments"][:10]:
    print(f"  [{seg['start']:.2f}s - {seg['end']:.2f}s] {seg['text']}")

# save the final json object
output_filename = "transcription_output.json"
with open(output_filename, "w", encoding="utf-8") as f:
    json.dump(final_output, f, indent=2, ensure_ascii=False)

print(f"\nSuccessfully saved detailed transcription to '{output_filename}'")

Hope this helps!

[Feature] Word level timestamps #3

Description

Feature Description

Use Case

Metadata

Metadata

Assignees

Labels

Type

Projects

Milestone

Relationships

Development

Issue actions