easytranscriber 0.1.0__tar.gz

easytranscriber-0.1.0/PKG-INFO

@@ -0,0 +1,34 @@
+Metadata-Version: 2.4
+Name: easytranscriber
+Version: 0.1.0
+Summary: Speech recognition with accurate word-level timestamps.
+Author: Faton Rekathati
+Project-URL: Repository, https://github.com/kb-labb/easytranscriber
+Requires-Python: >=3.10
+Description-Content-Type: text/markdown
+Requires-Dist: transformers>=4.45.0
+Requires-Dist: torch!=2.9.*,>=2.7.0
+Requires-Dist: torchaudio!=2.9.*,>=2.7.0
+Requires-Dist: tqdm>=4.66.1
+Requires-Dist: soundfile>=0.12.1
+Requires-Dist: nltk>=3.8.2
+Requires-Dist: pyannote-audio>=3.3.1
+Requires-Dist: silero-vad~=6.0
+Requires-Dist: ctranslate2>=4.4.0
+Requires-Dist: msgspec
+Requires-Dist: easyaligner==0.*
+
+<div align="center"><img width="1020" height="340" alt="image" src="https://github.com/user-attachments/assets/7f1bdf33-5161-40c1-b6a7-6f1f586e030b" /></div>
+
+
+`easytranscriber` is an automatic speech recognition library built for efficient, large-scale transcription with accurate word-level timestamps. The library is backend-agnostic, featuring modular, parallelizable, pipeline components (VAD, transcription, feature/emission extraction, forced alignment), with support for both `ctranslate2` and `Hugging Face` inference backends. Notable features include:
+
+* **GPU accelerated forced alignment**, using [Pytorch's forced alignment API](https://docs.pytorch.org/audio/main/tutorials/ctc_forced_alignment_api_tutorial.html). Forced alignment is based on a GPU implementation of the Viterbi algorithm ([Pratap et al., 2024](https://jmlr.org/papers/volume25/23-1318/23-1318.pdf#page=8)).
+* **Parallel loading and pre-fetching of audio files** for efficient data loading and batch processing.
+* **Flexible text normalization for improved alignment quality**. Users can supply custom regex-based text normalization functions to preprocess ASR outputs before alignment. A mapping from the original text to the normalized text is maintained internally. All of the applied normalizations and transformations are consequently **non-destructive and reversible after alignment**. 
+* **35% to 102% faster inference compared to [`WhisperX`](https://github.com/m-bain/whisperX)**. See the [benchmarks](#benchmarks) for more details.
+* Batch inference support for both wav2vec2 and Whisper models.
+
+### Benchmarks
+
+![Benchmarks](benchmarks/plots/all_speedup.png)

easytranscriber-0.1.0/README.md

@@ -0,0 +1,14 @@
+<div align="center"><img width="1020" height="340" alt="image" src="https://github.com/user-attachments/assets/7f1bdf33-5161-40c1-b6a7-6f1f586e030b" /></div>
+
+
+`easytranscriber` is an automatic speech recognition library built for efficient, large-scale transcription with accurate word-level timestamps. The library is backend-agnostic, featuring modular, parallelizable, pipeline components (VAD, transcription, feature/emission extraction, forced alignment), with support for both `ctranslate2` and `Hugging Face` inference backends. Notable features include:
+
+* **GPU accelerated forced alignment**, using [Pytorch's forced alignment API](https://docs.pytorch.org/audio/main/tutorials/ctc_forced_alignment_api_tutorial.html). Forced alignment is based on a GPU implementation of the Viterbi algorithm ([Pratap et al., 2024](https://jmlr.org/papers/volume25/23-1318/23-1318.pdf#page=8)).
+* **Parallel loading and pre-fetching of audio files** for efficient data loading and batch processing.
+* **Flexible text normalization for improved alignment quality**. Users can supply custom regex-based text normalization functions to preprocess ASR outputs before alignment. A mapping from the original text to the normalized text is maintained internally. All of the applied normalizations and transformations are consequently **non-destructive and reversible after alignment**. 
+* **35% to 102% faster inference compared to [`WhisperX`](https://github.com/m-bain/whisperX)**. See the [benchmarks](#benchmarks) for more details.
+* Batch inference support for both wav2vec2 and Whisper models.
+
+### Benchmarks
+
+![Benchmarks](benchmarks/plots/all_speedup.png)

easytranscriber-0.1.0/pyproject.toml

@@ -0,0 +1,57 @@
+[build-system]
+requires = ["setuptools>=67.0.0"]
+build-backend = "setuptools.build_meta"
+
+[project]
+version = "0.1.0"
+name = "easytranscriber"
+requires-python = ">= 3.10"
+description = "Speech recognition with accurate word-level timestamps."
+readme = "README.md"
+authors = [{ name = "Faton Rekathati" }]
+
+dependencies = [
+  "transformers>=4.45.0",
+  "torch>=2.7.0,!=2.9.*",
+  "torchaudio>=2.7.0,!=2.9.*",
+  "tqdm>=4.66.1",
+  "soundfile>=0.12.1",
+  "nltk>=3.8.2",
+  "pyannote-audio>=3.3.1",
+  "silero-vad~=6.0",
+  "ctranslate2>=4.4.0",
+  "msgspec",
+  "easyaligner==0.*"
+]
+
+[project.urls]
+Repository = "https://github.com/kb-labb/easytranscriber"
+
+[tool.uv.sources]
+torch = [
+  { index = "pytorch-cpu", marker = "sys_platform == 'darwin'" },
+  { index = "pytorch-cuda", marker = "sys_platform != 'darwin'" },
+]
+torchaudio = [
+  { index = "pytorch-cpu", marker = "sys_platform == 'darwin'" },
+  { index = "pytorch-cuda", marker = "sys_platform != 'darwin'" },
+]
+
+[[tool.uv.index]]
+name = "pytorch-cuda"
+url = "https://download.pytorch.org/whl/cu128"
+explicit = true
+
+[[tool.uv.index]]
+name = "pytorch-cpu"
+url = "https://download.pytorch.org/whl/cpu"
+explicit = true
+
+[tool.ruff]
+line-length = 99
+indent-width = 4
+target-version = "py310"
+
+[tool.ruff.lint]
+# Disable fix for unused imports (`F401`).
+unfixable = ["F401"]

easytranscriber-0.1.0/setup.cfg

@@ -0,0 +1,4 @@
+[egg_info]
+tag_build = 
+tag_date = 0
+

easytranscriber-0.1.0/src/easytranscriber/asr/ct2.py

@@ -0,0 +1,245 @@
+"""
+CTranslate2-based Whisper transcription module.
+
+This module provides a transcribe function using ctranslate2 for efficient
+Whisper inference, mirroring the HuggingFace implementation in hf.py.
+"""
+
+import logging
+from pathlib import Path
+
+import ctranslate2
+import numpy as np
+import torch
+from easyaligner.utils import save_metadata_json, save_metadata_msgpack
+from easytranscriber.data.collators import transcribe_collate_fn
+from tqdm import tqdm
+from transformers import WhisperProcessor
+
+logger = logging.getLogger(__name__)
+
+
+def transcribe(
+    model: ctranslate2.models.Whisper,
+    processor: WhisperProcessor,
+    file_dataloader: torch.utils.data.DataLoader,
+    language: str | None = None,
+    task: str = "transcribe",
+    batch_size: int = 8,
+    beam_size: int = 5,
+    patience: float = 1.0,
+    length_penalty: float = 1.0,
+    repetition_penalty: float = 1.0,
+    no_repeat_ngram_size: int = 0,
+    max_length: int = 448,
+    suppress_blank: bool = True,
+    num_workers: int = 2,
+    prefetch_factor: int = 2,
+    output_dir: str = "output/transcriptions",
+):
+    """
+    Transcribe audio files using CTranslate2 Whisper model.
+
+    This function processes audio files through a dataloader structure similar
+    to the HuggingFace implementation, but uses ctranslate2 for inference.
+
+    Parameters
+    ----------
+    model : ctranslate2.models.Whisper
+        CTranslate2 Whisper model.
+    processor : transformers.WhisperProcessor
+        WhisperProcessor for tokenization and decoding.
+    file_dataloader : torch.utils.data.DataLoader
+        DataLoader yielding audio file datasets.
+    language : str, optional
+        Language code (e.g., 'sv', 'en'). If None, auto-detect.
+    batch_size : int, optional
+        Batch size for feature processing.
+    task : str, optional
+        Task type - 'transcribe' or 'translate'.
+    beam_size : int, optional
+        Beam size for search. Default is 5.
+    patience : float, optional
+        Beam search patience factor. Default is 1.0.
+    length_penalty : float, optional
+        Length penalty for beam search. Default is 1.0.
+    repetition_penalty : float, optional
+        Repetition penalty. Default is 1.0.
+    no_repeat_ngram_size : int, optional
+        N-gram size for no repeat. Default is 0.
+    max_length : int, optional
+        Maximum output length. Default is 448.
+    suppress_blank : bool, optional
+        Whether to suppress blank tokens. Default is True.
+    num_workers : int, optional
+        Number of workers for feature dataloader (file dataloader is created outside
+        of this function).
+    prefetch_factor : int, optional
+        Prefetch factor for feature dataloader (file dataloader is created outside
+        of this function).
+    output_dir : str, optional
+        Directory to save transcription JSON files. Default is `output/transcriptions`.
+    """
+    for features in tqdm(file_dataloader, desc="Transcribing audio files"):
+        slice_dataset = features[0]["dataset"]
+        metadata = features[0]["dataset"].metadata
+        transcription_texts = []
+        language_detections = []
+
+        feature_dataloader = torch.utils.data.DataLoader(
+            slice_dataset,
+            batch_size=batch_size,
+            num_workers=num_workers,
+            prefetch_factor=prefetch_factor,
+            pin_memory=True,
+            collate_fn=transcribe_collate_fn,
+        )
+
+        logger.info(f"Transcribing {metadata.audio_path} ...")
+
+        for batch in feature_dataloader:
+            batch_features = batch["features"].numpy()  # Convert to numpy for ctranslate2
+            current_batch_size = batch_features.shape[0]
+
+            # Convert to ctranslate2 StorageView
+            features_ct2 = ctranslate2.StorageView.from_array(batch_features)
+
+            if language is not None:
+                # Build the prompt tokens for the batch
+                prompt_tokens = [
+                    "<|startoftranscript|>",
+                    f"<|{language}|>",
+                    f"<|{task}|>",
+                    "<|notimestamps|>",
+                ]
+                prompt_ids = processor.tokenizer.convert_tokens_to_ids(prompt_tokens)
+                prompt_ids = [prompt_ids] * current_batch_size
+            else:
+                languages = detect_language(model, features_ct2)
+                language_detections.extend(languages)
+                prompt_ids = []
+                for curr_lang in languages:
+                    prompt_tokens = [
+                        "<|startoftranscript|>",
+                        curr_lang["language"],
+                        f"<|{task}|>",
+                        "<|notimestamps|>",
+                    ]
+
+                    prompt_ids.append(processor.tokenizer.convert_tokens_to_ids(prompt_tokens))
+
+            # Generate transcriptions
+            outputs = model.generate(
+                features_ct2,
+                prompt_ids,
+                beam_size=beam_size,
+                patience=patience,
+                length_penalty=length_penalty,
+                repetition_penalty=repetition_penalty,
+                no_repeat_ngram_size=no_repeat_ngram_size,
+                max_length=max_length,
+                suppress_blank=suppress_blank,
+            )
+
+            # Decode the sequences
+            sequences = [result.sequences_ids[0] for result in outputs]
+            transcription = processor.batch_decode(sequences, skip_special_tokens=True)
+
+            transcription_texts.extend(transcription)
+
+        # Update metadata with transcriptions
+        for i, speech in enumerate(metadata.speeches):
+            for j, chunk in enumerate(speech.chunks):
+                chunk.text = transcription_texts[j].strip()
+                if len(language_detections) > 0:
+                    chunk.language = language_detections[j]["language"]
+                    chunk.language_prob = language_detections[j]["probability"]
+
+        # Save transcription to file
+        output_path = Path(output_dir) / Path(metadata.audio_path).with_suffix(".json")
+        output_path.parent.mkdir(parents=True, exist_ok=True)
+        save_metadata_json(metadata, output_dir=output_dir)
+
+
+def lang_detect_only(
+    model,
+    file_dataloader,
+    batch_size=8,
+    num_workers=2,
+    prefetch_factor=2,
+    output_dir=None,
+):
+    """
+    Run language detection only.
+
+    Parameters
+    ----------
+    model : ctranslate2.models.Whisper
+        CTranslate2 Whisper model.
+    file_dataloader : torch.utils.data.DataLoader
+        DataLoader yielding audio file datasets.
+    batch_size : int, optional
+        Batch size. Default is 8.
+    num_workers : int, optional
+        Number of workers. Default is 2.
+    prefetch_factor : int, optional
+        Prefetch factor. Default is 2.
+    output_dir : str, optional
+        Output directory. Default is None.
+    """
+    for features in file_dataloader:
+        slice_dataset = features[0]["dataset"]
+        metadata = features[0]["dataset"].metadata
+        language_detections = []
+
+        feature_dataloader = torch.utils.data.DataLoader(
+            slice_dataset,
+            batch_size=batch_size,
+            num_workers=num_workers,
+            prefetch_factor=prefetch_factor,
+            pin_memory=True,
+            collate_fn=transcribe_collate_fn,
+        )
+
+        for batch in feature_dataloader:
+            features_ct2 = batch["features"].numpy()
+            features_ct2 = ctranslate2.StorageView.from_array(features_ct2)
+            languages = detect_language(model, features_ct2)
+            language_detections.append(languages)
+
+        for i, speech in enumerate(metadata.speeches):
+            for j, chunk in enumerate(speech.chunks):
+                chunk.language = language_detections[j]["language"]
+                chunk.language_probability = language_detections[j]["probability"]
+
+        # Save transcription to file
+        output_path = Path(output_dir) / Path(metadata.audio_path).with_suffix(".json")
+        output_path.parent.mkdir(parents=True, exist_ok=True)
+        save_metadata_json(metadata, output_dir=output_dir)
+
+
+def detect_language(model: ctranslate2.models.Whisper, features: ctranslate2.StorageView) -> list:
+    """
+    Return the highest probability language for each chunk in the features batch.
+
+    Parameters
+    ----------
+    model : ctranslate2.models.Whisper
+        CTranslate2 Whisper model.
+    features : ctranslate2.StorageView
+        Input features.
+
+    Returns
+    -------
+    list
+        List of dicts containing 'language' and 'probability'.
+    """
+
+    # List of tuple[str, float] with language and probability
+    lang_probs = model.detect_language(features)
+
+    top_langs = []
+    for chunk in lang_probs:
+        top_langs.append({"language": chunk[0][0], "probability": chunk[0][1]})
+
+    return top_langs

easytranscriber-0.1.0/src/easytranscriber/asr/hf.py

@@ -0,0 +1,106 @@
+import logging
+from pathlib import Path
+
+import torch
+from easyaligner.data.dataset import AudioFileDataset
+from easyaligner.utils import save_metadata_json
+from easytranscriber.data.collators import transcribe_collate_fn
+from tqdm import tqdm
+from transformers import WhisperForConditionalGeneration, WhisperProcessor
+
+logger = logging.getLogger(__name__)
+
+
+def transcribe(
+    model: WhisperForConditionalGeneration,
+    processor: WhisperProcessor,
+    file_dataloader: torch.utils.data.DataLoader,
+    language: str | None = None,
+    task: str = "transcribe",
+    batch_size: int = 4,
+    beam_size: int = 5,
+    length_penalty: float = 1.0,
+    repetition_penalty: float = 1.0,
+    max_length: int = 250,
+    num_workers: int = 2,
+    prefetch_factor: int = 2,
+    output_dir: str = "output/transcriptions",
+    device: str = "cuda",
+):
+    """
+    Transcribe audio files using HuggingFace Whisper model.
+
+    Parameters
+    ----------
+    model : transformers.WhisperForConditionalGeneration
+        HuggingFace Whisper model.
+    processor : transformers.WhisperProcessor
+        HuggingFace Whisper processor.
+    file_dataloader : torch.utils.data.DataLoader
+        DataLoader yielding audio file datasets.
+    language : str, optional
+        Language code (e.g., 'sv', 'en'). Default is `None` (auto-detect).
+    batch_size : int, optional
+        Batch size for inference.
+    beam_size : int, optional
+        Number of beams for beam search. Default is 5.
+    length_penalty : float, optional
+        Length penalty. Default is 1.0.
+    repetition_penalty : float, optional
+        Repetition penalty. Default is 1.0.
+    max_length : int, optional
+        Maximum length of generated text. Default is 250.
+    num_workers : int, optional
+        Number of workers for feature dataloader.
+    prefetch_factor : int, optional
+        Prefetch factor for feature dataloader.
+    output_dir : str, optional
+        Directory to save transcription JSON files. Default is `output/transcriptions`.
+    device : str, optional
+        Device to run inference on. Default is `cuda`.
+    """
+    for features in tqdm(file_dataloader, desc="Transcribing audio files"):
+        slice_dataset = features[0]["dataset"]
+        metadata = features[0]["dataset"].metadata
+        transcription_texts = []
+
+        feature_dataloader = torch.utils.data.DataLoader(
+            slice_dataset,
+            batch_size=batch_size,
+            num_workers=num_workers,
+            prefetch_factor=prefetch_factor,
+            collate_fn=transcribe_collate_fn,
+        )
+
+        logger.info(f"Transcribing {metadata.audio_path} ...")
+
+        for batch in feature_dataloader:
+            with torch.inference_mode():
+                batch = batch["features"].to(device).half()
+                predicted_ids = model.generate(
+                    batch,
+                    return_dict_in_generate=True,
+                    task=task,
+                    language=language,
+                    output_scores=False,
+                    max_length=max_length,
+                    num_beams=beam_size,
+                    repetition_penalty=repetition_penalty,
+                    length_penalty=length_penalty,
+                    early_stopping=True,
+                )
+
+                transcription = processor.batch_decode(
+                    predicted_ids["sequences"], skip_special_tokens=True
+                )
+
+                transcription_texts.extend(transcription)
+
+        for i, speech in enumerate(metadata.speeches):
+            for j, chunk in enumerate(speech.chunks):
+                chunk.text = transcription_texts[j].strip()
+
+        # Write final transcription to file with msgspec serialization
+        output_path = Path(output_dir) / Path(metadata.audio_path).with_suffix(".json")
+        output_path.parent.mkdir(parents=True, exist_ok=True)
+        save_metadata_json(metadata, output_dir=output_dir)

easytranscriber-0.1.0/src/easytranscriber/audio.py

@@ -0,0 +1,148 @@
+import logging
+import subprocess
+from pathlib import Path
+from typing import Tuple
+
+import numpy as np
+
+logger = logging.getLogger(__name__)
+
+
+def read_audio_segment(
+    audio_path: str | Path,
+    start_sec: float,
+    duration_sec: float,
+    sample_rate: int = 16000,
+) -> np.ndarray:
+    """
+    Read a segment of audio using ffmpeg subprocess with seek.
+
+    Uses ffmpeg's fast seek (-ss before -i) to efficiently read only the
+    required segment, with resampling to the target sample rate and mono conversion.
+
+    Parameters
+    ----------
+    audio_path : str or Path
+        Path to the audio file.
+    start_sec : float
+        Start time in seconds.
+    duration_sec : float
+        Duration to read in seconds.
+    sample_rate : int, optional
+        Target sample rate for resampling.
+
+    Returns
+    -------
+    np.ndarray
+        Audio data as float32 numpy array.
+    """
+    cmd = [
+        "ffmpeg",
+        "-ss",
+        str(start_sec),  # Seek to position (before -i = fast seek)
+        "-i",
+        str(audio_path),
+        "-t",
+        str(duration_sec),  # Read this many seconds
+        "-ar",
+        str(sample_rate),  # Resample
+        "-ac",
+        "1",  # Mono
+        "-f",
+        "f32le",  # Raw float32 little-endian output
+        "-loglevel",
+        "error",
+        "pipe:1",  # Output to stdout
+    ]
+
+    try:
+        proc = subprocess.run(cmd, capture_output=True, check=True)
+        audio = np.frombuffer(proc.stdout, dtype=np.float32)
+        return audio
+    except subprocess.CalledProcessError as e:
+        logger.error(f"ffmpeg error reading {audio_path}: {e.stderr.decode()}")
+        raise
+
+
+def convert_audio_to_array(input_file: str, sample_rate: int = 16000) -> Tuple[np.ndarray, int]:
+    """
+    Convert audio to in-memory numpy array.
+
+    Parameters
+    ----------
+    input_file : str
+        Path to the input audio file.
+    sample_rate : int, optional
+        Target sample rate.
+
+    Returns
+    -------
+    Tuple[np.ndarray, int]
+        Tuple containing the audio array (int16) and the sample rate.
+
+    Raises
+    ------
+    RuntimeError
+        If ffmpeg command fails.
+    """
+    # fmt: off
+    command = [
+        "ffmpeg",
+        "-i", input_file,
+        "-f", "s16le",  # raw PCM 16-bit little endian
+        "-acodec", "pcm_s16le",
+        "-ac", "1",  # mono
+        "-ar", str(sample_rate),  # 16 kHz
+        "-loglevel", "error",  # suppress output
+        "-hide_banner",
+        "-nostats",
+    ]
+    # fmt: on
+
+    process = subprocess.Popen(command, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
+    out, err = process.communicate()
+
+    if process.returncode != 0:
+        raise RuntimeError(f"ffmpeg error: {err.decode()}")
+
+    # Convert byte output to numpy array
+    audio_array = np.frombuffer(out, dtype=np.int16)
+
+    return audio_array, sample_rate  # (samples, sample_rate)
+
+
+def convert_audio_to_wav(input_file: str, output_file: str) -> None:
+    """
+    Convert audio file to WAV format with 16kHz sample rate and mono channel.
+
+    Parameters
+    ----------
+    input_file : str
+        Path to the input audio file.
+    output_file : str
+        Path to the output WAV file.
+
+    Raises
+    ------
+    RuntimeError
+        If ffmpeg command fails.
+    """
+    # fmt: off
+    command = [
+        'ffmpeg',
+        '-i', input_file,
+        '-ar', '16000',  # Set the audio sample rate to 16kHz
+        '-ac', '1',      # Set the number of audio channels to 1 (mono)
+        '-c:a', 'pcm_s16le',
+        '-loglevel', 'warning',
+        '-hide_banner',
+        '-nostats',
+        '-nostdin',
+        output_file
+    ]
+    # fmt: on
+    process = subprocess.Popen(command, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
+    _, err = process.communicate()
+
+    if process.returncode != 0:
+        raise RuntimeError(f"ffmpeg error: {err.decode()}")

easytranscriber-0.1.0/src/easytranscriber/data/__init__.py

@@ -0,0 +1,3 @@
+from easytranscriber.data.dataset import StreamingAudioFileDataset
+
+__all__ = ["StreamingAudioFileDataset"]

easytranscriber-0.1.0/src/easytranscriber/data/collators.py

@@ -0,0 +1,30 @@
+import torch
+
+
+def transcribe_collate_fn(batch: list[dict]) -> dict:
+    """
+    Collate function for transcription.
+
+    Parameters
+    ----------
+    batch : list of dict
+        List of samples from the dataset.
+
+    Returns
+    -------
+    dict
+        Collated batch with 'features', 'start_times', and 'speech_ids'.
+    """
+    # Remove None values
+    speech_ids = [b["speech_id"] for b in batch if b is not None]
+    start_times = [b["start_time_global"] for b in batch if b is not None]
+    batch = [b["feature"] for b in batch if b is not None]
+
+    # Concat, keep batch dimension
+    batch = torch.cat(batch, dim=0)
+
+    return {
+        "features": batch,
+        "start_times": start_times,
+        "speech_ids": speech_ids,
+    }