Powerful ASR + diarization + speculative decoding with Hugging Face Inference Endpoints

Whisper is one of the best open source speech recognition models and definitely the one most widely used. Hugging Face Inference Endpoints make it very easy to deploy any Whisper model out of the box. However, if you’d like to introduce additional features, like a diarization pipeline to identify speakers, or assisted generation for speculative decoding, things get trickier. The reason is that you need to combine Whisper with additional models, while still exposing a single API endpoint.

We'll solve this challenge using a custom inference handler, which will implement the Automatic Speech Recognition (ASR) and Diarization pipeline on Inference Endpoints, as well as supporting speculative decoding. The implementation of the diarization pipeline is inspired by the famous Insanely Fast Whisper, and it uses a Pyannote model for diarization.

This will also be a demonstration of how flexible Inference Endpoints are and that you can host pretty much anything there. Here is the code to follow along. Note that during initialization of the endpoint, the whole repository gets mounted, so your handler.py can refer to other files in your repository if you prefer not to have all the logic in a single file. In this case, we decided to separate things into several files to keep things clean:

• handler.py contains initialization and inference code
• diarization_utils.py has all the diarization-related pre- and post-processing
• config.py has ModelSettings and InferenceConfig. ModelSettings define which models will be utilized in the pipeline (you don't have to use all of them), and InferenceConfig defines the default inference parameters

Starting with Pytorch 2.2, SDPA supports Flash Attention 2 out-of-the-box, so we'll use that version for faster inference.

The main modules

This is a high-level diagram of what the endpoint looks like under the hood:

Speculative decoding comes with restrictions:

If you do use an assistant model, a great choice for Whisper is a distilled version.