[ { "path": ".github/FUNDING.yml", "content": "custom: https://www.buymeacoffee.com/maxhbain\n" }, { "path": ".github/workflows/build-and-release.yml", "content": "name: Build and release\n\non:\n release:\n types: [published]\n\njobs:\n build:\n runs-on: ubuntu-latest\n steps:\n - name: Checkout\n uses: actions/checkout@v4\n\n - name: Install uv\n uses: astral-sh/setup-uv@v5\n with:\n version: \"0.5.14\"\n python-version: \"3.10\"\n\n - name: Check if lockfile is up to date\n run: uv lock --check\n\n - name: Build package\n run: uv build\n\n - name: Release to Github\n uses: softprops/action-gh-release@v2\n with:\n files: dist/*.whl\n\n - name: Publish package to PyPi\n run: uv publish\n env:\n UV_PUBLISH_TOKEN: ${{ secrets.PYPI_API_TOKEN }}\n" }, { "path": ".github/workflows/python-compatibility.yml", "content": "name: Python Compatibility Test\n\non:\n push:\n branches: [main]\n pull_request:\n branches: [main]\n workflow_dispatch: # Allows manual triggering from GitHub UI\n\njobs:\n test:\n runs-on: ubuntu-latest\n strategy:\n matrix:\n python-version: [\"3.10\", \"3.11\", \"3.12\", \"3.13\"]\n\n steps:\n - uses: actions/checkout@v4\n\n - name: Install uv\n uses: astral-sh/setup-uv@v5\n with:\n version: \"0.5.14\"\n python-version: ${{ matrix.python-version }}\n\n - name: Check if lockfile is up to date\n run: uv lock --check\n\n - name: Install the project\n run: uv sync --all-extras\n\n - name: Test import\n run: |\n uv run python -c \"import whisperx; print('Successfully imported whisperx')\"\n" }, { "path": ".gitignore", "content": "# Byte-compiled / optimized / DLL files\n__pycache__/\n*.py[cod]\n*$py.class\n\n# C extensions\n*.so\n\n# Distribution / packaging\n.Python\nbuild/\ndevelop-eggs/\ndist/\ndownloads/\neggs/\n.eggs/\nlib/\nlib64/\nparts/\nsdist/\nvar/\nwheels/\nshare/python-wheels/\n*.egg-info/\n.installed.cfg\n*.egg\nMANIFEST\n\n# PyInstaller\n# Usually these files are written by a python script from a template\n# before PyInstaller builds the exe, so as to inject date/other infos into it.\n*.manifest\n*.spec\n\n# Installer logs\npip-log.txt\npip-delete-this-directory.txt\n\n# Unit test / coverage reports\nhtmlcov/\n.tox/\n.nox/\n.coverage\n.coverage.*\n.cache\nnosetests.xml\ncoverage.xml\n*.cover\n*.py,cover\n.hypothesis/\n.pytest_cache/\ncover/\n\n# Translations\n*.mo\n*.pot\n\n# Django stuff:\n*.log\nlocal_settings.py\ndb.sqlite3\ndb.sqlite3-journal\n\n# Flask stuff:\ninstance/\n.webassets-cache\n\n# Scrapy stuff:\n.scrapy\n\n# Sphinx documentation\ndocs/_build/\n\n# PyBuilder\n.pybuilder/\ntarget/\n\n# Jupyter Notebook\n.ipynb_checkpoints\n\n# IPython\nprofile_default/\nipython_config.py\n\n# pyenv\n# For a library or package, you might want to ignore these files since the code is\n# intended to run in multiple environments; otherwise, check them in:\n# .python-version\n\n# pipenv\n# According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control.\n# However, in case of collaboration, if having platform-specific dependencies or dependencies\n# having no cross-platform support, pipenv may install dependencies that don't work, or not\n# install all needed dependencies.\n#Pipfile.lock\n\n# UV\n# Similar to Pipfile.lock, it is generally recommended to include uv.lock in version control.\n# This is especially recommended for binary packages to ensure reproducibility, and is more\n# commonly ignored for libraries.\n#uv.lock\n\n# poetry\n# Similar to Pipfile.lock, it is generally recommended to include poetry.lock in version control.\n# This is especially recommended for binary packages to ensure reproducibility, and is more\n# commonly ignored for libraries.\n# https://python-poetry.org/docs/basic-usage/#commit-your-poetrylock-file-to-version-control\n#poetry.lock\n\n# pdm\n# Similar to Pipfile.lock, it is generally recommended to include pdm.lock in version control.\n#pdm.lock\n# pdm stores project-wide configurations in .pdm.toml, but it is recommended to not include it\n# in version control.\n# https://pdm.fming.dev/latest/usage/project/#working-with-version-control\n.pdm.toml\n.pdm-python\n.pdm-build/\n\n# PEP 582; used by e.g. github.com/David-OConnor/pyflow and github.com/pdm-project/pdm\n__pypackages__/\n\n# Celery stuff\ncelerybeat-schedule\ncelerybeat.pid\n\n# SageMath parsed files\n*.sage.py\n\n# Environments\n.env\n.venv\nenv/\nvenv/\nENV/\nenv.bak/\nvenv.bak/\n\n# Spyder project settings\n.spyderproject\n.spyproject\n\n# Rope project settings\n.ropeproject\n\n# mkdocs documentation\n/site\n\n# mypy\n.mypy_cache/\n.dmypy.json\ndmypy.json\n\n# Pyre type checker\n.pyre/\n\n# pytype static type analyzer\n.pytype/\n\n# Cython debug symbols\ncython_debug/\n\n# PyCharm\n# JetBrains specific template is maintained in a separate JetBrains.gitignore that can\n# be found at https://github.com/github/gitignore/blob/main/Global/JetBrains.gitignore\n# and can be added to the global gitignore or merged into this file. For a more nuclear\n# option (not recommended) you can uncomment the following to ignore the entire idea folder.\n#.idea/\n\n# PyPI configuration file\n.pypirc" }, { "path": ".python-version", "content": "3.10\n" }, { "path": "CUDNN_TROUBLESHOOTING.md", "content": "# Troubleshooting cuDNN Loading Errors\n\nThis guide helps resolve common cuDNN-related errors when running WhisperX on GPU. These issues typically occur when the system can't locate cuDNN libraries or finds conflicting versions.\n\n## Unable to Load cuDNN Libraries\n\nIf you encounter the following error when running WhisperX:\n\n`Unable to load any of {libcudnn_cnn.so.9.1.0, libcudnn_cnn.so.9.1, libcudnn_cnn.so.9, libcudnn_cnn.so}`\n\nThis means the cuDNN libraries are installed (via whisperx dependencies) but aren't in a location where the system's dynamic linker can find them.\n\n### Solution 1: Add to LD_LIBRARY_PATH (Recommended)\n\nAdd this at the start of your Python script or notebook:\n\n```python\nimport os\n\n# Get current LD_LIBRARY_PATH\noriginal = os.environ.get(\"LD_LIBRARY_PATH\", \"\")\n\ncudnn_path = \"/usr/local/lib/python3.12/dist-packages/nvidia/cudnn/lib/\"\nos.environ['LD_LIBRARY_PATH'] = original + \":\" + cudnn_path\n```\n\n**Note:** Adjust the Python version (`python3.12`) to match your environment.\n\n### Solution 2: Symlink to LD_LIBRARY_PATH Directory\n\nIf Solution 1 didn't work and you still get the \"unable to load\" error, symlink the libraries to a directory that's already in your `LD_LIBRARY_PATH`:\n\n1. Check what's in your LD_LIBRARY_PATH: `echo \"$LD_LIBRARY_PATH\"`\n2. Assuming that there is only one path set. \n Symlink the downloaded libcudnn files to that path: \n `ln -s /usr/local/lib/python3.12/dist-packages/nvidia/cudnn/lib/libcudnn* \"$LD_LIBRARY_PATH\"/`\n\n **Note:** If `LD_LIBRARY_PATH` contains multiple paths (separated by `:`), pick one directory and use it instead of `\"$LD_LIBRARY_PATH\"`. For example: `/usr/lib/x86_64-linux-gnu/`\n\n## cuDNN Version Incompatibility\n\nIf you encounter this error:\n\n```\nRuntimeError: cuDNN version incompatibility: PyTorch was compiled against (9, 10, 2) but found runtime version (9, 2, 1)\n```\n\nThis means PyTorch is finding a different cuDNN version than the one it was compiled with. **PyTorch comes bundled with its own cuDNN**, but a conflicting cuDNN in `LD_LIBRARY_PATH` is taking precedence.\n\n### Solution: Remove Conflicting cuDNN from Path\n\nCheck if there's a conflicting cuDNN path:\n\n```bash\necho $LD_LIBRARY_PATH\n```\n\nIf you see paths pointing to older cuDNN installations (e.g., system-installed cuDNN or manually downloaded), try one of these:\n\n**Option 1: Clear LD_LIBRARY_PATH temporarily**\n\n```python\nimport os\n# Let PyTorch use its bundled cuDNN\nos.environ.pop('LD_LIBRARY_PATH', None)\n```\n\n**Option 2: Set LD_LIBRARY_PATH to only the correct version**\n\n```python\nimport os\n# Point only to the cuDNN that matches PyTorch's compiled version\nos.environ['LD_LIBRARY_PATH'] = \"/usr/local/lib/python3.12/dist-packages/nvidia/cudnn/lib/\"\n```\n\n**Note:** This error is unlikely on a clean install. If it occurs anyway, [open an issue](https://github.com/m-bain/whisperX/issues). If you've modified system libraries or CUDA/cuDNN, the options above should help resolve most cases.\n" }, { "path": "EXAMPLES.md", "content": "# More Examples\n\n## Other Languages\n\nFor non-english ASR, it is best to use the `large` whisper model. Alignment models are automatically picked by the chosen language from the default [lists](https://github.com/m-bain/whisperX/blob/main/whisperx/alignment.py#L18).\n\nCurrently support default models tested for {en, fr, de, es, it, ja, zh, nl}\n\n\nIf the detected language is not in this list, you need to find a phoneme-based ASR model from [huggingface model hub](https://huggingface.co/models) and test it on your data.\n\n### French\n whisperx --model large --language fr examples/sample_fr_01.wav\n\n\nhttps://user-images.githubusercontent.com/36994049/208298804-31c49d6f-6787-444e-a53f-e93c52706752.mov\n\n\n### German\n whisperx --model large --language de examples/sample_de_01.wav\n\n\nhttps://user-images.githubusercontent.com/36994049/208298811-e36002ba-3698-4731-97d4-0aebd07e0eb3.mov\n\n\n### Italian\n whisperx --model large --language de examples/sample_it_01.wav\n\n\nhttps://user-images.githubusercontent.com/36994049/208298819-6f462b2c-8cae-4c54-b8e1-90855794efc7.mov\n\n\n### Japanese\n whisperx --model large --language ja examples/sample_ja_01.wav\n\n\nhttps://user-images.githubusercontent.com/19920981/208731743-311f2360-b73b-4c60-809d-aaf3cd7e06f4.mov\n" }, { "path": "LICENSE", "content": "BSD 2-Clause License\n\nCopyright (c) 2024, Max Bain\n\nRedistribution and use in source and binary forms, with or without\nmodification, are permitted provided that the following conditions are met:\n\n1. Redistributions of source code must retain the above copyright notice, this\n list of conditions and the following disclaimer.\n\n2. Redistributions in binary form must reproduce the above copyright notice,\n this list of conditions and the following disclaimer in the documentation\n and/or other materials provided with the distribution.\n\nTHIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS \"AS IS\"\nAND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE\nIMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE\nDISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE\nFOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL\nDAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR\nSERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER\nCAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,\nOR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE\nOF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n" }, { "path": "MANIFEST.in", "content": "include whisperx/assets/*\ninclude LICENSE\ninclude requirements.txt\n" }, { "path": "README.md", "content": "

WhisperX

\n\n## Recall.ai - Meeting Transcription API\n\nIf you’re looking for a transcription API for meetings, consider checking out [Recall.ai's Meeting Transcription API](https://www.recall.ai/product/meeting-transcription-api?utm_source=github&utm_medium=sponsorship&utm_campaign=mbain-whisperx), an API that works with Zoom, Google Meet, Microsoft Teams, and more. Recall.ai diarizes by pulling the speaker data and separate audio streams from the meeting platforms, which means 100% accurate speaker diarization with actual speaker names.\n\n\n

\n \n \"GitHub\n \n \n \"GitHub\n \n \n \"GitHub\n \n \n \"ArXiv\n \n \n \"Twitter\"\n \n

\n\n\"whisperx-arch\"\n\n\n\n\n\nThis repository provides fast automatic speech recognition (70x realtime with large-v2) with word-level timestamps and speaker diarization.\n\n- ⚡️ Batched inference for 70x realtime transcription using whisper large-v2\n- 🪶 [faster-whisper](https://github.com/guillaumekln/faster-whisper) backend, requires <8GB gpu memory for large-v2 with beam_size=5\n- 🎯 Accurate word-level timestamps using wav2vec2 alignment\n- 👯‍♂️ Multispeaker ASR using speaker diarization from [pyannote-audio](https://github.com/pyannote/pyannote-audio) (speaker ID labels)\n- 🗣️ VAD preprocessing, reduces hallucination & batching with no WER degradation\n\n**Whisper** is an ASR model [developed by OpenAI](https://github.com/openai/whisper), trained on a large dataset of diverse audio. Whilst it does produces highly accurate transcriptions, the corresponding timestamps are at the utterance-level, not per word, and can be inaccurate by several seconds. OpenAI's whisper does not natively support batching.\n\n**Phoneme-Based ASR** A suite of models finetuned to recognise the smallest unit of speech distinguishing one word from another, e.g. the element p in \"tap\". A popular example model is [wav2vec2.0](https://huggingface.co/facebook/wav2vec2-large-960h-lv60-self).\n\n**Forced Alignment** refers to the process by which orthographic transcriptions are aligned to audio recordings to automatically generate phone level segmentation.\n\n**Voice Activity Detection (VAD)** is the detection of the presence or absence of human speech.\n\n**Speaker Diarization** is the process of partitioning an audio stream containing human speech into homogeneous segments according to the identity of each speaker.\n\n

New🚨

\n\n- 1st place at [Ego4d transcription challenge](https://eval.ai/web/challenges/challenge-page/1637/leaderboard/3931/WER) 🏆\n- _WhisperX_ accepted at INTERSPEECH 2023\n- v3 transcript segment-per-sentence: using nltk sent_tokenize for better subtitlting & better diarization\n- v3 released, 70x speed-up open-sourced. Using batched whisper with [faster-whisper](https://github.com/guillaumekln/faster-whisper) backend!\n- v2 released, code cleanup, imports whisper library VAD filtering is now turned on by default, as in the paper.\n- Paper drop🎓👨‍🏫! Please see our [ArxiV preprint](https://arxiv.org/abs/2303.00747) for benchmarking and details of WhisperX. We also introduce more efficient batch inference resulting in large-v2 with \\*60-70x REAL TIME speed.\n\n

Setup ⚙️

\n\n### 0. CUDA Installation\n\nTo use WhisperX with GPU acceleration, install the CUDA toolkit 12.8 before WhisperX. Skip this step if using only the CPU.\n\n- For **Linux** users, install the CUDA toolkit 12.8 following this guide:\n [CUDA Installation Guide for Linux](https://docs.nvidia.com/cuda/cuda-installation-guide-linux/).\n- For **Windows** users, download and install the CUDA toolkit 12.8:\n [CUDA Downloads](https://developer.nvidia.com/cuda-12-8-1-download-archive).\n\n### 1. Simple Installation (Recommended)\n\nThe easiest way to install WhisperX is through PyPi:\n\n```bash\npip install whisperx\n```\n\nOr if using [uvx](https://docs.astral.sh/uv/guides/tools/#running-tools):\n\n```bash\nuvx whisperx\n```\n\n### 2. Advanced Installation Options\n\nThese installation methods are for developers or users with specific needs. If you're not sure, stick with the simple installation above.\n\n#### Option A: Install from GitHub\n\nTo install directly from the GitHub repository:\n\n```bash\nuvx git+https://github.com/m-bain/whisperX.git\n```\n\n#### Option B: Developer Installation\n\nIf you want to modify the code or contribute to the project:\n\n```bash\ngit clone https://github.com/m-bain/whisperX.git\ncd whisperX\nuv sync --all-extras --dev\n```\n\n> **Note**: The development version may contain experimental features and bugs. Use the stable PyPI release for production environments.\n\nYou may also need to install ffmpeg, rust etc. Follow openAI instructions here https://github.com/openai/whisper#setup.\n\n### Speaker Diarization\n\nTo **enable Speaker Diarization**, include your Hugging Face access token (read) that you can generate from [Here](https://huggingface.co/settings/tokens) after the `--hf_token` argument and accept the user agreement for the [speaker-diarization-community-1](https://huggingface.co/pyannote/speaker-diarization-community-1) model.\n\n

Usage 💬 (command line)

\n\n### English\n\nRun whisper on example segment (using default params, whisper small) add `--highlight_words True` to visualise word timings in the .srt file.\n\n whisperx path/to/audio.wav\n\nResult using _WhisperX_ with forced alignment to wav2vec2.0 large:\n\nhttps://user-images.githubusercontent.com/36994049/208253969-7e35fe2a-7541-434a-ae91-8e919540555d.mp4\n\nCompare this to original whisper out the box, where many transcriptions are out of sync:\n\nhttps://user-images.githubusercontent.com/36994049/207743923-b4f0d537-29ae-4be2-b404-bb941db73652.mov\n\nFor increased timestamp accuracy, at the cost of higher gpu mem, use bigger models (bigger alignment model not found to be that helpful, see paper) e.g.\n\n whisperx path/to/audio.wav --model large-v2 --align_model WAV2VEC2_ASR_LARGE_LV60K_960H --batch_size 4\n\nTo label the transcript with speaker ID's (set number of speakers if known e.g. `--min_speakers 2` `--max_speakers 2`):\n\n whisperx path/to/audio.wav --model large-v2 --diarize --highlight_words True\n\nTo run on CPU instead of GPU (and for running on Mac OS X):\n\n whisperx path/to/audio.wav --compute_type int8 --device cpu\n\n### Other languages\n\nThe phoneme ASR alignment model is _language-specific_, for tested languages these models are [automatically picked from torchaudio pipelines or huggingface](https://github.com/m-bain/whisperX/blob/f2da2f858e99e4211fe4f64b5f2938b007827e17/whisperx/alignment.py#L24-L58).\nJust pass in the `--language` code, and use the whisper `--model large`.\n\nCurrently default models provided for `{en, fr, de, es, it}` via torchaudio pipelines and many other languages via Hugging Face. Please find the list of currently supported languages under `DEFAULT_ALIGN_MODELS_HF` on [alignment.py](https://github.com/m-bain/whisperX/blob/main/whisperx/alignment.py). If the detected language is not in this list, you need to find a phoneme-based ASR model from [huggingface model hub](https://huggingface.co/models) and test it on your data.\n\n#### E.g. German\n\n whisperx --model large-v2 --language de path/to/audio.wav\n\nhttps://user-images.githubusercontent.com/36994049/208298811-e36002ba-3698-4731-97d4-0aebd07e0eb3.mov\n\nSee more examples in other languages [here](EXAMPLES.md).\n\n## Python usage 🐍\n\n```python\nimport whisperx\nimport gc\nfrom whisperx.diarize import DiarizationPipeline\n\ndevice = \"cuda\"\naudio_file = \"audio.mp3\"\nbatch_size = 16 # reduce if low on GPU mem\ncompute_type = \"float16\" # change to \"int8\" if low on GPU mem (may reduce accuracy)\n\n# 1. Transcribe with original whisper (batched)\nmodel = whisperx.load_model(\"large-v2\", device, compute_type=compute_type)\n\n# save model to local path (optional)\n# model_dir = \"/path/\"\n# model = whisperx.load_model(\"large-v2\", device, compute_type=compute_type, download_root=model_dir)\n\naudio = whisperx.load_audio(audio_file)\nresult = model.transcribe(audio, batch_size=batch_size)\nprint(result[\"segments\"]) # before alignment\n\n# delete model if low on GPU resources\n# import gc; import torch; gc.collect(); torch.cuda.empty_cache(); del model\n\n# 2. Align whisper output\nmodel_a, metadata = whisperx.load_align_model(language_code=result[\"language\"], device=device)\nresult = whisperx.align(result[\"segments\"], model_a, metadata, audio, device, return_char_alignments=False)\n\nprint(result[\"segments\"]) # after alignment\n\n# delete model if low on GPU resources\n# import gc; import torch; gc.collect(); torch.cuda.empty_cache(); del model_a\n\n# 3. Assign speaker labels\ndiarize_model = DiarizationPipeline(token=YOUR_HF_TOKEN, device=device)\n\n# add min/max number of speakers if known\ndiarize_segments = diarize_model(audio)\n# diarize_model(audio, min_speakers=min_speakers, max_speakers=max_speakers)\n\nresult = whisperx.assign_word_speakers(diarize_segments, result)\nprint(diarize_segments)\nprint(result[\"segments\"]) # segments are now assigned speaker IDs\n```\n\n## Demos 🚀\n\n[![Replicate (large-v3](https://img.shields.io/static/v1?label=Replicate+WhisperX+large-v3&message=Demo+%26+Cloud+API&color=blue)](https://replicate.com/victor-upmeet/whisperx)\n[![Replicate (large-v2](https://img.shields.io/static/v1?label=Replicate+WhisperX+large-v2&message=Demo+%26+Cloud+API&color=blue)](https://replicate.com/daanelson/whisperx)\n[![Replicate (medium)](https://img.shields.io/static/v1?label=Replicate+WhisperX+medium&message=Demo+%26+Cloud+API&color=blue)](https://replicate.com/carnifexer/whisperx)\n\nIf you don't have access to your own GPUs, use the links above to try out WhisperX.\n\n

Technical Details 👷‍♂️

\n\nFor specific details on the batching and alignment, the effect of VAD, as well as the chosen alignment model, see the preprint [paper](https://www.robots.ox.ac.uk/~vgg/publications/2023/Bain23/bain23.pdf).\n\nTo reduce GPU memory requirements, try any of the following (2. & 3. can affect quality):\n\n1. reduce batch size, e.g. `--batch_size 4`\n2. use a smaller ASR model `--model base`\n3. Use lighter compute type `--compute_type int8`\n\nTranscription differences from openai's whisper:\n\n1. Transcription without timestamps. To enable single pass batching, whisper inference is performed `--without_timestamps True`, this ensures 1 forward pass per sample in the batch. However, this can cause discrepancies the default whisper output.\n2. VAD-based segment transcription, unlike the buffered transcription of openai's. In the WhisperX paper we show this reduces WER, and enables accurate batched inference\n3. `--condition_on_prev_text` is set to `False` by default (reduces hallucination)\n\n

Limitations ⚠️

\n\n- Transcript words which do not contain characters in the alignment models dictionary e.g. \"2014.\" or \"£13.60\" cannot be aligned and therefore are not given a timing.\n- Overlapping speech is not handled particularly well by whisper nor whisperx\n- Diarization is far from perfect\n- Language specific wav2vec2 model is needed\n\n

Contribute 🧑‍🏫

\n\nIf you are multilingual, a major way you can contribute to this project is to find phoneme models on huggingface (or train your own) and test them on speech for the target language. If the results look good send a pull request and some examples showing its success.\n\nBug finding and pull requests are also highly appreciated to keep this project going, since it's already diverging from the original research scope.\n\n

TODO 🗓

\n\n- [x] Multilingual init\n\n- [x] Automatic align model selection based on language detection\n\n- [x] Python usage\n\n- [x] Incorporating speaker diarization\n\n- [x] Model flush, for low gpu mem resources\n\n- [x] Faster-whisper backend\n\n- [x] Add max-line etc. see (openai's whisper utils.py)\n\n- [x] Sentence-level segments (nltk toolbox)\n\n- [x] Improve alignment logic\n\n- [ ] update examples with diarization and word highlighting\n\n- [ ] Subtitle .ass output <- bring this back (removed in v3)\n\n- [ ] Add benchmarking code (TEDLIUM for spd/WER & word segmentation)\n\n- [x] Allow silero-vad as alternative VAD option\n\n- [ ] Improve diarization (word level). _Harder than first thought..._\n\n

Contact/Support 📇

\n\nContact maxhbain@gmail.com for queries.\n\n\"Buy\n\n

Acknowledgements 🙏

\n\nThis work, and my PhD, is supported by the [VGG (Visual Geometry Group)](https://www.robots.ox.ac.uk/~vgg/) and the University of Oxford.\n\nOf course, this is builds on [openAI's whisper](https://github.com/openai/whisper).\nBorrows important alignment code from [PyTorch tutorial on forced alignment](https://pytorch.org/tutorials/intermediate/forced_alignment_with_torchaudio_tutorial.html)\nAnd uses the wonderful pyannote VAD / Diarization https://github.com/pyannote/pyannote-audio\n\nValuable VAD & Diarization Models from:\n\n- [pyannote-audio](https://github.com/pyannote/pyannote-audio) — Speaker diarization powered by the [speaker-diarization-community-1](https://huggingface.co/pyannote/speaker-diarization-community-1) model, licensed under [CC-BY-4.0](https://creativecommons.org/licenses/by/4.0/) by [pyannoteAI](https://www.pyannote.ai)\n- [silero-vad](https://github.com/snakers4/silero-vad)\n\nGreat backend from [faster-whisper](https://github.com/guillaumekln/faster-whisper) and [CTranslate2](https://github.com/OpenNMT/CTranslate2)\n\nThose who have [supported this work financially](https://www.buymeacoffee.com/maxhbain) 🙏\n\nFinally, thanks to the OS [contributors](https://github.com/m-bain/whisperX/graphs/contributors) of this project, keeping it going and identifying bugs.\n\n

Citation

\nIf you use this in your research, please cite the paper:\n\n```bibtex\n@article{bain2022whisperx,\n title={WhisperX: Time-Accurate Speech Transcription of Long-Form Audio},\n author={Bain, Max and Huh, Jaesung and Han, Tengda and Zisserman, Andrew},\n journal={INTERSPEECH 2023},\n year={2023}\n}\n```\n" }, { "path": "pyproject.toml", "content": "[project]\nurls = { repository = \"https://github.com/m-bain/whisperx\" }\nauthors = [{ name = \"Max Bain\" }]\nname = \"whisperx\"\nversion = \"3.8.2\"\ndescription = \"Time-Accurate Automatic Speech Recognition using Whisper.\"\nreadme = \"README.md\"\nrequires-python = \">=3.10, <3.14\"\nlicense = { text = \"BSD-2-Clause\" }\n\ndependencies = [\n \"ctranslate2>=4.5.0\",\n \"faster-whisper>=1.1.1\",\n \"nltk>=3.9.1\",\n \"numpy>=2.1.0\",\n \"omegaconf>=2.3.0\",\n \"pandas>=2.2.3\",\n \"pyannote-audio>=4.0.0\",\n \"huggingface-hub<1.0.0\",\n \"torch~=2.8.0\",\n \"torchaudio~=2.8.0\",\n \"transformers>=4.48.0\",\n \"triton>=3.3.0; sys_platform == 'linux' and platform_machine == 'x86_64'\" # only install triton on x86_64 Linux\n]\n\n\n[project.scripts]\nwhisperx = \"whisperx.__main__:cli\"\n\n[build-system]\nrequires = [\"setuptools\"]\n\n[tool.setuptools]\ninclude-package-data = true\n\n[tool.setuptools.packages.find]\nwhere = [\".\"]\ninclude = [\"whisperx*\"]\n\n# torchcodec (transitive dep of pyannote-audio >=4) has no wheels for Linux aarch64\n[tool.uv]\noverride-dependencies = [\n \"torchcodec>=0.6.0; (sys_platform == 'linux' and platform_machine == 'x86_64') or sys_platform == 'darwin' or sys_platform == 'win32'\",\n]\n\n[tool.uv.sources]\ntorch = [\n { index = \"pytorch-cpu\", marker = \"sys_platform == 'darwin'\" },\n { index = \"pytorch-cpu\", marker = \"platform_machine != 'x86_64' and sys_platform != 'darwin'\" },\n { index = \"pytorch\", marker = \"platform_machine == 'x86_64' and sys_platform != 'darwin'\" },\n]\ntorchaudio = [\n { index = \"pytorch-cpu\", marker = \"sys_platform == 'darwin'\" },\n { index = \"pytorch-cpu\", marker = \"platform_machine != 'x86_64' and sys_platform != 'darwin'\" },\n { index = \"pytorch\", marker = \"platform_machine == 'x86_64' and sys_platform != 'darwin'\" },\n]\ntriton = [\n { index = \"pytorch\", marker = \"sys_platform == 'linux'\" },\n]\n\n[[tool.uv.index]]\nname = \"pytorch\"\nurl = \"https://download.pytorch.org/whl/cu128\"\nexplicit = true\n\n[[tool.uv.index]]\nname = \"pytorch-cpu\"\nurl = \"https://download.pytorch.org/whl/cpu\"\nexplicit = true\n" }, { "path": "whisperx/SubtitlesProcessor.py", "content": "import math\r\nfrom whisperx.conjunctions import get_conjunctions, get_comma\r\n\r\ndef normal_round(n):\r\n if n - math.floor(n) < 0.5:\r\n return math.floor(n)\r\n return math.ceil(n)\r\n\r\n\r\ndef format_timestamp(seconds: float, is_vtt: bool = False):\r\n\r\n assert seconds >= 0, \"non-negative timestamp expected\"\r\n milliseconds = round(seconds * 1000.0)\r\n\r\n hours = milliseconds // 3_600_000\r\n milliseconds -= hours * 3_600_000\r\n\r\n minutes = milliseconds // 60_000\r\n milliseconds -= minutes * 60_000\r\n\r\n seconds = milliseconds // 1_000\r\n milliseconds -= seconds * 1_000\r\n\r\n separator = '.' if is_vtt else ','\r\n \r\n hours_marker = f\"{hours:02d}:\"\r\n return (\r\n f\"{hours_marker}{minutes:02d}:{seconds:02d}{separator}{milliseconds:03d}\"\r\n )\r\n\r\n\r\n\r\nclass SubtitlesProcessor:\r\n def __init__(self, segments, lang, max_line_length = 45, min_char_length_splitter = 30, is_vtt = False):\r\n self.comma = get_comma(lang)\r\n self.conjunctions = set(get_conjunctions(lang))\r\n self.segments = segments\r\n self.lang = lang\r\n self.max_line_length = max_line_length\r\n self.min_char_length_splitter = min_char_length_splitter\r\n self.is_vtt = is_vtt\r\n complex_script_languages = ['th', 'lo', 'my', 'km', 'am', 'ko', 'ja', 'zh', 'ti', 'ta', 'te', 'kn', 'ml', 'hi', 'ne', 'mr', 'ar', 'fa', 'ur', 'ka']\r\n if self.lang in complex_script_languages:\r\n self.max_line_length = 30\r\n self.min_char_length_splitter = 20\r\n\r\n def estimate_timestamp_for_word(self, words, i, next_segment_start_time=None):\r\n k = 0.25\r\n has_prev_end = i > 0 and 'end' in words[i - 1]\r\n has_next_start = i < len(words) - 1 and 'start' in words[i + 1]\r\n\r\n if has_prev_end:\r\n words[i]['start'] = words[i - 1]['end']\r\n if has_next_start:\r\n words[i]['end'] = words[i + 1]['start']\r\n else:\r\n if next_segment_start_time:\r\n words[i]['end'] = next_segment_start_time if next_segment_start_time - words[i - 1]['end'] <= 1 else next_segment_start_time - 0.5\r\n else:\r\n words[i]['end'] = words[i]['start'] + len(words[i]['word']) * k\r\n\r\n elif has_next_start:\r\n words[i]['start'] = words[i + 1]['start'] - len(words[i]['word']) * k\r\n words[i]['end'] = words[i + 1]['start']\r\n\r\n else:\r\n if next_segment_start_time:\r\n words[i]['start'] = next_segment_start_time - 1\r\n words[i]['end'] = next_segment_start_time - 0.5\r\n else:\r\n words[i]['start'] = 0\r\n words[i]['end'] = 0\r\n\r\n\r\n\r\n def process_segments(self, advanced_splitting=True):\r\n subtitles = []\r\n for i, segment in enumerate(self.segments):\r\n next_segment_start_time = self.segments[i + 1]['start'] if i + 1 < len(self.segments) else None\r\n \r\n if advanced_splitting:\r\n\r\n split_points = self.determine_advanced_split_points(segment, next_segment_start_time)\r\n subtitles.extend(self.generate_subtitles_from_split_points(segment, split_points, next_segment_start_time))\r\n else:\r\n words = segment['words']\r\n for i, word in enumerate(words):\r\n if 'start' not in word or 'end' not in word:\r\n self.estimate_timestamp_for_word(words, i, next_segment_start_time)\r\n\r\n subtitles.append({\r\n 'start': segment['start'],\r\n 'end': segment['end'],\r\n 'text': segment['text']\r\n })\r\n\r\n return subtitles\r\n\r\n def determine_advanced_split_points(self, segment, next_segment_start_time=None):\r\n split_points = []\r\n last_split_point = 0\r\n char_count = 0\r\n\r\n words = segment.get('words', segment['text'].split())\r\n add_space = 0 if self.lang in ['zh', 'ja'] else 1\r\n\r\n total_char_count = sum(len(word['word']) if isinstance(word, dict) else len(word) + add_space for word in words)\r\n char_count_after = total_char_count\r\n\r\n for i, word in enumerate(words):\r\n word_text = word['word'] if isinstance(word, dict) else word\r\n word_length = len(word_text) + add_space\r\n char_count += word_length\r\n char_count_after -= word_length\r\n\r\n char_count_before = char_count - word_length\r\n\r\n if isinstance(word, dict) and ('start' not in word or 'end' not in word):\r\n self.estimate_timestamp_for_word(words, i, next_segment_start_time)\r\n\r\n if char_count >= self.max_line_length:\r\n midpoint = normal_round((last_split_point + i) / 2)\r\n if char_count_before >= self.min_char_length_splitter:\r\n split_points.append(midpoint)\r\n last_split_point = midpoint + 1\r\n char_count = sum(len(words[j]['word']) if isinstance(words[j], dict) else len(words[j]) + add_space for j in range(last_split_point, i + 1))\r\n\r\n elif word_text.endswith(self.comma) and char_count_before >= self.min_char_length_splitter and char_count_after >= self.min_char_length_splitter:\r\n split_points.append(i)\r\n last_split_point = i + 1\r\n char_count = 0\r\n\r\n elif word_text.lower() in self.conjunctions and char_count_before >= self.min_char_length_splitter and char_count_after >= self.min_char_length_splitter:\r\n split_points.append(i - 1)\r\n last_split_point = i\r\n char_count = word_length\r\n\r\n return split_points\r\n\r\n \r\n def generate_subtitles_from_split_points(self, segment, split_points, next_start_time=None):\r\n subtitles = []\r\n \r\n words = segment.get('words', segment['text'].split())\r\n total_word_count = len(words)\r\n total_time = segment['end'] - segment['start']\r\n elapsed_time = segment['start']\r\n prefix = ' ' if self.lang not in ['zh', 'ja'] else ''\r\n start_idx = 0\r\n for split_point in split_points:\r\n\r\n fragment_words = words[start_idx:split_point + 1]\r\n current_word_count = len(fragment_words)\r\n \r\n\r\n if isinstance(fragment_words[0], dict):\r\n start_time = fragment_words[0]['start']\r\n end_time = fragment_words[-1]['end']\r\n next_start_time_for_word = words[split_point + 1]['start'] if split_point + 1 < len(words) else None\r\n if next_start_time_for_word and (next_start_time_for_word - end_time) <= 0.8:\r\n end_time = next_start_time_for_word\r\n else:\r\n fragment = prefix.join(fragment_words).strip()\r\n current_duration = (current_word_count / total_word_count) * total_time\r\n start_time = elapsed_time\r\n end_time = elapsed_time + current_duration\r\n elapsed_time += current_duration\r\n\r\n\r\n subtitles.append({\r\n 'start': start_time,\r\n 'end': end_time,\r\n 'text': fragment if not isinstance(fragment_words[0], dict) else prefix.join(word['word'] for word in fragment_words)\r\n })\r\n \r\n start_idx = split_point + 1\r\n\r\n # Handle the last fragment\r\n if start_idx < len(words):\r\n fragment_words = words[start_idx:]\r\n current_word_count = len(fragment_words)\r\n \r\n if isinstance(fragment_words[0], dict):\r\n start_time = fragment_words[0]['start']\r\n end_time = fragment_words[-1]['end']\r\n else:\r\n fragment = prefix.join(fragment_words).strip()\r\n current_duration = (current_word_count / total_word_count) * total_time\r\n start_time = elapsed_time\r\n end_time = elapsed_time + current_duration\r\n\r\n if next_start_time and (next_start_time - end_time) <= 0.8:\r\n end_time = next_start_time\r\n\r\n subtitles.append({\r\n 'start': start_time,\r\n 'end': end_time if end_time is not None else segment['end'],\r\n 'text': fragment if not isinstance(fragment_words[0], dict) else prefix.join(word['word'] for word in fragment_words)\r\n })\r\n \r\n return subtitles\r\n \r\n\r\n\r\n def save(self, filename=\"subtitles.srt\", advanced_splitting=True):\r\n \r\n subtitles = self.process_segments(advanced_splitting)\r\n\r\n def write_subtitle(file, idx, start_time, end_time, text):\r\n\r\n file.write(f\"{idx}\\n\")\r\n file.write(f\"{start_time} --> {end_time}\\n\")\r\n file.write(text + \"\\n\\n\")\r\n\r\n with open(filename, 'w', encoding='utf-8') as file:\r\n if self.is_vtt:\r\n file.write(\"WEBVTT\\n\\n\")\r\n \r\n if advanced_splitting:\r\n for idx, subtitle in enumerate(subtitles, 1):\r\n start_time = format_timestamp(subtitle['start'], self.is_vtt)\r\n end_time = format_timestamp(subtitle['end'], self.is_vtt)\r\n text = subtitle['text'].strip()\r\n write_subtitle(file, idx, start_time, end_time, text)\r\n\r\n return len(subtitles)" }, { "path": "whisperx/__init__.py", "content": "import importlib\n\n\ndef _lazy_import(name):\n module = importlib.import_module(f\"whisperx.{name}\")\n return module\n\n\ndef load_align_model(*args, **kwargs):\n alignment = _lazy_import(\"alignment\")\n return alignment.load_align_model(*args, **kwargs)\n\n\ndef align(*args, **kwargs):\n alignment = _lazy_import(\"alignment\")\n return alignment.align(*args, **kwargs)\n\n\ndef load_model(*args, **kwargs):\n asr = _lazy_import(\"asr\")\n return asr.load_model(*args, **kwargs)\n\n\ndef load_audio(*args, **kwargs):\n audio = _lazy_import(\"audio\")\n return audio.load_audio(*args, **kwargs)\n\n\ndef assign_word_speakers(*args, **kwargs):\n diarize = _lazy_import(\"diarize\")\n return diarize.assign_word_speakers(*args, **kwargs)\n\n\ndef setup_logging(*args, **kwargs):\n \"\"\"\n Configure logging for WhisperX.\n\n Args:\n level: Logging level (debug, info, warning, error, critical). Default: warning\n log_file: Optional path to log file. If None, logs only to console.\n \"\"\"\n logging_module = _lazy_import(\"log_utils\")\n return logging_module.setup_logging(*args, **kwargs)\n\n\ndef get_logger(*args, **kwargs):\n \"\"\"\n Get a logger instance for the given module.\n\n Args:\n name: Logger name (typically __name__ from calling module)\n\n Returns:\n Logger instance configured with WhisperX settings\n \"\"\"\n logging_module = _lazy_import(\"log_utils\")\n return logging_module.get_logger(*args, **kwargs)\n" }, { "path": "whisperx/__main__.py", "content": "import argparse\nimport importlib.metadata\nimport platform\n\nimport torch\n\nfrom whisperx.utils import (LANGUAGES, TO_LANGUAGE_CODE, optional_float,\n optional_int, str2bool)\nfrom whisperx.log_utils import setup_logging\n\n\ndef cli():\n # fmt: off\n parser = argparse.ArgumentParser(formatter_class=argparse.ArgumentDefaultsHelpFormatter)\n parser.add_argument(\"audio\", nargs=\"+\", type=str, help=\"audio file(s) to transcribe\")\n parser.add_argument(\"--model\", default=\"small\", help=\"name of the Whisper model to use\")\n parser.add_argument(\"--model_cache_only\", type=str2bool, default=False, help=\"If True, will not attempt to download models, instead using cached models from --model_dir\")\n parser.add_argument(\"--model_dir\", type=str, default=None, help=\"the path to save model files; uses ~/.cache/whisper by default\")\n parser.add_argument(\"--device\", default=\"cuda\" if torch.cuda.is_available() else \"cpu\", help=\"device type to use for PyTorch inference (e.g. cpu, cuda)\")\n parser.add_argument(\"--device_index\", default=0, type=int, help=\"device index to use for FasterWhisper inference\")\n parser.add_argument(\"--batch_size\", default=8, type=int, help=\"the preferred batch size for inference\")\n parser.add_argument(\"--compute_type\", default=\"default\", type=str, choices=[\"default\", \"float16\", \"float32\", \"int8\"], help=\"compute type for computation; 'default' uses float16 on GPU, float32 on CPU\")\n\n parser.add_argument(\"--output_dir\", \"-o\", type=str, default=\".\", help=\"directory to save the outputs\")\n parser.add_argument(\"--output_format\", \"-f\", type=str, default=\"all\", choices=[\"all\", \"srt\", \"vtt\", \"txt\", \"tsv\", \"json\", \"aud\"], help=\"format of the output file; if not specified, all available formats will be produced\")\n parser.add_argument(\"--verbose\", type=str2bool, default=True, help=\"whether to print out the progress and debug messages\")\n parser.add_argument(\"--log-level\", type=str, default=None, choices=[\"debug\", \"info\", \"warning\", \"error\", \"critical\"], help=\"logging level (overrides --verbose if set)\")\n\n parser.add_argument(\"--task\", type=str, default=\"transcribe\", choices=[\"transcribe\", \"translate\"], help=\"whether to perform X->X speech recognition ('transcribe') or X->English translation ('translate')\")\n parser.add_argument(\"--language\", type=str, default=None, choices=sorted(LANGUAGES.keys()) + sorted([k.title() for k in TO_LANGUAGE_CODE.keys()]), help=\"language spoken in the audio, specify None to perform language detection\")\n\n # alignment params\n parser.add_argument(\"--align_model\", default=None, help=\"Name of phoneme-level ASR model to do alignment\")\n parser.add_argument(\"--interpolate_method\", default=\"nearest\", choices=[\"nearest\", \"linear\", \"ignore\"], help=\"For word .srt, method to assign timestamps to non-aligned words, or merge them into neighbouring.\")\n parser.add_argument(\"--no_align\", action='store_true', help=\"Do not perform phoneme alignment\")\n parser.add_argument(\"--return_char_alignments\", action='store_true', help=\"Return character-level alignments in the output json file\")\n\n # vad params\n parser.add_argument(\"--vad_method\", type=str, default=\"pyannote\", choices=[\"pyannote\", \"silero\"], help=\"VAD method to be used\")\n parser.add_argument(\"--vad_onset\", type=float, default=0.500, help=\"Onset threshold for VAD (see pyannote.audio), reduce this if speech is not being detected\")\n parser.add_argument(\"--vad_offset\", type=float, default=0.363, help=\"Offset threshold for VAD (see pyannote.audio), reduce this if speech is not being detected.\")\n parser.add_argument(\"--chunk_size\", type=int, default=30, help=\"Chunk size for merging VAD segments. Default is 30, reduce this if the chunk is too long.\")\n\n # diarization params\n parser.add_argument(\"--diarize\", action=\"store_true\", help=\"Apply diarization to assign speaker labels to each segment/word\")\n parser.add_argument(\"--min_speakers\", default=None, type=int, help=\"Minimum number of speakers to in audio file\")\n parser.add_argument(\"--max_speakers\", default=None, type=int, help=\"Maximum number of speakers to in audio file\")\n parser.add_argument(\"--diarize_model\", default=\"pyannote/speaker-diarization-community-1\", type=str, help=\"Name of the speaker diarization model to use\")\n parser.add_argument(\"--speaker_embeddings\", action=\"store_true\", help=\"Include speaker embeddings in JSON output (only works with --diarize)\")\n\n parser.add_argument(\"--temperature\", type=float, default=0, help=\"temperature to use for sampling\")\n parser.add_argument(\"--best_of\", type=optional_int, default=5, help=\"number of candidates when sampling with non-zero temperature\")\n parser.add_argument(\"--beam_size\", type=optional_int, default=5, help=\"number of beams in beam search, only applicable when temperature is zero\")\n parser.add_argument(\"--patience\", type=float, default=1.0, help=\"optional patience value to use in beam decoding, as in https://arxiv.org/abs/2204.05424, the default (1.0) is equivalent to conventional beam search\")\n parser.add_argument(\"--length_penalty\", type=float, default=1.0, help=\"optional token length penalty coefficient (alpha) as in https://arxiv.org/abs/1609.08144, uses simple length normalization by default\")\n\n parser.add_argument(\"--suppress_tokens\", type=str, default=\"-1\", help=\"comma-separated list of token ids to suppress during sampling; '-1' will suppress most special characters except common punctuations\")\n parser.add_argument(\"--suppress_numerals\", action=\"store_true\", help=\"whether to suppress numeric symbols and currency symbols during sampling, since wav2vec2 cannot align them correctly\")\n\n parser.add_argument(\"--initial_prompt\", type=str, default=None, help=\"optional text to provide as a prompt for the first window.\")\n parser.add_argument(\"--hotwords\", type=str, default=None, help=\"hotwords/hint phrases to the model (e.g. \\\"WhisperX, PyAnnote, GPU\\\"); improves recognition of rare/technical terms\")\n parser.add_argument(\"--condition_on_previous_text\", type=str2bool, default=False, help=\"if True, provide the previous output of the model as a prompt for the next window; disabling may make the text inconsistent across windows, but the model becomes less prone to getting stuck in a failure loop\")\n parser.add_argument(\"--fp16\", type=str2bool, default=True, help=\"whether to perform inference in fp16; True by default\")\n\n parser.add_argument(\"--temperature_increment_on_fallback\", type=optional_float, default=0.2, help=\"temperature to increase when falling back when the decoding fails to meet either of the thresholds below\")\n parser.add_argument(\"--compression_ratio_threshold\", type=optional_float, default=2.4, help=\"if the gzip compression ratio is higher than this value, treat the decoding as failed\")\n parser.add_argument(\"--logprob_threshold\", type=optional_float, default=-1.0, help=\"if the average log probability is lower than this value, treat the decoding as failed\")\n parser.add_argument(\"--no_speech_threshold\", type=optional_float, default=0.6, help=\"if the probability of the <|nospeech|> token is higher than this value AND the decoding has failed due to `logprob_threshold`, consider the segment as silence\")\n\n parser.add_argument(\"--max_line_width\", type=optional_int, default=None, help=\"(not possible with --no_align) the maximum number of characters in a line before breaking the line\")\n parser.add_argument(\"--max_line_count\", type=optional_int, default=None, help=\"(not possible with --no_align) the maximum number of lines in a segment\")\n parser.add_argument(\"--highlight_words\", type=str2bool, default=False, help=\"(not possible with --no_align) underline each word as it is spoken in srt and vtt\")\n parser.add_argument(\"--segment_resolution\", type=str, default=\"sentence\", choices=[\"sentence\", \"chunk\"], help=\"(not possible with --no_align) the maximum number of characters in a line before breaking the line\")\n\n parser.add_argument(\"--threads\", type=optional_int, default=0, help=\"number of threads used by torch for CPU inference; supercedes MKL_NUM_THREADS/OMP_NUM_THREADS\")\n\n parser.add_argument(\"--hf_token\", type=str, default=None, help=\"Hugging Face Access Token to access PyAnnote gated models\")\n\n parser.add_argument(\"--print_progress\", type=str2bool, default = False, help = \"if True, progress will be printed in transcribe() and align() methods.\")\n parser.add_argument(\"--version\", \"-V\", action=\"version\", version=f\"%(prog)s {importlib.metadata.version('whisperx')}\",help=\"Show whisperx version information and exit\")\n parser.add_argument(\"--python-version\", \"-P\", action=\"version\", version=f\"Python {platform.python_version()} ({platform.python_implementation()})\",help=\"Show python version information and exit\")\n # fmt: on\n\n args = parser.parse_args().__dict__\n\n log_level = args.get(\"log_level\")\n verbose = args.get(\"verbose\")\n\n if log_level is not None:\n setup_logging(level=log_level)\n elif verbose:\n setup_logging(level=\"info\")\n else:\n setup_logging(level=\"warning\")\n\n from whisperx.transcribe import transcribe_task\n\n transcribe_task(args, parser)\n\n\nif __name__ == \"__main__\":\n cli()\n" }, { "path": "whisperx/alignment.py", "content": "\"\"\"\nForced Alignment with Whisper\nC. Max Bain\n\"\"\"\nfrom dataclasses import dataclass\nfrom typing import Iterable, Optional, Union, List\n\nimport numpy as np\nimport pandas as pd\nimport torch\nimport torchaudio\nfrom transformers import Wav2Vec2ForCTC, Wav2Vec2Processor\n\nfrom whisperx.audio import SAMPLE_RATE, load_audio\nfrom whisperx.utils import interpolate_nans, PUNKT_LANGUAGES\nfrom whisperx.schema import (\n AlignedTranscriptionResult,\n SingleSegment,\n SingleAlignedSegment,\n SingleWordSegment,\n SegmentData,\n ProgressCallback,\n)\nimport nltk\nfrom nltk.data import load as nltk_load\nfrom whisperx.log_utils import get_logger\n\nlogger = get_logger(__name__)\n\nLANGUAGES_WITHOUT_SPACES = [\"ja\", \"zh\"]\n\nDEFAULT_ALIGN_MODELS_TORCH = {\n \"en\": \"WAV2VEC2_ASR_BASE_960H\",\n \"fr\": \"VOXPOPULI_ASR_BASE_10K_FR\",\n \"de\": \"VOXPOPULI_ASR_BASE_10K_DE\",\n \"es\": \"VOXPOPULI_ASR_BASE_10K_ES\",\n \"it\": \"VOXPOPULI_ASR_BASE_10K_IT\",\n}\n\nDEFAULT_ALIGN_MODELS_HF = {\n \"ja\": \"jonatasgrosman/wav2vec2-large-xlsr-53-japanese\",\n \"zh\": \"jonatasgrosman/wav2vec2-large-xlsr-53-chinese-zh-cn\",\n \"nl\": \"jonatasgrosman/wav2vec2-large-xlsr-53-dutch\",\n \"uk\": \"Yehor/wav2vec2-xls-r-300m-uk-with-small-lm\",\n \"pt\": \"jonatasgrosman/wav2vec2-large-xlsr-53-portuguese\",\n \"ar\": \"jonatasgrosman/wav2vec2-large-xlsr-53-arabic\",\n \"cs\": \"comodoro/wav2vec2-xls-r-300m-cs-250\",\n \"ru\": \"jonatasgrosman/wav2vec2-large-xlsr-53-russian\",\n \"pl\": \"jonatasgrosman/wav2vec2-large-xlsr-53-polish\",\n \"hu\": \"jonatasgrosman/wav2vec2-large-xlsr-53-hungarian\",\n \"fi\": \"jonatasgrosman/wav2vec2-large-xlsr-53-finnish\",\n \"fa\": \"jonatasgrosman/wav2vec2-large-xlsr-53-persian\",\n \"el\": \"jonatasgrosman/wav2vec2-large-xlsr-53-greek\",\n \"tr\": \"mpoyraz/wav2vec2-xls-r-300m-cv7-turkish\",\n \"da\": \"saattrupdan/wav2vec2-xls-r-300m-ftspeech\",\n \"he\": \"imvladikon/wav2vec2-xls-r-300m-hebrew\",\n \"vi\": 'nguyenvulebinh/wav2vec2-base-vi-vlsp2020',\n \"ko\": \"kresnik/wav2vec2-large-xlsr-korean\",\n \"ur\": \"kingabzpro/wav2vec2-large-xls-r-300m-Urdu\",\n \"te\": \"anuragshas/wav2vec2-large-xlsr-53-telugu\",\n \"hi\": \"theainerd/Wav2Vec2-large-xlsr-hindi\",\n \"ca\": \"softcatala/wav2vec2-large-xlsr-catala\",\n \"ml\": \"gvs/wav2vec2-large-xlsr-malayalam\",\n \"no\": \"NbAiLab/nb-wav2vec2-1b-bokmaal-v2\",\n \"nn\": \"NbAiLab/nb-wav2vec2-1b-nynorsk\",\n \"sk\": \"comodoro/wav2vec2-xls-r-300m-sk-cv8\",\n \"sl\": \"anton-l/wav2vec2-large-xlsr-53-slovenian\",\n \"hr\": \"classla/wav2vec2-xls-r-parlaspeech-hr\",\n \"ro\": \"gigant/romanian-wav2vec2\",\n \"eu\": \"stefan-it/wav2vec2-large-xlsr-53-basque\",\n \"gl\": \"ifrz/wav2vec2-large-xlsr-galician\",\n \"ka\": \"xsway/wav2vec2-large-xlsr-georgian\",\n \"lv\": \"jimregan/wav2vec2-large-xlsr-latvian-cv\",\n \"tl\": \"Khalsuu/filipino-wav2vec2-l-xls-r-300m-official\",\n \"sv\": \"KBLab/wav2vec2-large-voxrex-swedish\",\n}\n\n\ndef load_align_model(language_code: str, device: str, model_name: Optional[str] = None, model_dir=None, model_cache_only: bool = False):\n if model_name is None:\n # use default model\n if language_code in DEFAULT_ALIGN_MODELS_TORCH:\n model_name = DEFAULT_ALIGN_MODELS_TORCH[language_code]\n elif language_code in DEFAULT_ALIGN_MODELS_HF:\n model_name = DEFAULT_ALIGN_MODELS_HF[language_code]\n else:\n logger.error(f\"No default alignment model for language: {language_code}. \"\n f\"Please find a wav2vec2.0 model finetuned on this language at https://huggingface.co/models, \"\n f\"then pass the model name via --align_model [MODEL_NAME]\")\n raise ValueError(f\"No default align-model for language: {language_code}\")\n\n if model_name in torchaudio.pipelines.__all__:\n pipeline_type = \"torchaudio\"\n bundle = torchaudio.pipelines.__dict__[model_name]\n align_model = bundle.get_model(dl_kwargs={\"model_dir\": model_dir}).to(device)\n labels = bundle.get_labels()\n align_dictionary = {c.lower(): i for i, c in enumerate(labels)}\n else:\n try:\n processor = Wav2Vec2Processor.from_pretrained(model_name, cache_dir=model_dir, local_files_only=model_cache_only)\n align_model = Wav2Vec2ForCTC.from_pretrained(model_name, cache_dir=model_dir, local_files_only=model_cache_only)\n except Exception as e:\n print(e)\n print(f\"Error loading model from huggingface, check https://huggingface.co/models for finetuned wav2vec2.0 models\")\n raise ValueError(f'The chosen align_model \"{model_name}\" could not be found in huggingface (https://huggingface.co/models) or torchaudio (https://pytorch.org/audio/stable/pipelines.html#id14)')\n pipeline_type = \"huggingface\"\n align_model = align_model.to(device)\n labels = processor.tokenizer.get_vocab()\n align_dictionary = {char.lower(): code for char,code in processor.tokenizer.get_vocab().items()}\n\n align_metadata = {\"language\": language_code, \"dictionary\": align_dictionary, \"type\": pipeline_type}\n\n return align_model, align_metadata\n\n\ndef align(\n transcript: Iterable[SingleSegment],\n model: torch.nn.Module,\n align_model_metadata: dict,\n audio: Union[str, np.ndarray, torch.Tensor],\n device: str,\n interpolate_method: str = \"nearest\",\n return_char_alignments: bool = False,\n print_progress: bool = False,\n combined_progress: bool = False,\n progress_callback: ProgressCallback = None,\n) -> AlignedTranscriptionResult:\n \"\"\"\n Align phoneme recognition predictions to known transcription.\n \"\"\"\n\n if not torch.is_tensor(audio):\n if isinstance(audio, str):\n audio = load_audio(audio)\n audio = torch.from_numpy(audio)\n if len(audio.shape) == 1:\n audio = audio.unsqueeze(0)\n\n MAX_DURATION = audio.shape[1] / SAMPLE_RATE\n\n model_dictionary = align_model_metadata[\"dictionary\"]\n model_lang = align_model_metadata[\"language\"]\n model_type = align_model_metadata[\"type\"]\n\n # 1. Preprocess to keep only characters in dictionary\n total_segments = len(transcript)\n # Store temporary processing values\n segment_data: dict[int, SegmentData] = {}\n for sdx, segment in enumerate(transcript):\n # strip spaces at beginning / end, but keep track of the amount.\n if print_progress:\n base_progress = ((sdx + 1) / total_segments) * 100\n percent_complete = (50 + base_progress / 2) if combined_progress else base_progress\n print(f\"Progress: {percent_complete:.2f}%...\")\n\n num_leading = len(segment[\"text\"]) - len(segment[\"text\"].lstrip())\n num_trailing = len(segment[\"text\"]) - len(segment[\"text\"].rstrip())\n text = segment[\"text\"]\n\n # split into words\n if model_lang not in LANGUAGES_WITHOUT_SPACES:\n per_word = text.split(\" \")\n else:\n per_word = text\n\n clean_char, clean_cdx = [], []\n for cdx, char in enumerate(text):\n char_ = char.lower()\n # wav2vec2 models use \"|\" character to represent spaces\n if model_lang not in LANGUAGES_WITHOUT_SPACES:\n char_ = char_.replace(\" \", \"|\")\n\n # ignore whitespace at beginning and end of transcript\n if cdx < num_leading:\n pass\n elif cdx > len(text) - num_trailing - 1:\n pass\n elif char_ in model_dictionary.keys():\n clean_char.append(char_)\n clean_cdx.append(cdx)\n\n clean_wdx = []\n for wdx, wrd in enumerate(per_word):\n if any([c in model_dictionary.keys() for c in wrd.lower()]):\n clean_wdx.append(wdx)\n\n # Use language-specific Punkt model if available otherwise we fallback to English.\n punkt_lang = PUNKT_LANGUAGES.get(model_lang, 'english')\n try:\n sentence_splitter = nltk_load(f'tokenizers/punkt_tab/{punkt_lang}.pickle')\n except LookupError:\n nltk.download('punkt_tab', quiet=True)\n sentence_splitter = nltk_load(f'tokenizers/punkt_tab/{punkt_lang}.pickle')\n sentence_spans = list(sentence_splitter.span_tokenize(text))\n\n segment_data[sdx] = {\n \"clean_char\": clean_char,\n \"clean_cdx\": clean_cdx,\n \"clean_wdx\": clean_wdx,\n \"sentence_spans\": sentence_spans\n }\n\n aligned_segments: List[SingleAlignedSegment] = []\n\n # 2. Get prediction matrix from alignment model & align\n for sdx, segment in enumerate(transcript):\n\n t1 = segment[\"start\"]\n t2 = segment[\"end\"]\n text = segment[\"text\"]\n avg_logprob = segment.get(\"avg_logprob\")\n\n aligned_seg: SingleAlignedSegment = {\n \"start\": t1,\n \"end\": t2,\n \"text\": text,\n \"words\": [],\n \"chars\": None,\n }\n\n if avg_logprob is not None:\n aligned_seg[\"avg_logprob\"] = avg_logprob\n\n if return_char_alignments:\n aligned_seg[\"chars\"] = []\n\n # check we can align\n if len(segment_data[sdx][\"clean_char\"]) == 0:\n logger.warning(f'Failed to align segment (\"{segment[\"text\"]}\"): no characters in this segment found in model dictionary, resorting to original')\n aligned_segments.append(aligned_seg)\n continue\n\n if t1 >= MAX_DURATION:\n logger.warning(f'Failed to align segment (\"{segment[\"text\"]}\"): original start time longer than audio duration, skipping')\n aligned_segments.append(aligned_seg)\n continue\n\n text_clean = \"\".join(segment_data[sdx][\"clean_char\"])\n tokens = [model_dictionary[c] for c in text_clean]\n\n f1 = int(t1 * SAMPLE_RATE)\n f2 = int(t2 * SAMPLE_RATE)\n\n # TODO: Probably can get some speedup gain with batched inference here\n waveform_segment = audio[:, f1:f2]\n # Handle the minimum input length for wav2vec2 models\n if waveform_segment.shape[-1] < 400:\n lengths = torch.as_tensor([waveform_segment.shape[-1]]).to(device)\n waveform_segment = torch.nn.functional.pad(\n waveform_segment, (0, 400 - waveform_segment.shape[-1])\n )\n else:\n lengths = None\n\n with torch.inference_mode():\n if model_type == \"torchaudio\":\n emissions, _ = model(waveform_segment.to(device), lengths=lengths)\n elif model_type == \"huggingface\":\n emissions = model(waveform_segment.to(device)).logits\n else:\n raise NotImplementedError(f\"Align model of type {model_type} not supported.\")\n emissions = torch.log_softmax(emissions, dim=-1)\n\n emission = emissions[0].cpu().detach()\n\n blank_id = 0\n for char, code in model_dictionary.items():\n if char == '[pad]' or char == '':\n blank_id = code\n\n trellis = get_trellis(emission, tokens, blank_id)\n path = backtrack(trellis, emission, tokens, blank_id)\n\n if path is None:\n logger.warning(f'Failed to align segment (\"{segment[\"text\"]}\"): backtrack failed, resorting to original')\n aligned_segments.append(aligned_seg)\n continue\n\n char_segments = merge_repeats(path, text_clean)\n\n duration = t2 - t1\n ratio = duration * waveform_segment.size(0) / (trellis.size(0) - 1)\n\n # assign timestamps to aligned characters\n char_segments_arr = []\n word_idx = 0\n for cdx, char in enumerate(text):\n start, end, score = None, None, None\n if cdx in segment_data[sdx][\"clean_cdx\"]:\n char_seg = char_segments[segment_data[sdx][\"clean_cdx\"].index(cdx)]\n start = round(char_seg.start * ratio + t1, 3)\n end = round(char_seg.end * ratio + t1, 3)\n score = round(char_seg.score, 3)\n\n char_segments_arr.append(\n {\n \"char\": char,\n \"start\": start,\n \"end\": end,\n \"score\": score,\n \"word-idx\": word_idx,\n }\n )\n\n # increment word_idx, nltk word tokenization would probably be more robust here, but us space for now...\n if model_lang in LANGUAGES_WITHOUT_SPACES:\n word_idx += 1\n elif cdx == len(text) - 1 or text[cdx+1] == \" \":\n word_idx += 1\n\n char_segments_arr = pd.DataFrame(char_segments_arr)\n\n aligned_subsegments = []\n # assign sentence_idx to each character index\n char_segments_arr[\"sentence-idx\"] = None\n for sdx2, (sstart, send) in enumerate(segment_data[sdx][\"sentence_spans\"]):\n curr_chars = char_segments_arr.loc[(char_segments_arr.index >= sstart) & (char_segments_arr.index <= send)]\n char_segments_arr.loc[(char_segments_arr.index >= sstart) & (char_segments_arr.index <= send), \"sentence-idx\"] = sdx2\n\n sentence_text = text[sstart:send]\n sentence_start = curr_chars[\"start\"].min()\n end_chars = curr_chars[curr_chars[\"char\"] != ' ']\n sentence_end = end_chars[\"end\"].max()\n sentence_words = []\n\n for word_idx in curr_chars[\"word-idx\"].unique():\n word_chars = curr_chars.loc[curr_chars[\"word-idx\"] == word_idx]\n word_text = \"\".join(word_chars[\"char\"].tolist()).strip()\n if len(word_text) == 0:\n continue\n\n # dont use space character for alignment\n word_chars = word_chars[word_chars[\"char\"] != \" \"]\n\n word_start = word_chars[\"start\"].min()\n word_end = word_chars[\"end\"].max()\n word_score = round(word_chars[\"score\"].mean(), 3)\n\n # -1 indicates unalignable\n word_segment = {\"word\": word_text}\n\n if not np.isnan(word_start):\n word_segment[\"start\"] = word_start\n if not np.isnan(word_end):\n word_segment[\"end\"] = word_end\n if not np.isnan(word_score):\n word_segment[\"score\"] = word_score\n\n sentence_words.append(word_segment)\n\n subsegment = {\n \"text\": sentence_text,\n \"start\": sentence_start,\n \"end\": sentence_end,\n \"words\": sentence_words,\n }\n if avg_logprob is not None:\n subsegment[\"avg_logprob\"] = avg_logprob\n aligned_subsegments.append(subsegment)\n\n if return_char_alignments:\n curr_chars = curr_chars[[\"char\", \"start\", \"end\", \"score\"]]\n curr_chars.fillna(-1, inplace=True)\n curr_chars = curr_chars.to_dict(\"records\")\n curr_chars = [{key: val for key, val in char.items() if val != -1} for char in curr_chars]\n aligned_subsegments[-1][\"chars\"] = curr_chars\n\n aligned_subsegments = pd.DataFrame(aligned_subsegments)\n aligned_subsegments[\"start\"] = interpolate_nans(aligned_subsegments[\"start\"], method=interpolate_method)\n aligned_subsegments[\"end\"] = interpolate_nans(aligned_subsegments[\"end\"], method=interpolate_method)\n # concatenate sentences with same timestamps\n agg_dict = {\"text\": \" \".join, \"words\": \"sum\"}\n if model_lang in LANGUAGES_WITHOUT_SPACES:\n agg_dict[\"text\"] = \"\".join\n if return_char_alignments:\n agg_dict[\"chars\"] = \"sum\"\n if avg_logprob is not None:\n agg_dict[\"avg_logprob\"] = \"first\"\n aligned_subsegments= aligned_subsegments.groupby([\"start\", \"end\"], as_index=False).agg(agg_dict)\n aligned_subsegments = aligned_subsegments.to_dict('records')\n if progress_callback is not None:\n progress_callback(((sdx + 1) / total_segments) * 100)\n\n aligned_segments += aligned_subsegments\n\n # create word_segments list\n word_segments: List[SingleWordSegment] = []\n for segment in aligned_segments:\n word_segments += segment[\"words\"]\n\n return {\"segments\": aligned_segments, \"word_segments\": word_segments}\n\n\"\"\"\nsource: https://pytorch.org/tutorials/intermediate/forced_alignment_with_torchaudio_tutorial.html\n\"\"\"\n\n\ndef get_trellis(emission, tokens, blank_id=0):\n num_frame = emission.size(0)\n num_tokens = len(tokens)\n\n # Trellis has extra dimensions for both time axis and tokens.\n # The extra dim for tokens represents (start-of-sentence)\n # The extra dim for time axis is for simplification of the code.\n trellis = torch.empty((num_frame + 1, num_tokens + 1))\n trellis[0, 0] = 0\n trellis[1:, 0] = torch.cumsum(emission[:, blank_id], 0)\n trellis[0, -num_tokens:] = -float(\"inf\")\n trellis[-num_tokens:, 0] = float(\"inf\")\n\n for t in range(num_frame):\n trellis[t + 1, 1:] = torch.maximum(\n # Score for staying at the same token\n trellis[t, 1:] + emission[t, blank_id],\n # Score for changing to the next token\n trellis[t, :-1] + emission[t, tokens],\n )\n return trellis\n\n\n@dataclass\nclass Point:\n token_index: int\n time_index: int\n score: float\n\n\ndef backtrack(trellis, emission, tokens, blank_id=0):\n # Note:\n # j and t are indices for trellis, which has extra dimensions\n # for time and tokens at the beginning.\n # When referring to time frame index `T` in trellis,\n # the corresponding index in emission is `T-1`.\n # Similarly, when referring to token index `J` in trellis,\n # the corresponding index in transcript is `J-1`.\n j = trellis.size(1) - 1\n t_start = torch.argmax(trellis[:, j]).item()\n\n path = []\n for t in range(t_start, 0, -1):\n # 1. Figure out if the current position was stay or change\n # Note (again):\n # `emission[J-1]` is the emission at time frame `J` of trellis dimension.\n # Score for token staying the same from time frame J-1 to T.\n stayed = trellis[t - 1, j] + emission[t - 1, blank_id]\n # Score for token changing from C-1 at T-1 to J at T.\n changed = trellis[t - 1, j - 1] + emission[t - 1, tokens[j - 1]]\n\n # 2. Store the path with frame-wise probability.\n prob = emission[t - 1, tokens[j - 1] if changed > stayed else blank_id].exp().item()\n # Return token index and time index in non-trellis coordinate.\n path.append(Point(j - 1, t - 1, prob))\n\n # 3. Update the token\n if changed > stayed:\n j -= 1\n if j == 0:\n break\n else:\n # failed\n return None\n\n return path[::-1]\n\n\n# Merge the labels\n@dataclass\nclass Segment:\n label: str\n start: int\n end: int\n score: float\n\n def __repr__(self):\n return f\"{self.label}\\t({self.score:4.2f}): [{self.start:5d}, {self.end:5d})\"\n\n @property\n def length(self):\n return self.end - self.start\n\ndef merge_repeats(path, transcript):\n i1, i2 = 0, 0\n segments = []\n while i1 < len(path):\n while i2 < len(path) and path[i1].token_index == path[i2].token_index:\n i2 += 1\n score = sum(path[k].score for k in range(i1, i2)) / (i2 - i1)\n segments.append(\n Segment(\n transcript[path[i1].token_index],\n path[i1].time_index,\n path[i2 - 1].time_index + 1,\n score,\n )\n )\n i1 = i2\n return segments\n\ndef merge_words(segments, separator=\"|\"):\n words = []\n i1, i2 = 0, 0\n while i1 < len(segments):\n if i2 >= len(segments) or segments[i2].label == separator:\n if i1 != i2:\n segs = segments[i1:i2]\n word = \"\".join([seg.label for seg in segs])\n score = sum(seg.score * seg.length for seg in segs) / sum(seg.length for seg in segs)\n words.append(Segment(word, segments[i1].start, segments[i2 - 1].end, score))\n i1 = i2 + 1\n i2 = i1\n else:\n i2 += 1\n return words" }, { "path": "whisperx/asr.py", "content": "import os\nfrom typing import List, Optional, Union\nfrom dataclasses import replace\n\nimport ctranslate2\nimport faster_whisper\nimport numpy as np\nimport torch\nfrom faster_whisper.tokenizer import Tokenizer\nfrom faster_whisper.transcribe import TranscriptionOptions, get_ctranslate2_storage\nfrom transformers import Pipeline\nfrom transformers.pipelines.pt_utils import PipelineIterator\n\nfrom whisperx.audio import N_SAMPLES, SAMPLE_RATE, load_audio, log_mel_spectrogram\nfrom whisperx.schema import SingleSegment, TranscriptionResult, ProgressCallback\nfrom whisperx.vads import Vad, Silero, Pyannote\nfrom whisperx.log_utils import get_logger\n\nlogger = get_logger(__name__)\n\n\ndef find_numeral_symbol_tokens(tokenizer):\n numeral_symbol_tokens = []\n for i in range(tokenizer.eot):\n token = tokenizer.decode([i]).removeprefix(\" \")\n has_numeral_symbol = any(c in \"0123456789%$£\" for c in token)\n if has_numeral_symbol:\n numeral_symbol_tokens.append(i)\n return numeral_symbol_tokens\n\nclass WhisperModel(faster_whisper.WhisperModel):\n '''\n FasterWhisperModel provides batched inference for faster-whisper.\n Currently only works in non-timestamp mode and fixed prompt for all samples in batch.\n '''\n\n def generate_segment_batched(\n self,\n features: np.ndarray,\n tokenizer: Tokenizer,\n options: TranscriptionOptions,\n encoder_output=None,\n ):\n batch_size = features.shape[0]\n all_tokens = []\n prompt_reset_since = 0\n if options.initial_prompt is not None:\n initial_prompt = \" \" + options.initial_prompt.strip()\n initial_prompt_tokens = tokenizer.encode(initial_prompt)\n all_tokens.extend(initial_prompt_tokens)\n previous_tokens = all_tokens[prompt_reset_since:]\n prompt = self.get_prompt(\n tokenizer,\n previous_tokens,\n without_timestamps=options.without_timestamps,\n prefix=options.prefix,\n hotwords=options.hotwords\n )\n\n encoder_output = self.encode(features)\n \n result = self.model.generate(\n encoder_output,\n [prompt] * batch_size,\n beam_size=options.beam_size,\n patience=options.patience,\n length_penalty=options.length_penalty,\n max_length=self.max_length,\n suppress_blank=options.suppress_blank,\n suppress_tokens=options.suppress_tokens,\n no_repeat_ngram_size=options.no_repeat_ngram_size,\n repetition_penalty=options.repetition_penalty,\n return_scores=True,\n )\n\n tokens_batch = [x.sequences_ids[0] for x in result]\n\n avg_logprobs = []\n for res in result:\n seq_len = len(res.sequences_ids[0])\n cum_logprob = res.scores[0] * (seq_len ** options.length_penalty)\n avg_logprobs.append(cum_logprob / (seq_len + 1))\n\n def decode_batch(tokens: List[List[int]]) -> List[str]:\n res = []\n for tk in tokens:\n res.append([token for token in tk if token < tokenizer.eot])\n # text_tokens = [token for token in tokens if token < self.eot]\n return tokenizer.tokenizer.decode_batch(res)\n\n text = decode_batch(tokens_batch)\n\n return {'text': text, 'avg_logprob': avg_logprobs}\n\n def encode(self, features: np.ndarray) -> ctranslate2.StorageView:\n # When the model is running on multiple GPUs, the encoder output should be moved\n # to the CPU since we don't know which GPU will handle the next job.\n to_cpu = self.model.device == \"cuda\" and len(self.model.device_index) > 1\n # unsqueeze if batch size = 1\n if len(features.shape) == 2:\n features = np.expand_dims(features, 0)\n features = get_ctranslate2_storage(features)\n\n return self.model.encode(features, to_cpu=to_cpu)\n\nclass FasterWhisperPipeline(Pipeline):\n \"\"\"\n Huggingface Pipeline wrapper for FasterWhisperModel.\n \"\"\"\n # TODO:\n # - add support for timestamp mode\n # - add support for custom inference kwargs\n\n def __init__(\n self,\n model: WhisperModel,\n vad,\n vad_params: dict,\n options: TranscriptionOptions,\n tokenizer: Optional[Tokenizer] = None,\n device: Union[int, str, \"torch.device\"] = -1,\n framework=\"pt\",\n language: Optional[str] = None,\n suppress_numerals: bool = False,\n **kwargs,\n ):\n self.model = model\n self.tokenizer = tokenizer\n self.options = options\n self.preset_language = language\n self.suppress_numerals = suppress_numerals\n self._batch_size = kwargs.pop(\"batch_size\", None)\n self._num_workers = 1\n self._preprocess_params, self._forward_params, self._postprocess_params = self._sanitize_parameters(**kwargs)\n self.call_count = 0\n self.framework = framework\n if self.framework == \"pt\":\n if isinstance(device, torch.device):\n self.device = device\n elif isinstance(device, str):\n self.device = torch.device(device)\n elif device < 0:\n self.device = torch.device(\"cpu\")\n else:\n self.device = torch.device(f\"cuda:{device}\")\n else:\n self.device = device\n\n super(Pipeline, self).__init__()\n self.vad_model = vad\n self._vad_params = vad_params\n\n def _sanitize_parameters(self, **kwargs):\n preprocess_kwargs = {}\n if \"tokenizer\" in kwargs:\n preprocess_kwargs[\"maybe_arg\"] = kwargs[\"maybe_arg\"]\n return preprocess_kwargs, {}, {}\n\n def preprocess(self, audio):\n audio = audio['inputs']\n model_n_mels = self.model.feat_kwargs.get(\"feature_size\")\n features = log_mel_spectrogram(\n audio,\n n_mels=model_n_mels if model_n_mels is not None else 80,\n padding=N_SAMPLES - audio.shape[0],\n )\n return {'inputs': features}\n\n def _forward(self, model_inputs):\n outputs = self.model.generate_segment_batched(model_inputs['inputs'], self.tokenizer, self.options)\n return outputs\n\n def postprocess(self, model_outputs):\n return model_outputs\n\n def get_iterator(\n self,\n inputs,\n num_workers: int,\n batch_size: int,\n preprocess_params: dict,\n forward_params: dict,\n postprocess_params: dict,\n ):\n dataset = PipelineIterator(inputs, self.preprocess, preprocess_params)\n if \"TOKENIZERS_PARALLELISM\" not in os.environ:\n os.environ[\"TOKENIZERS_PARALLELISM\"] = \"false\"\n # TODO hack by collating feature_extractor and image_processor\n\n def stack(items):\n return {'inputs': torch.stack([x['inputs'] for x in items])}\n dataloader = torch.utils.data.DataLoader(dataset, num_workers=num_workers, batch_size=batch_size, collate_fn=stack)\n model_iterator = PipelineIterator(dataloader, self.forward, forward_params, loader_batch_size=batch_size)\n final_iterator = PipelineIterator(model_iterator, self.postprocess, postprocess_params)\n return final_iterator\n\n def transcribe(\n self,\n audio: Union[str, np.ndarray],\n batch_size: Optional[int] = None,\n num_workers=0,\n language: Optional[str] = None,\n task: Optional[str] = None,\n chunk_size=30,\n print_progress=False,\n combined_progress=False,\n verbose=False,\n progress_callback: ProgressCallback = None,\n ) -> TranscriptionResult:\n if isinstance(audio, str):\n audio = load_audio(audio)\n\n def data(audio, segments):\n for seg in segments:\n f1 = int(seg['start'] * SAMPLE_RATE)\n f2 = int(seg['end'] * SAMPLE_RATE)\n # print(f2-f1)\n yield {'inputs': audio[f1:f2]}\n\n # Pre-process audio and merge chunks as defined by the respective VAD child class \n # In case vad_model is manually assigned (see 'load_model') follow the functionality of pyannote toolkit\n if issubclass(type(self.vad_model), Vad):\n waveform = self.vad_model.preprocess_audio(audio)\n merge_chunks = self.vad_model.merge_chunks\n else:\n waveform = Pyannote.preprocess_audio(audio)\n merge_chunks = Pyannote.merge_chunks\n\n vad_segments = self.vad_model({\"waveform\": waveform, \"sample_rate\": SAMPLE_RATE})\n vad_segments = merge_chunks(\n vad_segments,\n chunk_size,\n onset=self._vad_params[\"vad_onset\"],\n offset=self._vad_params[\"vad_offset\"],\n )\n if self.tokenizer is None:\n language = language or self.detect_language(audio)\n task = task or \"transcribe\"\n self.tokenizer = Tokenizer(\n self.model.hf_tokenizer,\n self.model.model.is_multilingual,\n task=task,\n language=language,\n )\n else:\n language = language or self.tokenizer.language_code\n task = task or self.tokenizer.task\n if task != self.tokenizer.task or language != self.tokenizer.language_code:\n self.tokenizer = Tokenizer(\n self.model.hf_tokenizer,\n self.model.model.is_multilingual,\n task=task,\n language=language,\n )\n\n if self.suppress_numerals:\n previous_suppress_tokens = self.options.suppress_tokens\n numeral_symbol_tokens = find_numeral_symbol_tokens(self.tokenizer)\n logger.info(\"Suppressing numeral and symbol tokens\")\n new_suppressed_tokens = numeral_symbol_tokens + self.options.suppress_tokens\n new_suppressed_tokens = list(set(new_suppressed_tokens))\n self.options = replace(self.options, suppress_tokens=new_suppressed_tokens)\n\n segments: List[SingleSegment] = []\n batch_size = batch_size or self._batch_size\n total_segments = len(vad_segments)\n for idx, out in enumerate(self.__call__(data(audio, vad_segments), batch_size=batch_size, num_workers=num_workers)):\n if print_progress:\n base_progress = ((idx + 1) / total_segments) * 100\n percent_complete = base_progress / 2 if combined_progress else base_progress\n print(f\"Progress: {percent_complete:.2f}%...\")\n if progress_callback is not None:\n progress_callback(((idx + 1) / total_segments) * 100)\n text = out['text']\n avg_logprob = out['avg_logprob']\n if batch_size in [0, 1, None]:\n text = text[0]\n avg_logprob = avg_logprob[0]\n if verbose:\n print(f\"Transcript: [{round(vad_segments[idx]['start'], 3)} --> {round(vad_segments[idx]['end'], 3)}] {text}\")\n segments.append(\n {\n \"text\": text,\n \"start\": round(vad_segments[idx]['start'], 3),\n \"end\": round(vad_segments[idx]['end'], 3),\n \"avg_logprob\": avg_logprob,\n }\n )\n\n # revert the tokenizer if multilingual inference is enabled\n if self.preset_language is None:\n self.tokenizer = None\n\n # revert suppressed tokens if suppress_numerals is enabled\n if self.suppress_numerals:\n self.options = replace(self.options, suppress_tokens=previous_suppress_tokens)\n\n return {\"segments\": segments, \"language\": language}\n\n def detect_language(self, audio: np.ndarray) -> str:\n if audio.shape[0] < N_SAMPLES:\n logger.warning(\"Audio is shorter than 30s, language detection may be inaccurate\")\n model_n_mels = self.model.feat_kwargs.get(\"feature_size\")\n segment = log_mel_spectrogram(audio[: N_SAMPLES],\n n_mels=model_n_mels if model_n_mels is not None else 80,\n padding=0 if audio.shape[0] >= N_SAMPLES else N_SAMPLES - audio.shape[0])\n encoder_output = self.model.encode(segment)\n results = self.model.model.detect_language(encoder_output)\n language_token, language_probability = results[0][0]\n language = language_token[2:-2]\n logger.info(f\"Detected language: {language} ({language_probability:.2f}) in first 30s of audio\")\n return language\n\n\ndef load_model(\n whisper_arch: str,\n device: str,\n device_index=0,\n compute_type=\"default\",\n asr_options: Optional[dict] = None,\n language: Optional[str] = None,\n vad_model: Optional[Vad]= None,\n vad_method: Optional[str] = \"pyannote\",\n vad_options: Optional[dict] = None,\n model: Optional[WhisperModel] = None,\n task=\"transcribe\",\n download_root: Optional[str] = None,\n local_files_only=False,\n threads=4,\n use_auth_token: Optional[Union[str, bool]] = None,\n) -> FasterWhisperPipeline:\n \"\"\"Load a Whisper model for inference.\n Args:\n whisper_arch - The name of the Whisper model to load.\n device - The device to load the model on.\n compute_type - The compute type to use for the model.\n Use \"default\" to automatically select based on device (float16 for GPU, float32 for CPU).\n vad_model - The vad model to manually assign.\n vad_method - The vad method to use. vad_model has a higher priority if it is not None.\n options - A dictionary of options to use for the model.\n language - The language of the model. (use English for now)\n model - The WhisperModel instance to use.\n download_root - The root directory to download the model to.\n local_files_only - If `True`, avoid downloading the file and return the path to the local cached file if it exists.\n threads - The number of cpu threads to use per worker, e.g. will be multiplied by num workers.\n Returns:\n A Whisper pipeline.\n \"\"\"\n\n if compute_type == \"default\":\n compute_type = \"float16\" if device == \"cuda\" else \"float32\"\n logger.info(f\"Compute type not specified, defaulting to {compute_type} for device {device}\")\n\n if whisper_arch.endswith(\".en\"):\n language = \"en\"\n\n model = model or WhisperModel(whisper_arch,\n device=device,\n device_index=device_index,\n compute_type=compute_type,\n download_root=download_root,\n local_files_only=local_files_only,\n cpu_threads=threads,\n use_auth_token=use_auth_token)\n if language is not None:\n tokenizer = Tokenizer(model.hf_tokenizer, model.model.is_multilingual, task=task, language=language)\n else:\n logger.info(\"No language specified, language will be detected for each audio file (increases inference time)\")\n tokenizer = None\n\n default_asr_options = {\n \"beam_size\": 5,\n \"best_of\": 5,\n \"patience\": 1,\n \"length_penalty\": 1,\n \"repetition_penalty\": 1,\n \"no_repeat_ngram_size\": 0,\n \"temperatures\": [0.0, 0.2, 0.4, 0.6, 0.8, 1.0],\n \"compression_ratio_threshold\": 2.4,\n \"log_prob_threshold\": -1.0,\n \"no_speech_threshold\": 0.6,\n \"condition_on_previous_text\": False,\n \"prompt_reset_on_temperature\": 0.5,\n \"initial_prompt\": None,\n \"prefix\": None,\n \"suppress_blank\": True,\n \"suppress_tokens\": [-1],\n \"without_timestamps\": True,\n \"max_initial_timestamp\": 0.0,\n \"word_timestamps\": False,\n \"prepend_punctuations\": \"\\\"'“¿([{-\",\n \"append_punctuations\": \"\\\"'.。,，!！?？:：”)]}、\",\n \"multilingual\": model.model.is_multilingual,\n \"suppress_numerals\": False,\n \"max_new_tokens\": None,\n \"clip_timestamps\": None,\n \"hallucination_silence_threshold\": None,\n \"hotwords\": None,\n }\n\n if asr_options is not None:\n default_asr_options.update(asr_options)\n\n suppress_numerals = default_asr_options[\"suppress_numerals\"]\n del default_asr_options[\"suppress_numerals\"]\n\n default_asr_options = TranscriptionOptions(**default_asr_options)\n\n default_vad_options = {\n \"chunk_size\": 30, # needed by silero since binarization happens before merge_chunks\n \"vad_onset\": 0.500,\n \"vad_offset\": 0.363\n }\n\n if vad_options is not None:\n default_vad_options.update(vad_options)\n\n # Note: manually assigned vad_model has higher priority than vad_method!\n if vad_model is not None:\n print(\"Use manually assigned vad_model. vad_method is ignored.\")\n vad_model = vad_model\n else:\n if vad_method == \"silero\":\n vad_model = Silero(**default_vad_options)\n elif vad_method == \"pyannote\":\n if device == 'cuda':\n device_vad = f'cuda:{device_index}'\n else:\n device_vad = device\n vad_model = Pyannote(torch.device(device_vad), token=None, **default_vad_options)\n else:\n raise ValueError(f\"Invalid vad_method: {vad_method}\")\n\n return FasterWhisperPipeline(\n model=model,\n vad=vad_model,\n options=default_asr_options,\n tokenizer=tokenizer,\n language=language,\n suppress_numerals=suppress_numerals,\n vad_params=default_vad_options,\n )\n" }, { "path": "whisperx/audio.py", "content": "import os\nimport subprocess\nfrom functools import lru_cache\nfrom typing import Optional, Union\n\nimport numpy as np\nimport torch\nimport torch.nn.functional as F\n\nfrom whisperx.utils import exact_div\n\n# hard-coded audio hyperparameters\nSAMPLE_RATE = 16000\nN_FFT = 400\nHOP_LENGTH = 160\nCHUNK_LENGTH = 30\nN_SAMPLES = CHUNK_LENGTH * SAMPLE_RATE # 480000 samples in a 30-second chunk\nN_FRAMES = exact_div(N_SAMPLES, HOP_LENGTH) # 3000 frames in a mel spectrogram input\n\nN_SAMPLES_PER_TOKEN = HOP_LENGTH * 2 # the initial convolutions has stride 2\nFRAMES_PER_SECOND = exact_div(SAMPLE_RATE, HOP_LENGTH) # 10ms per audio frame\nTOKENS_PER_SECOND = exact_div(SAMPLE_RATE, N_SAMPLES_PER_TOKEN) # 20ms per audio token\n\n\ndef load_audio(file: str, sr: int = SAMPLE_RATE) -> np.ndarray:\n \"\"\"\n Open an audio file and read as mono waveform, resampling as necessary\n\n Parameters\n ----------\n file: str\n The audio file to open\n\n sr: int\n The sample rate to resample the audio if necessary\n\n Returns\n -------\n A NumPy array containing the audio waveform, in float32 dtype.\n \"\"\"\n try:\n # Launches a subprocess to decode audio while down-mixing and resampling as necessary.\n # Requires the ffmpeg CLI to be installed.\n cmd = [\n \"ffmpeg\",\n \"-nostdin\",\n \"-threads\",\n \"0\",\n \"-i\",\n file,\n \"-f\",\n \"s16le\",\n \"-ac\",\n \"1\",\n \"-acodec\",\n \"pcm_s16le\",\n \"-ar\",\n str(sr),\n \"-\",\n ]\n out = subprocess.run(cmd, capture_output=True, check=True).stdout\n except subprocess.CalledProcessError as e:\n raise RuntimeError(f\"Failed to load audio: {e.stderr.decode()}\") from e\n\n return np.frombuffer(out, np.int16).flatten().astype(np.float32) / 32768.0\n\n\ndef pad_or_trim(array, length: int = N_SAMPLES, *, axis: int = -1):\n \"\"\"\n Pad or trim the audio array to N_SAMPLES, as expected by the encoder.\n \"\"\"\n if torch.is_tensor(array):\n if array.shape[axis] > length:\n array = array.index_select(\n dim=axis, index=torch.arange(length, device=array.device)\n )\n\n if array.shape[axis] < length:\n pad_widths = [(0, 0)] * array.ndim\n pad_widths[axis] = (0, length - array.shape[axis])\n array = F.pad(array, [pad for sizes in pad_widths[::-1] for pad in sizes])\n else:\n if array.shape[axis] > length:\n array = array.take(indices=range(length), axis=axis)\n\n if array.shape[axis] < length:\n pad_widths = [(0, 0)] * array.ndim\n pad_widths[axis] = (0, length - array.shape[axis])\n array = np.pad(array, pad_widths)\n\n return array\n\n\n@lru_cache(maxsize=None)\ndef mel_filters(device, n_mels: int) -> torch.Tensor:\n \"\"\"\n load the mel filterbank matrix for projecting STFT into a Mel spectrogram.\n Allows decoupling librosa dependency; saved using:\n\n np.savez_compressed(\n \"mel_filters.npz\",\n mel_80=librosa.filters.mel(sr=16000, n_fft=400, n_mels=80),\n )\n \"\"\"\n assert n_mels in [80, 128], f\"Unsupported n_mels: {n_mels}\"\n with np.load(\n os.path.join(os.path.dirname(__file__), \"assets\", \"mel_filters.npz\")\n ) as f:\n return torch.from_numpy(f[f\"mel_{n_mels}\"]).to(device)\n\n\ndef log_mel_spectrogram(\n audio: Union[str, np.ndarray, torch.Tensor],\n n_mels: int,\n padding: int = 0,\n device: Optional[Union[str, torch.device]] = None,\n):\n \"\"\"\n Compute the log-Mel spectrogram of\n\n Parameters\n ----------\n audio: Union[str, np.ndarray, torch.Tensor], shape = (*)\n The path to audio or either a NumPy array or Tensor containing the audio waveform in 16 kHz\n\n n_mels: int\n The number of Mel-frequency filters, only 80 is supported\n\n padding: int\n Number of zero samples to pad to the right\n\n device: Optional[Union[str, torch.device]]\n If given, the audio tensor is moved to this device before STFT\n\n Returns\n -------\n torch.Tensor, shape = (80, n_frames)\n A Tensor that contains the Mel spectrogram\n \"\"\"\n if not torch.is_tensor(audio):\n if isinstance(audio, str):\n audio = load_audio(audio)\n audio = torch.from_numpy(audio)\n\n if device is not None:\n audio = audio.to(device)\n if padding > 0:\n audio = F.pad(audio, (0, padding))\n window = torch.hann_window(N_FFT).to(audio.device)\n stft = torch.stft(audio, N_FFT, HOP_LENGTH, window=window, return_complex=True)\n magnitudes = stft[..., :-1].abs() ** 2\n\n filters = mel_filters(audio.device, n_mels)\n mel_spec = filters @ magnitudes\n\n log_spec = torch.clamp(mel_spec, min=1e-10).log10()\n log_spec = torch.maximum(log_spec, log_spec.max() - 8.0)\n log_spec = (log_spec + 4.0) / 4.0\n return log_spec\n" }, { "path": "whisperx/conjunctions.py", "content": "# conjunctions.py\r\n\r\nfrom typing import Set\r\n\r\n\r\nconjunctions_by_language = {\r\n 'en': {'and', 'whether', 'or', 'as', 'but', 'so', 'for', 'nor', 'which', 'yet', 'although', 'since', 'unless', 'when', 'while', 'because', 'if', 'how', 'that', 'than', 'who', 'where', 'what', 'near', 'before', 'after', 'across', 'through', 'until', 'once', 'whereas', 'even', 'both', 'either', 'neither', 'though'},\r\n 'fr': {'et', 'ou', 'mais', 'parce', 'bien', 'pendant', 'quand', 'où', 'comme', 'si', 'que', 'avant', 'après', 'aussitôt', 'jusqu’à', 'à', 'malgré', 'donc', 'tant', 'puisque', 'ni', 'soit', 'bien', 'encore', 'dès', 'lorsque'},\r\n 'de': {'und', 'oder', 'aber', 'weil', 'obwohl', 'während', 'wenn', 'wo', 'wie', 'dass', 'bevor', 'nachdem', 'sobald', 'bis', 'außer', 'trotzdem', 'also', 'sowie', 'indem', 'weder', 'sowohl', 'zwar', 'jedoch'},\r\n 'es': {'y', 'o', 'pero', 'porque', 'aunque', 'sin', 'mientras', 'cuando', 'donde', 'como', 'si', 'que', 'antes', 'después', 'tan', 'hasta', 'a', 'a', 'por', 'ya', 'ni', 'sino'},\r\n 'it': {'e', 'o', 'ma', 'perché', 'anche', 'mentre', 'quando', 'dove', 'come', 'se', 'che', 'prima', 'dopo', 'appena', 'fino', 'a', 'nonostante', 'quindi', 'poiché', 'né', 'ossia', 'cioè'},\r\n 'ja': {'そして', 'または', 'しかし', 'なぜなら', 'もし', 'それとも', 'だから', 'それに', 'なのに', 'そのため', 'かつ', 'それゆえに', 'ならば', 'もしくは', 'ため'},\r\n 'zh': {'和', '或', '但是', '因为', '任何', '也', '虽然', '而且', '所以', '如果', '除非', '尽管', '既然', '即使', '只要', '直到', '然后', '因此', '不但', '而是', '不过'},\r\n 'nl': {'en', 'of', 'maar', 'omdat', 'hoewel', 'terwijl', 'wanneer', 'waar', 'zoals', 'als', 'dat', 'voordat', 'nadat', 'zodra', 'totdat', 'tenzij', 'ondanks', 'dus', 'zowel', 'noch', 'echter', 'toch'},\r\n 'uk': {'та', 'або', 'але', 'тому', 'хоча', 'поки', 'бо', 'коли', 'де', 'як', 'якщо', 'що', 'перш', 'після', 'доки', 'незважаючи', 'тому', 'ані'},\r\n 'pt': {'e', 'ou', 'mas', 'porque', 'embora', 'enquanto', 'quando', 'onde', 'como', 'se', 'que', 'antes', 'depois', 'assim', 'até', 'a', 'apesar', 'portanto', 'já', 'pois', 'nem', 'senão'},\r\n 'ar': {'و', 'أو', 'لكن', 'لأن', 'مع', 'بينما', 'عندما', 'حيث', 'كما', 'إذا', 'الذي', 'قبل', 'بعد', 'فور', 'حتى', 'إلا', 'رغم', 'لذلك', 'بما'},\r\n 'cs': {'a', 'nebo', 'ale', 'protože', 'ačkoli', 'zatímco', 'když', 'kde', 'jako', 'pokud', 'že', 'než', 'poté', 'jakmile', 'dokud', 'pokud ne', 'navzdory', 'tak', 'stejně', 'ani', 'tudíž'},\r\n 'ru': {'и', 'или', 'но', 'потому', 'хотя', 'пока', 'когда', 'где', 'как', 'если', 'что', 'перед', 'после', 'несмотря', 'таким', 'также', 'ни', 'зато'},\r\n 'pl': {'i', 'lub', 'ale', 'ponieważ', 'chociaż', 'podczas', 'kiedy', 'gdzie', 'jak', 'jeśli', 'że', 'zanim', 'po', 'jak tylko', 'dopóki', 'chyba', 'pomimo', 'więc', 'tak', 'ani', 'czyli'},\r\n 'hu': {'és', 'vagy', 'de', 'mert', 'habár', 'míg', 'amikor', 'ahol', 'ahogy', 'ha', 'hogy', 'mielőtt', 'miután', 'amint', 'amíg', 'hacsak', 'ellenére', 'tehát', 'úgy', 'sem', 'vagyis'},\r\n 'fi': {'ja', 'tai', 'mutta', 'koska', 'vaikka', 'kun', 'missä', 'kuten', 'jos', 'että', 'ennen', 'sen jälkeen', 'heti', 'kunnes', 'ellei', 'huolimatta', 'siis', 'sekä', 'eikä', 'vaan'},\r\n 'fa': {'و', 'یا', 'اما', 'چون', 'اگرچه', 'در حالی', 'وقتی', 'کجا', 'چگونه', 'اگر', 'که', 'قبل', 'پس', 'به محض', 'تا زمانی', 'مگر', 'با وجود', 'پس', 'همچنین', 'نه'},\r\n 'el': {'και', 'ή', 'αλλά', 'επειδή', 'αν', 'ενώ', 'όταν', 'όπου', 'όπως', 'αν', 'που', 'προτού', 'αφού', 'μόλις', 'μέχρι', 'εκτός', 'παρά', 'έτσι', 'όπως', 'ούτε', 'δηλαδή'},\r\n 'tr': {'ve', 'veya', 'ama', 'çünkü', 'her ne', 'iken', 'nerede', 'nasıl', 'eğer', 'ki', 'önce', 'sonra', 'hemen', 'kadar', 'rağmen', 'hem', 'ne', 'yani'},\r\n 'da': {'og', 'eller', 'men', 'fordi', 'selvom', 'mens', 'når', 'hvor', 'som', 'hvis', 'at', 'før', 'efter', 'indtil', 'medmindre', 'således', 'ligesom', 'hverken', 'altså'},\r\n 'he': {'ו', 'או', 'אבל', 'כי', 'אף', 'בזמן', 'כאשר', 'היכן', 'כיצד', 'אם', 'ש', 'לפני', 'אחרי', 'ברגע', 'עד', 'אלא', 'למרות', 'לכן', 'כמו', 'לא', 'אז'},\r\n 'vi': {'và', 'hoặc', 'nhưng', 'bởi', 'mặc', 'trong', 'khi', 'ở', 'như', 'nếu', 'rằng', 'trước', 'sau', 'ngay', 'cho', 'trừ', 'mặc', 'vì', 'giống', 'cũng', 'tức'},\r\n 'ko': {'그리고', '또는','그런데','그래도', '이나', '결국', '마지막으로', '마찬가지로', '반면에', '아니면', '거나', '또는', '그럼에도', '그렇기', '때문에', '덧붙이자면', '게다가', '그러나', '고', '그래서', '랑', '한다면', '하지만', '무엇', '왜냐하면', '비록', '동안', '언제', '어디서', '어떻게', '만약', '그', '전에', '후에', '즉시', '까지', '아니라면', '불구하고', '따라서', '같은', '도'},\r\n 'ur': {'اور', 'یا', 'مگر', 'کیونکہ', 'اگرچہ', 'جبکہ', 'جب', 'کہاں', 'کس طرح', 'اگر', 'کہ', 'سے پہلے', 'کے بعد', 'جیسے ہی', 'تک', 'اگر نہیں تو', 'کے باوجود', 'اس لئے', 'جیسے', 'نہ'},\r\n 'hi': {'और', 'या', 'पर', 'तो', 'न', 'फिर', 'हालांकि', 'चूंकि', 'अगर', 'कैसे', 'वह', 'से', 'जो', 'जहां', 'क्या', 'नजदीक', 'पहले', 'बाद', 'के', 'पार', 'माध्यम', 'तक', 'एक', 'जबकि', 'यहां', 'तक', 'दोनों', 'या', 'न', 'हालांकि'}\r\n\r\n}\r\n\r\ncommas_by_language = {\r\n 'ja': '、', \r\n 'zh': '，',\r\n 'fa': '،', \r\n 'ur': '،' \r\n}\r\n\r\ndef get_conjunctions(lang_code: str) -> Set[str]:\r\n return conjunctions_by_language.get(lang_code, set())\r\n\r\n\r\ndef get_comma(lang_code: str) -> str:\r\n return commas_by_language.get(lang_code, \",\")\r\n" }, { "path": "whisperx/diarize.py", "content": "import numpy as np\nimport pandas as pd\nfrom pyannote.audio import Pipeline\nfrom typing import Optional, Union, List, Tuple\nimport torch\n\nfrom whisperx.audio import load_audio, SAMPLE_RATE\nfrom whisperx.schema import TranscriptionResult, AlignedTranscriptionResult, ProgressCallback\nfrom whisperx.log_utils import get_logger\n\nlogger = get_logger(__name__)\n\n\nclass IntervalTree:\n \"\"\"\n Simple interval tree for fast overlap queries using sorted array + binary search.\n\n Uses O(n) space and provides O(log n) query time instead of O(n) linear scan.\n This achieves ~228x speedup for speaker assignment in long-form content.\n \"\"\"\n\n def __init__(self, intervals: List[Tuple[float, float, str]]):\n \"\"\"\n Initialize the interval tree with diarization segments.\n\n Args:\n intervals: List of (start, end, speaker) tuples\n \"\"\"\n if not intervals:\n self.starts = np.array([])\n self.ends = np.array([])\n self.speakers: List[str] = []\n return\n\n # Sort intervals by start time for binary search\n sorted_intervals = sorted(intervals, key=lambda x: x[0])\n self.starts = np.array([i[0] for i in sorted_intervals], dtype=np.float64)\n self.ends = np.array([i[1] for i in sorted_intervals], dtype=np.float64)\n self.speakers = [i[2] for i in sorted_intervals]\n\n def query(self, start: float, end: float) -> List[Tuple[str, float]]:\n \"\"\"\n Find all intervals that overlap with [start, end] and compute intersection.\n\n Args:\n start: Query interval start time\n end: Query interval end time\n\n Returns:\n List of (speaker, intersection_duration) tuples for overlapping segments\n \"\"\"\n if len(self.starts) == 0:\n return []\n\n # Binary search to find candidate intervals\n # Only intervals with start < end could overlap\n right_idx = np.searchsorted(self.starts, end, side='left')\n if right_idx == 0:\n return []\n\n # Check candidates for actual overlap\n candidates = slice(0, right_idx)\n overlaps = (self.starts[candidates] < end) & (self.ends[candidates] > start)\n\n results = []\n for idx in np.where(overlaps)[0]:\n intersection = min(self.ends[idx], end) - max(self.starts[idx], start)\n if intersection > 0:\n results.append((self.speakers[idx], intersection))\n return results\n\n def find_nearest(self, time: float) -> Optional[str]:\n \"\"\"\n Find the speaker of the nearest segment to a given time point.\n\n Args:\n time: Time point to find nearest segment for\n\n Returns:\n Speaker ID of nearest segment, or None if no segments exist\n \"\"\"\n if len(self.starts) == 0:\n return None\n\n # Calculate midpoints of all segments\n mids = (self.starts + self.ends) / 2\n nearest_idx = np.argmin(np.abs(mids - time))\n return self.speakers[nearest_idx]\n\n\nclass DiarizationPipeline:\n def __init__(\n self,\n model_name=None,\n token=None,\n device: Optional[Union[str, torch.device]] = \"cpu\",\n cache_dir=None,\n ):\n if isinstance(device, str):\n device = torch.device(device)\n model_config = model_name or \"pyannote/speaker-diarization-community-1\"\n logger.info(f\"Loading diarization model: {model_config}\")\n self.model = Pipeline.from_pretrained(model_config, token=token, cache_dir=cache_dir).to(device)\n\n def __call__(\n self,\n audio: Union[str, np.ndarray],\n num_speakers: Optional[int] = None,\n min_speakers: Optional[int] = None,\n max_speakers: Optional[int] = None,\n return_embeddings: bool = False,\n progress_callback: ProgressCallback = None,\n ) -> Union[tuple[pd.DataFrame, Optional[dict[str, list[float]]]], pd.DataFrame]:\n \"\"\"\n Perform speaker diarization on audio.\n\n Args:\n audio: Path to audio file or audio array\n num_speakers: Exact number of speakers (if known)\n min_speakers: Minimum number of speakers to detect\n max_speakers: Maximum number of speakers to detect\n return_embeddings: Whether to return speaker embeddings\n progress_callback: Optional callable receiving a float (0-100) with progress percentage\n\n Returns:\n If return_embeddings is True:\n Tuple of (diarization dataframe, speaker embeddings dictionary)\n Otherwise:\n Just the diarization dataframe\n \"\"\"\n if isinstance(audio, str):\n audio = load_audio(audio)\n audio_data = {\n 'waveform': torch.from_numpy(audio[None, :]),\n 'sample_rate': SAMPLE_RATE\n }\n\n hook = None\n if progress_callback is not None:\n # pyannote's diarization has two progress-trackable steps, each with\n # its own completed/total counter that resets between steps. Map each\n # step into a sub-range so progress is monotonic and meaningful.\n _STEP_RANGES = {\n \"segmentation\": (0.0, 50.0),\n \"embeddings\": (50.0, 99.0),\n }\n last_pct = [0.0]\n def hook(step_name, step_artifact, file=None, total=None, completed=None):\n if total is not None and completed is not None and total > 0:\n offset, end = _STEP_RANGES.get(step_name, (0.0, 99.0))\n pct = offset + min(completed / total, 1.0) * (end - offset)\n if pct > last_pct[0]:\n last_pct[0] = pct\n progress_callback(pct)\n\n output = self.model(\n audio_data,\n num_speakers=num_speakers,\n min_speakers=min_speakers,\n max_speakers=max_speakers,\n **({\"hook\": hook} if hook is not None else {}),\n )\n\n if progress_callback is not None:\n progress_callback(100.0)\n\n diarization = output.speaker_diarization\n embeddings = output.speaker_embeddings if return_embeddings else None\n\n diarize_df = pd.DataFrame(diarization.itertracks(yield_label=True), columns=['segment', 'label', 'speaker'])\n diarize_df['start'] = diarize_df['segment'].apply(lambda x: x.start)\n diarize_df['end'] = diarize_df['segment'].apply(lambda x: x.end)\n\n if return_embeddings and embeddings is not None:\n speaker_embeddings = {speaker: embeddings[s].tolist() for s, speaker in enumerate(diarization.labels())}\n return diarize_df, speaker_embeddings\n \n # For backwards compatibility\n if return_embeddings:\n return diarize_df, None\n else:\n return diarize_df\n\n\ndef assign_word_speakers(\n diarize_df: pd.DataFrame,\n transcript_result: Union[AlignedTranscriptionResult, TranscriptionResult],\n speaker_embeddings: Optional[dict[str, list[float]]] = None,\n fill_nearest: bool = False,\n) -> Union[AlignedTranscriptionResult, TranscriptionResult]:\n \"\"\"\n Assign speakers to words and segments in the transcript.\n\n Uses an interval tree for O(log n) overlap queries instead of O(n) linear scan,\n achieving ~228x speedup for long-form content (3+ hour podcasts).\n\n Args:\n diarize_df: Diarization dataframe from DiarizationPipeline\n transcript_result: Transcription result to augment with speaker labels\n speaker_embeddings: Optional dictionary mapping speaker IDs to embedding vectors\n fill_nearest: If True, assign speakers even when there's no direct time overlap\n\n Returns:\n Updated transcript_result with speaker assignments and optionally embeddings\n \"\"\"\n transcript_segments = transcript_result.get(\"segments\", [])\n if not transcript_segments or diarize_df is None or len(diarize_df) == 0:\n return transcript_result\n\n # Build interval tree from diarization segments for O(log n) queries\n intervals = [\n (row['start'], row['end'], row['speaker'])\n for _, row in diarize_df.iterrows()\n ]\n tree = IntervalTree(intervals)\n\n for seg in transcript_segments:\n seg_start = seg.get('start', 0.0)\n seg_end = seg.get('end', 0.0)\n\n # Query overlapping segments using interval tree\n overlaps = tree.query(seg_start, seg_end)\n\n if overlaps:\n # Sum intersection durations per speaker and pick the dominant one\n speaker_intersections: dict[str, float] = {}\n for speaker, intersection in overlaps:\n speaker_intersections[speaker] = speaker_intersections.get(speaker, 0.0) + intersection\n seg['speaker'] = max(speaker_intersections.items(), key=lambda x: x[1])[0]\n elif fill_nearest:\n # Find nearest segment if no overlap\n seg_mid = (seg_start + seg_end) / 2\n nearest_speaker = tree.find_nearest(seg_mid)\n if nearest_speaker:\n seg['speaker'] = nearest_speaker\n\n # Assign speaker to words\n if 'words' in seg:\n for word in seg['words']:\n if 'start' not in word:\n continue\n\n word_start = word['start']\n word_end = word.get('end', word_start)\n\n word_overlaps = tree.query(word_start, word_end)\n\n if word_overlaps:\n speaker_intersections = {}\n for speaker, intersection in word_overlaps:\n speaker_intersections[speaker] = speaker_intersections.get(speaker, 0.0) + intersection\n word['speaker'] = max(speaker_intersections.items(), key=lambda x: x[1])[0]\n elif fill_nearest:\n word_mid = (word_start + word_end) / 2\n nearest_speaker = tree.find_nearest(word_mid)\n if nearest_speaker:\n word['speaker'] = nearest_speaker\n\n # Add speaker embeddings to the result if provided\n if speaker_embeddings is not None:\n transcript_result[\"speaker_embeddings\"] = speaker_embeddings\n\n return transcript_result\n\n\nclass Segment:\n def __init__(self, start:int, end:int, speaker:Optional[str]=None):\n self.start = start\n self.end = end\n self.speaker = speaker\n" }, { "path": "whisperx/log_utils.py", "content": "import logging\nimport sys\nfrom typing import Optional\n\n_LOG_FORMAT = \"%(asctime)s - %(name)s - %(levelname)s - %(message)s\"\n_DATE_FORMAT = \"%Y-%m-%d %H:%M:%S\"\n\n\ndef setup_logging(\n level: str = \"info\",\n log_file: Optional[str] = None,\n) -> None:\n \"\"\"\n Configure logging for WhisperX.\n\n Args:\n level: Logging level (debug, info, warning, error, critical). Default: info\n log_file: Optional path to log file. If None, logs only to console.\n \"\"\"\n logger = logging.getLogger(\"whisperx\")\n\n logger.handlers.clear()\n\n try:\n log_level = getattr(logging, level.upper())\n except AttributeError:\n log_level = logging.WARNING\n logger.setLevel(log_level)\n\n formatter = logging.Formatter(_LOG_FORMAT, datefmt=_DATE_FORMAT)\n\n console_handler = logging.StreamHandler(sys.stdout)\n console_handler.setLevel(log_level)\n console_handler.setFormatter(formatter)\n\n logger.addHandler(console_handler)\n\n if log_file:\n try:\n file_handler = logging.FileHandler(log_file)\n file_handler.setLevel(log_level)\n file_handler.setFormatter(formatter)\n logger.addHandler(file_handler)\n except (OSError) as e:\n logger.warning(f\"Failed to create log file '{log_file}': {e}\")\n logger.warning(\"Continuing with console logging only\")\n\n # Don't propagate to root logger to avoid duplicate messages\n logger.propagate = False\n\n\ndef get_logger(name: str) -> logging.Logger:\n \"\"\"\n Get a logger instance for the given module.\n\n Args:\n name: Logger name (typically __name__ from calling module)\n\n Returns:\n Logger instance configured with WhisperX settings\n \"\"\"\n whisperx_logger = logging.getLogger(\"whisperx\")\n if not whisperx_logger.handlers:\n setup_logging()\n\n logger_name = \"whisperx\" if name == \"__main__\" else name\n return logging.getLogger(logger_name)\n" }, { "path": "whisperx/schema.py", "content": "from typing import Callable, TypedDict, Optional, List, Tuple\n\nProgressCallback = Optional[Callable[[float], None]]\n\ntry:\n from typing import NotRequired\nexcept ImportError:\n from typing_extensions import NotRequired\n\n\nclass SingleWordSegment(TypedDict):\n \"\"\"\n A single word of a speech.\n \"\"\"\n word: str\n start: float\n end: float\n score: float\n\nclass SingleCharSegment(TypedDict):\n \"\"\"\n A single char of a speech.\n \"\"\"\n char: str\n start: float\n end: float\n score: float\n\n\nclass SingleSegment(TypedDict):\n \"\"\"\n A single segment (up to multiple sentences) of a speech.\n \"\"\"\n\n start: float\n end: float\n text: str\n avg_logprob: NotRequired[float]\n\n\nclass SegmentData(TypedDict):\n \"\"\"\n Temporary processing data used during alignment.\n Contains cleaned and preprocessed data for each segment.\n \"\"\"\n clean_char: List[str] # Cleaned characters that exist in model dictionary\n clean_cdx: List[int] # Original indices of cleaned characters\n clean_wdx: List[int] # Indices of words containing valid characters\n sentence_spans: List[Tuple[int, int]] # Start and end indices of sentences\n\n\nclass SingleAlignedSegment(TypedDict):\n \"\"\"\n A single segment (up to multiple sentences) of a speech with word alignment.\n \"\"\"\n\n start: float\n end: float\n text: str\n avg_logprob: NotRequired[float]\n words: List[SingleWordSegment]\n chars: Optional[List[SingleCharSegment]]\n\n\nclass TranscriptionResult(TypedDict):\n \"\"\"\n A list of segments and word segments of a speech.\n \"\"\"\n segments: List[SingleSegment]\n language: str\n\n\nclass AlignedTranscriptionResult(TypedDict):\n \"\"\"\n A list of segments and word segments of a speech.\n \"\"\"\n segments: List[SingleAlignedSegment]\n word_segments: List[SingleWordSegment]\n" }, { "path": "whisperx/transcribe.py", "content": "import argparse\nimport gc\nimport os\nimport warnings\n\nimport numpy as np\nimport torch\n\nfrom whisperx.alignment import align, load_align_model\nfrom whisperx.asr import load_model\nfrom whisperx.audio import load_audio\nfrom whisperx.diarize import DiarizationPipeline, assign_word_speakers\nfrom whisperx.schema import AlignedTranscriptionResult, TranscriptionResult\nfrom whisperx.utils import LANGUAGES, TO_LANGUAGE_CODE, get_writer\nfrom whisperx.log_utils import get_logger\n\nlogger = get_logger(__name__)\n\n\ndef transcribe_task(args: dict, parser: argparse.ArgumentParser):\n \"\"\"Transcription task to be called from CLI.\n\n Args:\n args: Dictionary of command-line arguments.\n parser: argparse.ArgumentParser object.\n \"\"\"\n # fmt: off\n\n model_name: str = args.pop(\"model\")\n batch_size: int = args.pop(\"batch_size\")\n model_dir: str = args.pop(\"model_dir\")\n model_cache_only: bool = args.pop(\"model_cache_only\")\n output_dir: str = args.pop(\"output_dir\")\n output_format: str = args.pop(\"output_format\")\n device: str = args.pop(\"device\")\n device_index: int = args.pop(\"device_index\")\n compute_type: str = args.pop(\"compute_type\")\n verbose: bool = args.pop(\"verbose\")\n\n # model_flush: bool = args.pop(\"model_flush\")\n os.makedirs(output_dir, exist_ok=True)\n\n align_model: str = args.pop(\"align_model\")\n interpolate_method: str = args.pop(\"interpolate_method\")\n no_align: bool = args.pop(\"no_align\")\n task: str = args.pop(\"task\")\n if task == \"translate\":\n # translation cannot be aligned\n no_align = True\n\n return_char_alignments: bool = args.pop(\"return_char_alignments\")\n\n hf_token: str = args.pop(\"hf_token\")\n vad_method: str = args.pop(\"vad_method\")\n vad_onset: float = args.pop(\"vad_onset\")\n vad_offset: float = args.pop(\"vad_offset\")\n\n chunk_size: int = args.pop(\"chunk_size\")\n\n diarize: bool = args.pop(\"diarize\")\n min_speakers: int = args.pop(\"min_speakers\")\n max_speakers: int = args.pop(\"max_speakers\")\n diarize_model_name: str = args.pop(\"diarize_model\")\n print_progress: bool = args.pop(\"print_progress\")\n return_speaker_embeddings: bool = args.pop(\"speaker_embeddings\")\n\n if return_speaker_embeddings and not diarize:\n warnings.warn(\"--speaker_embeddings has no effect without --diarize\")\n\n if args[\"language\"] is not None:\n args[\"language\"] = args[\"language\"].lower()\n if args[\"language\"] not in LANGUAGES:\n if args[\"language\"] in TO_LANGUAGE_CODE:\n args[\"language\"] = TO_LANGUAGE_CODE[args[\"language\"]]\n else:\n raise ValueError(f\"Unsupported language: {args['language']}\")\n\n if model_name.endswith(\".en\") and args[\"language\"] != \"en\":\n if args[\"language\"] is not None:\n warnings.warn(\n f\"{model_name} is an English-only model but received '{args['language']}'; using English instead.\"\n )\n args[\"language\"] = \"en\"\n align_language = (\n args[\"language\"] if args[\"language\"] is not None else \"en\"\n ) # default to loading english if not specified\n\n temperature = args.pop(\"temperature\")\n if (increment := args.pop(\"temperature_increment_on_fallback\")) is not None:\n temperature = tuple(np.arange(temperature, 1.0 + 1e-6, increment))\n else:\n temperature = [temperature]\n\n faster_whisper_threads = 4\n if (threads := args.pop(\"threads\")) > 0:\n torch.set_num_threads(threads)\n faster_whisper_threads = threads\n\n asr_options = {\n \"beam_size\": args.pop(\"beam_size\"),\n \"patience\": args.pop(\"patience\"),\n \"length_penalty\": args.pop(\"length_penalty\"),\n \"temperatures\": temperature,\n \"compression_ratio_threshold\": args.pop(\"compression_ratio_threshold\"),\n \"log_prob_threshold\": args.pop(\"logprob_threshold\"),\n \"no_speech_threshold\": args.pop(\"no_speech_threshold\"),\n \"condition_on_previous_text\": False,\n \"initial_prompt\": args.pop(\"initial_prompt\"),\n \"hotwords\": args.pop(\"hotwords\"),\n \"suppress_tokens\": [int(x) for x in args.pop(\"suppress_tokens\").split(\",\")],\n \"suppress_numerals\": args.pop(\"suppress_numerals\"),\n }\n\n writer = get_writer(output_format, output_dir)\n word_options = [\"highlight_words\", \"max_line_count\", \"max_line_width\"]\n if no_align:\n for option in word_options:\n if args[option]:\n parser.error(f\"--{option} not possible with --no_align\")\n if args[\"max_line_count\"] and not args[\"max_line_width\"]:\n warnings.warn(\"--max_line_count has no effect without --max_line_width\")\n writer_args = {arg: args.pop(arg) for arg in word_options}\n\n # Part 1: VAD & ASR Loop\n results = []\n # model = load_model(model_name, device=device, download_root=model_dir)\n model = load_model(\n model_name,\n device=device,\n device_index=device_index,\n download_root=model_dir,\n compute_type=compute_type,\n language=args[\"language\"],\n asr_options=asr_options,\n vad_method=vad_method,\n vad_options={\n \"chunk_size\": chunk_size,\n \"vad_onset\": vad_onset,\n \"vad_offset\": vad_offset,\n },\n task=task,\n local_files_only=model_cache_only,\n threads=faster_whisper_threads,\n use_auth_token=hf_token,\n )\n\n for audio_path in args.pop(\"audio\"):\n audio = load_audio(audio_path)\n # >> VAD & ASR\n logger.info(\"Performing transcription...\")\n result: TranscriptionResult = model.transcribe(\n audio,\n batch_size=batch_size,\n chunk_size=chunk_size,\n print_progress=print_progress,\n verbose=verbose,\n )\n results.append((result, audio_path))\n\n # Unload Whisper and VAD\n del model\n gc.collect()\n torch.cuda.empty_cache()\n\n # Part 2: Align Loop\n if not no_align:\n tmp_results = results\n results = []\n align_model, align_metadata = load_align_model(\n align_language, device, model_name=align_model, model_dir=model_dir, model_cache_only=model_cache_only\n )\n for result, audio_path in tmp_results:\n # >> Align\n if len(tmp_results) > 1:\n input_audio = audio_path\n else:\n # lazily load audio from part 1\n input_audio = audio\n\n if align_model is not None and len(result[\"segments\"]) > 0:\n if result.get(\"language\", \"en\") != align_metadata[\"language\"]:\n # load new language\n logger.info(\n f\"New language found ({result['language']})! Previous was ({align_metadata['language']}), loading new alignment model for new language...\"\n )\n align_model, align_metadata = load_align_model(\n result[\"language\"], device, model_dir=model_dir, model_cache_only=model_cache_only\n )\n logger.info(\"Performing alignment...\")\n result: AlignedTranscriptionResult = align(\n result[\"segments\"],\n align_model,\n align_metadata,\n input_audio,\n device,\n interpolate_method=interpolate_method,\n return_char_alignments=return_char_alignments,\n print_progress=print_progress,\n )\n\n results.append((result, audio_path))\n\n # Unload align model\n del align_model\n gc.collect()\n torch.cuda.empty_cache()\n\n # >> Diarize\n if diarize:\n if hf_token is None:\n logger.warning(\n \"No --hf_token provided, needs to be saved in environment variable, otherwise will throw error loading diarization model\"\n )\n tmp_results = results\n logger.info(\"Performing diarization...\")\n logger.info(f\"Using model: {diarize_model_name}\")\n results = []\n diarize_model = DiarizationPipeline(model_name=diarize_model_name, token=hf_token, device=device, cache_dir=model_dir)\n for result, input_audio_path in tmp_results:\n diarize_result = diarize_model(\n input_audio_path, \n min_speakers=min_speakers, \n max_speakers=max_speakers, \n return_embeddings=return_speaker_embeddings\n )\n\n if return_speaker_embeddings:\n diarize_segments, speaker_embeddings = diarize_result\n else:\n diarize_segments = diarize_result\n speaker_embeddings = None\n\n result = assign_word_speakers(diarize_segments, result, speaker_embeddings)\n results.append((result, input_audio_path))\n # >> Write\n for result, audio_path in results:\n result[\"language\"] = align_language\n writer(result, audio_path, writer_args)\n" }, { "path": "whisperx/utils.py", "content": "import json\nimport os\nimport re\nimport sys\nimport zlib\nfrom typing import Callable, Optional, TextIO\n\nLANGUAGES = {\n \"en\": \"english\",\n \"zh\": \"chinese\",\n \"de\": \"german\",\n \"es\": \"spanish\",\n \"ru\": \"russian\",\n \"ko\": \"korean\",\n \"fr\": \"french\",\n \"ja\": \"japanese\",\n \"pt\": \"portuguese\",\n \"tr\": \"turkish\",\n \"pl\": \"polish\",\n \"ca\": \"catalan\",\n \"nl\": \"dutch\",\n \"ar\": \"arabic\",\n \"sv\": \"swedish\",\n \"it\": \"italian\",\n \"id\": \"indonesian\",\n \"hi\": \"hindi\",\n \"fi\": \"finnish\",\n \"vi\": \"vietnamese\",\n \"he\": \"hebrew\",\n \"uk\": \"ukrainian\",\n \"el\": \"greek\",\n \"ms\": \"malay\",\n \"cs\": \"czech\",\n \"ro\": \"romanian\",\n \"da\": \"danish\",\n \"hu\": \"hungarian\",\n \"ta\": \"tamil\",\n \"no\": \"norwegian\",\n \"th\": \"thai\",\n \"ur\": \"urdu\",\n \"hr\": \"croatian\",\n \"bg\": \"bulgarian\",\n \"lt\": \"lithuanian\",\n \"la\": \"latin\",\n \"mi\": \"maori\",\n \"ml\": \"malayalam\",\n \"cy\": \"welsh\",\n \"sk\": \"slovak\",\n \"te\": \"telugu\",\n \"fa\": \"persian\",\n \"lv\": \"latvian\",\n \"bn\": \"bengali\",\n \"sr\": \"serbian\",\n \"az\": \"azerbaijani\",\n \"sl\": \"slovenian\",\n \"kn\": \"kannada\",\n \"et\": \"estonian\",\n \"mk\": \"macedonian\",\n \"br\": \"breton\",\n \"eu\": \"basque\",\n \"is\": \"icelandic\",\n \"hy\": \"armenian\",\n \"ne\": \"nepali\",\n \"mn\": \"mongolian\",\n \"bs\": \"bosnian\",\n \"kk\": \"kazakh\",\n \"sq\": \"albanian\",\n \"sw\": \"swahili\",\n \"gl\": \"galician\",\n \"mr\": \"marathi\",\n \"pa\": \"punjabi\",\n \"si\": \"sinhala\",\n \"km\": \"khmer\",\n \"sn\": \"shona\",\n \"yo\": \"yoruba\",\n \"so\": \"somali\",\n \"af\": \"afrikaans\",\n \"oc\": \"occitan\",\n \"ka\": \"georgian\",\n \"be\": \"belarusian\",\n \"tg\": \"tajik\",\n \"sd\": \"sindhi\",\n \"gu\": \"gujarati\",\n \"am\": \"amharic\",\n \"yi\": \"yiddish\",\n \"lo\": \"lao\",\n \"uz\": \"uzbek\",\n \"fo\": \"faroese\",\n \"ht\": \"haitian creole\",\n \"ps\": \"pashto\",\n \"tk\": \"turkmen\",\n \"nn\": \"nynorsk\",\n \"mt\": \"maltese\",\n \"sa\": \"sanskrit\",\n \"lb\": \"luxembourgish\",\n \"my\": \"myanmar\",\n \"bo\": \"tibetan\",\n \"tl\": \"tagalog\",\n \"mg\": \"malagasy\",\n \"as\": \"assamese\",\n \"tt\": \"tatar\",\n \"haw\": \"hawaiian\",\n \"ln\": \"lingala\",\n \"ha\": \"hausa\",\n \"ba\": \"bashkir\",\n \"jw\": \"javanese\",\n \"su\": \"sundanese\",\n \"yue\": \"cantonese\",\n}\n\n# language code lookup by name, with a few language aliases\nTO_LANGUAGE_CODE = {\n **{language: code for code, language in LANGUAGES.items()},\n \"burmese\": \"my\",\n \"valencian\": \"ca\",\n \"flemish\": \"nl\",\n \"haitian\": \"ht\",\n \"letzeburgesch\": \"lb\",\n \"pushto\": \"ps\",\n \"panjabi\": \"pa\",\n \"moldavian\": \"ro\",\n \"moldovan\": \"ro\",\n \"sinhalese\": \"si\",\n \"castilian\": \"es\",\n}\n\nLANGUAGES_WITHOUT_SPACES = [\"ja\", \"zh\"]\n\n# Mapping of language codes to NLTK Punkt tokenizer model names\nPUNKT_LANGUAGES = {\n 'cs': 'czech',\n 'da': 'danish', \n 'de': 'german',\n 'el': 'greek',\n 'en': 'english',\n 'es': 'spanish',\n 'et': 'estonian',\n 'fi': 'finnish',\n 'fr': 'french',\n 'it': 'italian',\n 'nl': 'dutch',\n 'no': 'norwegian',\n 'pl': 'polish',\n 'pt': 'portuguese',\n 'sl': 'slovene',\n 'sv': 'swedish',\n 'tr': 'turkish',\n \"ml\": \"malayalam\",\n \"ru\": \"russian\",\n}\n\nsystem_encoding = sys.getdefaultencoding()\n\nif system_encoding != \"utf-8\":\n\n def make_safe(string):\n # replaces any character not representable using the system default encoding with an '?',\n # avoiding UnicodeEncodeError (https://github.com/openai/whisper/discussions/729).\n return string.encode(system_encoding, errors=\"replace\").decode(system_encoding)\n\nelse:\n\n def make_safe(string):\n # utf-8 can encode any Unicode code point, so no need to do the round-trip encoding\n return string\n\n\ndef exact_div(x, y):\n assert x % y == 0\n return x // y\n\n\ndef str2bool(string):\n str2val = {\"True\": True, \"False\": False}\n if string in str2val:\n return str2val[string]\n else:\n raise ValueError(f\"Expected one of {set(str2val.keys())}, got {string}\")\n\n\ndef optional_int(string):\n return None if string == \"None\" else int(string)\n\n\ndef optional_float(string):\n return None if string == \"None\" else float(string)\n\n\ndef compression_ratio(text) -> float:\n text_bytes = text.encode(\"utf-8\")\n return len(text_bytes) / len(zlib.compress(text_bytes))\n\n\ndef format_timestamp(\n seconds: float, always_include_hours: bool = False, decimal_marker: str = \".\"\n):\n assert seconds >= 0, \"non-negative timestamp expected\"\n milliseconds = round(seconds * 1000.0)\n\n hours = milliseconds // 3_600_000\n milliseconds -= hours * 3_600_000\n\n minutes = milliseconds // 60_000\n milliseconds -= minutes * 60_000\n\n seconds = milliseconds // 1_000\n milliseconds -= seconds * 1_000\n\n hours_marker = f\"{hours:02d}:\" if always_include_hours or hours > 0 else \"\"\n return (\n f\"{hours_marker}{minutes:02d}:{seconds:02d}{decimal_marker}{milliseconds:03d}\"\n )\n\n\nclass ResultWriter:\n extension: str\n\n def __init__(self, output_dir: str):\n self.output_dir = output_dir\n\n def __call__(self, result: dict, audio_path: str, options: dict):\n audio_basename = os.path.basename(audio_path)\n audio_basename = os.path.splitext(audio_basename)[0]\n output_path = os.path.join(\n self.output_dir, audio_basename + \".\" + self.extension\n )\n\n with open(output_path, \"w\", encoding=\"utf-8\") as f:\n self.write_result(result, file=f, options=options)\n\n def write_result(self, result: dict, file: TextIO, options: dict):\n raise NotImplementedError\n\n\nclass WriteTXT(ResultWriter):\n extension: str = \"txt\"\n\n def write_result(self, result: dict, file: TextIO, options: dict):\n for segment in result[\"segments\"]:\n speaker = segment.get(\"speaker\")\n text = segment[\"text\"].strip()\n if speaker is not None:\n print(f\"[{speaker}]: {text}\", file=file, flush=True)\n else:\n print(text, file=file, flush=True)\n\n\nclass SubtitlesWriter(ResultWriter):\n always_include_hours: bool\n decimal_marker: str\n\n def iterate_result(self, result: dict, options: dict):\n raw_max_line_width: Optional[int] = options[\"max_line_width\"]\n max_line_count: Optional[int] = options[\"max_line_count\"]\n highlight_words: bool = options[\"highlight_words\"]\n max_line_width = 1000 if raw_max_line_width is None else raw_max_line_width\n preserve_segments = max_line_count is None or raw_max_line_width is None\n\n if len(result[\"segments\"]) == 0:\n return\n\n def iterate_subtitles():\n line_len = 0\n line_count = 1\n # the next subtitle to yield (a list of word timings with whitespace)\n subtitle: list[dict] = []\n times: list[tuple] = []\n last = result[\"segments\"][0][\"start\"]\n for segment in result[\"segments\"]:\n for i, original_timing in enumerate(segment[\"words\"]):\n timing = original_timing.copy()\n long_pause = not preserve_segments\n if \"start\" in timing:\n long_pause = long_pause and timing[\"start\"] - last > 3.0\n else:\n long_pause = False\n has_room = line_len + len(timing[\"word\"]) <= max_line_width\n seg_break = i == 0 and len(subtitle) > 0 and preserve_segments\n if line_len > 0 and has_room and not long_pause and not seg_break:\n # line continuation\n line_len += len(timing[\"word\"])\n else:\n # new line\n timing[\"word\"] = timing[\"word\"].strip()\n if (\n len(subtitle) > 0\n and max_line_count is not None\n and (long_pause or line_count >= max_line_count)\n or seg_break\n ):\n # subtitle break\n yield subtitle, times\n subtitle = []\n times = []\n line_count = 1\n elif line_len > 0:\n # line break\n line_count += 1\n timing[\"word\"] = \"\\n\" + timing[\"word\"]\n line_len = len(timing[\"word\"].strip())\n subtitle.append(timing)\n times.append((segment[\"start\"], segment[\"end\"], segment.get(\"speaker\")))\n if \"start\" in timing:\n last = timing[\"start\"]\n if len(subtitle) > 0:\n yield subtitle, times\n\n if \"words\" in result[\"segments\"][0]:\n for subtitle, times in iterate_subtitles():\n speaker = times[0][2]\n\n # Derive cue times from word-level timestamps when available,\n # falling back to segment-level times for fully unalignable subtitles.\n word_starts = [w[\"start\"] for w in subtitle if \"start\" in w]\n word_ends = [w[\"end\"] for w in subtitle if \"end\" in w]\n if word_starts and word_ends:\n subtitle_start = self.format_timestamp(min(word_starts))\n subtitle_end = self.format_timestamp(max(word_ends))\n else:\n subtitle_start = self.format_timestamp(times[0][0])\n subtitle_end = self.format_timestamp(times[0][1])\n if result[\"language\"] in LANGUAGES_WITHOUT_SPACES:\n subtitle_text = \"\".join([word[\"word\"] for word in subtitle])\n else:\n subtitle_text = \" \".join([word[\"word\"] for word in subtitle])\n has_timing = any([\"start\" in word for word in subtitle])\n\n # add [$SPEAKER_ID]: to each subtitle if speaker is available\n prefix = \"\"\n if speaker is not None:\n prefix = f\"[{speaker}]: \"\n\n if highlight_words and has_timing:\n last = subtitle_start\n all_words = [timing[\"word\"] for timing in subtitle]\n for i, this_word in enumerate(subtitle):\n if \"start\" in this_word:\n start = self.format_timestamp(this_word[\"start\"])\n end = self.format_timestamp(this_word[\"end\"])\n if last != start:\n yield last, start, prefix + subtitle_text\n\n yield start, end, prefix + \" \".join(\n [\n re.sub(r\"^(\\s*)(.*)$\", r\"\\1\\2\", word)\n if j == i\n else word\n for j, word in enumerate(all_words)\n ]\n )\n last = end\n else:\n yield subtitle_start, subtitle_end, prefix + subtitle_text\n else:\n for segment in result[\"segments\"]:\n segment_start = self.format_timestamp(segment[\"start\"])\n segment_end = self.format_timestamp(segment[\"end\"])\n segment_text = segment[\"text\"].strip().replace(\"-->\", \"->\")\n if \"speaker\" in segment:\n segment_text = f\"[{segment['speaker']}]: {segment_text}\"\n yield segment_start, segment_end, segment_text\n\n def format_timestamp(self, seconds: float):\n return format_timestamp(\n seconds=seconds,\n always_include_hours=self.always_include_hours,\n decimal_marker=self.decimal_marker,\n )\n\n\nclass WriteVTT(SubtitlesWriter):\n extension: str = \"vtt\"\n always_include_hours: bool = False\n decimal_marker: str = \".\"\n\n def write_result(self, result: dict, file: TextIO, options: dict):\n print(\"WEBVTT\\n\", file=file)\n for start, end, text in self.iterate_result(result, options):\n print(f\"{start} --> {end}\\n{text}\\n\", file=file, flush=True)\n\n\nclass WriteSRT(SubtitlesWriter):\n extension: str = \"srt\"\n always_include_hours: bool = True\n decimal_marker: str = \",\"\n\n def write_result(self, result: dict, file: TextIO, options: dict):\n for i, (start, end, text) in enumerate(\n self.iterate_result(result, options), start=1\n ):\n print(f\"{i}\\n{start} --> {end}\\n{text}\\n\", file=file, flush=True)\n\n\nclass WriteTSV(ResultWriter):\n \"\"\"\n Write a transcript to a file in TSV (tab-separated values) format containing lines like:\n \\t\\t\n\n Using integer milliseconds as start and end times means there's no chance of interference from\n an environment setting a language encoding that causes the decimal in a floating point number\n to appear as a comma; also is faster and more efficient to parse & store, e.g., in C++.\n \"\"\"\n\n extension: str = \"tsv\"\n\n def write_result(self, result: dict, file: TextIO, options: dict):\n print(\"start\", \"end\", \"text\", sep=\"\\t\", file=file)\n for segment in result[\"segments\"]:\n print(round(1000 * segment[\"start\"]), file=file, end=\"\\t\")\n print(round(1000 * segment[\"end\"]), file=file, end=\"\\t\")\n print(segment[\"text\"].strip().replace(\"\\t\", \" \"), file=file, flush=True)\n\nclass WriteAudacity(ResultWriter):\n \"\"\"\n Write a transcript to a text file that audacity can import as labels.\n The extension used is \"aud\" to distinguish it from the txt file produced by WriteTXT.\n Yet this is not an audacity project but only a label file!\n \n Please note : Audacity uses seconds in timestamps not ms! \n Also there is no header expected.\n\n If speaker is provided it is prepended to the text between double square brackets [[]].\n \"\"\"\n\n extension: str = \"aud\" \n\n def write_result(self, result: dict, file: TextIO, options: dict):\n ARROW = \"\t\"\n for segment in result[\"segments\"]:\n print(segment[\"start\"], file=file, end=ARROW)\n print(segment[\"end\"], file=file, end=ARROW)\n print( ( (\"[[\" + segment[\"speaker\"] + \"]]\") if \"speaker\" in segment else \"\") + segment[\"text\"].strip().replace(\"\\t\", \" \"), file=file, flush=True)\n\n \n\nclass WriteJSON(ResultWriter):\n extension: str = \"json\"\n\n def write_result(self, result: dict, file: TextIO, options: dict):\n json.dump(result, file, ensure_ascii=False)\n\n\ndef get_writer(\n output_format: str, output_dir: str\n) -> Callable[[dict, str, dict], None]:\n writers = {\n \"txt\": WriteTXT,\n \"vtt\": WriteVTT,\n \"srt\": WriteSRT,\n \"tsv\": WriteTSV,\n \"json\": WriteJSON,\n }\n optional_writers = {\n \"aud\": WriteAudacity,\n }\n\n if output_format == \"all\":\n all_writers = [writer(output_dir) for writer in writers.values()]\n\n def write_all(result: dict, file: str, options: dict):\n for writer in all_writers:\n writer(result, file, options)\n\n return write_all\n\n if output_format in optional_writers:\n return optional_writers[output_format](output_dir)\n return writers[output_format](output_dir)\n\ndef interpolate_nans(x, method='nearest'):\n if x.notnull().sum() > 1:\n return x.interpolate(method=method).ffill().bfill()\n else:\n return x.ffill().bfill()\n" }, { "path": "whisperx/vads/__init__.py", "content": "from whisperx.vads.pyannote import Pyannote as Pyannote\nfrom whisperx.vads.silero import Silero as Silero\nfrom whisperx.vads.vad import Vad as Vad\n" }, { "path": "whisperx/vads/pyannote.py", "content": "import os\nfrom typing import Callable, Text, Union\nfrom typing import Optional\n\nimport numpy as np\nimport torch\nfrom pyannote.audio import Model\nfrom pyannote.audio.core.io import AudioFile\nfrom pyannote.audio.pipelines import VoiceActivityDetection\nfrom pyannote.audio.pipelines.utils import PipelineModel\nfrom pyannote.core import Annotation, SlidingWindowFeature\nfrom pyannote.core import Segment\n\nfrom whisperx.diarize import Segment as SegmentX\nfrom whisperx.vads.vad import Vad\nfrom whisperx.log_utils import get_logger\n\nlogger = get_logger(__name__)\n\n\ndef load_vad_model(device, vad_onset=0.500, vad_offset=0.363, token=None, model_fp=None):\n model_dir = torch.hub._get_torch_home()\n\n main_dir = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))\n\n os.makedirs(model_dir, exist_ok = True)\n if model_fp is None:\n # Dynamically resolve the path to the model file\n model_fp = os.path.join(main_dir, \"assets\", \"pytorch_model.bin\")\n model_fp = os.path.abspath(model_fp) # Ensure the path is absolute\n else:\n model_fp = os.path.abspath(model_fp) # Ensure any provided path is absolute\n\n # Check if the resolved model file exists\n if not os.path.exists(model_fp):\n raise FileNotFoundError(f\"Model file not found at {model_fp}\")\n\n if os.path.exists(model_fp) and not os.path.isfile(model_fp):\n raise RuntimeError(f\"{model_fp} exists and is not a regular file\")\n\n vad_model = Model.from_pretrained(model_fp, token=token)\n hyperparameters = {\"onset\": vad_onset,\n \"offset\": vad_offset,\n \"min_duration_on\": 0.1,\n \"min_duration_off\": 0.1}\n vad_pipeline = VoiceActivitySegmentation(segmentation=vad_model, device=torch.device(device))\n vad_pipeline.instantiate(hyperparameters)\n\n return vad_pipeline\n\nclass Binarize:\n \"\"\"Binarize detection scores using hysteresis thresholding, with min-cut operation\n to ensure not segments are longer than max_duration.\n\n Parameters\n ----------\n onset : float, optional\n Onset threshold. Defaults to 0.5.\n offset : float, optional\n Offset threshold. Defaults to `onset`.\n min_duration_on : float, optional\n Remove active regions shorter than that many seconds. Defaults to 0s.\n min_duration_off : float, optional\n Fill inactive regions shorter than that many seconds. Defaults to 0s.\n pad_onset : float, optional\n Extend active regions by moving their start time by that many seconds.\n Defaults to 0s.\n pad_offset : float, optional\n Extend active regions by moving their end time by that many seconds.\n Defaults to 0s.\n max_duration: float\n The maximum length of an active segment, divides segment at timestamp with lowest score.\n Reference\n ---------\n Gregory Gelly and Jean-Luc Gauvain. \"Minimum Word Error Training of\n RNN-based Voice Activity Detection\", InterSpeech 2015.\n\n Modified by Max Bain to include WhisperX's min-cut operation\n https://arxiv.org/abs/2303.00747\n\n Pyannote-audio\n \"\"\"\n\n def __init__(\n self,\n onset: float = 0.5,\n offset: Optional[float] = None,\n min_duration_on: float = 0.0,\n min_duration_off: float = 0.0,\n pad_onset: float = 0.0,\n pad_offset: float = 0.0,\n max_duration: float = float('inf')\n ):\n\n super().__init__()\n\n self.onset = onset\n self.offset = offset or onset\n\n self.pad_onset = pad_onset\n self.pad_offset = pad_offset\n\n self.min_duration_on = min_duration_on\n self.min_duration_off = min_duration_off\n\n self.max_duration = max_duration\n\n def __call__(self, scores: SlidingWindowFeature) -> Annotation:\n \"\"\"Binarize detection scores\n Parameters\n ----------\n scores : SlidingWindowFeature\n Detection scores.\n Returns\n -------\n active : Annotation\n Binarized scores.\n \"\"\"\n\n num_frames, num_classes = scores.data.shape\n frames = scores.sliding_window\n timestamps = [frames[i].middle for i in range(num_frames)]\n\n # annotation meant to store 'active' regions\n active = Annotation()\n for k, k_scores in enumerate(scores.data.T):\n\n label = k if scores.labels is None else scores.labels[k]\n\n # initial state\n start = timestamps[0]\n is_active = k_scores[0] > self.onset\n curr_scores = [k_scores[0]]\n curr_timestamps = [start]\n t = start\n for t, y in zip(timestamps[1:], k_scores[1:]):\n # currently active\n if is_active:\n curr_duration = t - start\n if curr_duration > self.max_duration:\n search_after = len(curr_scores) // 2\n # divide segment\n min_score_div_idx = search_after + np.argmin(curr_scores[search_after:])\n min_score_t = curr_timestamps[min_score_div_idx]\n region = Segment(start - self.pad_onset, min_score_t + self.pad_offset)\n active[region, k] = label\n start = curr_timestamps[min_score_div_idx]\n curr_scores = curr_scores[min_score_div_idx + 1:]\n curr_timestamps = curr_timestamps[min_score_div_idx + 1:]\n # switching from active to inactive\n elif y < self.offset:\n region = Segment(start - self.pad_onset, t + self.pad_offset)\n active[region, k] = label\n start = t\n is_active = False\n curr_scores = []\n curr_timestamps = []\n curr_scores.append(y)\n curr_timestamps.append(t)\n # currently inactive\n else:\n # switching from inactive to active\n if y > self.onset:\n start = t\n is_active = True\n\n # if active at the end, add final region\n if is_active:\n region = Segment(start - self.pad_onset, t + self.pad_offset)\n active[region, k] = label\n\n # because of padding, some active regions might be overlapping: merge them.\n # also: fill same speaker gaps shorter than min_duration_off\n if self.pad_offset > 0.0 or self.pad_onset > 0.0 or self.min_duration_off > 0.0:\n if self.max_duration < float(\"inf\"):\n raise NotImplementedError(f\"This would break current max_duration param\")\n active = active.support(collar=self.min_duration_off)\n\n # remove tracks shorter than min_duration_on\n if self.min_duration_on > 0:\n for segment, track in list(active.itertracks()):\n if segment.duration < self.min_duration_on:\n del active[segment, track]\n\n return active\n\n\nclass VoiceActivitySegmentation(VoiceActivityDetection):\n def __init__(\n self,\n segmentation: PipelineModel = \"pyannote/segmentation\",\n fscore: bool = False,\n token: Union[Text, None] = None,\n **inference_kwargs,\n ):\n\n super().__init__(segmentation=segmentation, fscore=fscore, token=token, **inference_kwargs)\n\n def apply(self, file: AudioFile, hook: Optional[Callable] = None) -> Annotation:\n \"\"\"Apply voice activity detection\n\n Parameters\n ----------\n file : AudioFile\n Processed file.\n hook : callable, optional\n Hook called after each major step of the pipeline with the following\n signature: hook(\"step_name\", step_artefact, file=file)\n\n Returns\n -------\n speech : Annotation\n Speech regions.\n \"\"\"\n\n # setup hook (e.g. for debugging purposes)\n hook = self.setup_hook(file, hook=hook)\n\n # apply segmentation model (only if needed)\n # output shape is (num_chunks, num_frames, 1)\n if self.training:\n if self.CACHED_SEGMENTATION in file:\n segmentations = file[self.CACHED_SEGMENTATION]\n else:\n segmentations = self._segmentation(file)\n file[self.CACHED_SEGMENTATION] = segmentations\n else:\n segmentations: SlidingWindowFeature = self._segmentation(file)\n\n return segmentations\n\n\nclass Pyannote(Vad):\n\n def __init__(self, device, token=None, model_fp=None, **kwargs):\n logger.info(\"Performing voice activity detection using Pyannote...\")\n super().__init__(kwargs['vad_onset'])\n self.vad_pipeline = load_vad_model(device, token=token, model_fp=model_fp)\n\n def __call__(self, audio: AudioFile, **kwargs):\n return self.vad_pipeline(audio)\n\n @staticmethod\n def preprocess_audio(audio):\n return torch.from_numpy(audio).unsqueeze(0)\n\n @staticmethod\n def merge_chunks(segments,\n chunk_size,\n onset: float = 0.5,\n offset: Optional[float] = None,\n ):\n assert chunk_size > 0\n binarize = Binarize(max_duration=chunk_size, onset=onset, offset=offset)\n segments = binarize(segments)\n segments_list = []\n for speech_turn in segments.get_timeline():\n segments_list.append(SegmentX(speech_turn.start, speech_turn.end, \"UNKNOWN\"))\n\n if len(segments_list) == 0:\n logger.warning(\"No active speech found in audio\")\n return []\n assert segments_list, \"segments_list is empty.\"\n return Vad.merge_chunks(segments_list, chunk_size, onset, offset)\n" }, { "path": "whisperx/vads/silero.py", "content": "from io import IOBase\nfrom pathlib import Path\nfrom typing import Mapping, Text\nfrom typing import Optional\nfrom typing import Union\n\nimport torch\n\nfrom whisperx.diarize import Segment as SegmentX\nfrom whisperx.vads.vad import Vad\nfrom whisperx.log_utils import get_logger\n\nlogger = get_logger(__name__)\n\nAudioFile = Union[Text, Path, IOBase, Mapping]\n\n\nclass Silero(Vad):\n # check again default values\n def __init__(self, **kwargs):\n logger.info(\"Performing voice activity detection using Silero...\")\n super().__init__(kwargs['vad_onset'])\n\n self.vad_onset = kwargs['vad_onset']\n self.chunk_size = kwargs['chunk_size']\n self.vad_pipeline, vad_utils = torch.hub.load(repo_or_dir='snakers4/silero-vad',\n model='silero_vad',\n force_reload=False,\n onnx=False,\n trust_repo=True)\n (self.get_speech_timestamps, _, self.read_audio, _, _) = vad_utils\n\n def __call__(self, audio: AudioFile, **kwargs):\n \"\"\"use silero to get segments of speech\"\"\"\n # Only accept 16000 Hz for now.\n # Note: Silero models support both 8000 and 16000 Hz. Although other values are not directly supported,\n # multiples of 16000 (e.g. 32000 or 48000) are cast to 16000 inside of the JIT model!\n sample_rate = audio[\"sample_rate\"]\n if sample_rate != 16000:\n raise ValueError(\"Only 16000Hz sample rate is allowed\")\n\n timestamps = self.get_speech_timestamps(audio[\"waveform\"],\n model=self.vad_pipeline,\n sampling_rate=sample_rate,\n max_speech_duration_s=self.chunk_size,\n threshold=self.vad_onset\n # min_silence_duration_ms = self.min_duration_off/1000\n # min_speech_duration_ms = self.min_duration_on/1000\n # ...\n # See silero documentation for full option list\n )\n return [SegmentX(i['start'] / sample_rate, i['end'] / sample_rate, \"UNKNOWN\") for i in timestamps]\n\n @staticmethod\n def preprocess_audio(audio):\n return audio\n\n @staticmethod\n def merge_chunks(segments_list,\n chunk_size,\n onset: float = 0.5,\n offset: Optional[float] = None,\n ):\n assert chunk_size > 0\n if len(segments_list) == 0:\n logger.warning(\"No active speech found in audio\")\n return []\n assert segments_list, \"segments_list is empty.\"\n return Vad.merge_chunks(segments_list, chunk_size, onset, offset)\n" }, { "path": "whisperx/vads/vad.py", "content": "from typing import Optional\n\nimport pandas as pd\nfrom pyannote.core import Annotation, Segment\n\n\nclass Vad:\n def __init__(self, vad_onset):\n if not (0 < vad_onset < 1):\n raise ValueError(\n \"vad_onset is a decimal value between 0 and 1.\"\n )\n\n @staticmethod\n def preprocess_audio(audio):\n pass\n\n # keep merge_chunks as static so it can be also used by manually assigned vad_model (see 'load_model')\n @staticmethod\n def merge_chunks(segments,\n chunk_size,\n onset: float,\n offset: Optional[float]):\n \"\"\"\n Merge operation described in paper\n \"\"\"\n curr_end = 0\n merged_segments = []\n seg_idxs: list[tuple]= []\n speaker_idxs: list[Optional[str]] = []\n\n curr_start = segments[0].start\n for seg in segments:\n if seg.end - curr_start > chunk_size and curr_end - curr_start > 0:\n merged_segments.append({\n \"start\": curr_start,\n \"end\": curr_end,\n \"segments\": seg_idxs,\n })\n curr_start = seg.start\n seg_idxs = []\n speaker_idxs = []\n curr_end = seg.end\n seg_idxs.append((seg.start, seg.end))\n speaker_idxs.append(seg.speaker)\n # add final\n merged_segments.append({\n \"start\": curr_start,\n \"end\": curr_end,\n \"segments\": seg_idxs,\n })\n\n return merged_segments\n\n" } ]