Hugging Face Inference Endpoints ASR and Diarization Pipeline

Hugging Face has released a custom inference handler that enables the deployment of a modular pipeline combining Automatic Speech Recognition (ASR), speaker diarization, and speculative decoding via Hugging Face Inference Endpoints. This implementation allows developers to expose a single API endpoint that integrates multiple models, specifically leveraging OpenAI's Whisper for transcription and Pyannote for speaker identification.

Modular Pipeline Architecture

The pipeline is designed as a modular system where components can be enabled or disabled based on the specific use case. The core components include:

  • ASR Module: Utilizes Whisper models for high-quality speech-to-text transcription.
  • Diarization Module: Employs the Pyannote speaker-diarization-3.1 model, a state-of-the-art open-source implementation, to identify and partition transcriptions by speaker.
  • Speculative Decoding: Accelerates inference by using a smaller assistant model (such as a distilled Whisper model) to suggest generations that are then validated by the larger main model.

Technical Requirements and Constraints

To optimize performance, the implementation utilizes PyTorch 2.2, which provides out-of-the-box support for Flash Attention 2 via SDPA.

Speculative decoding introduces specific technical restrictions:

  • Architecture Match: The decoder portion of the assistant model must share the same architecture as the main model.
  • Batch Size: Speculative decoding requires a batch size of 1. For production environments requiring larger batches, standard inference may be faster than assisted generation.

Performance Benchmarks

Benchmarks conducted on an NVIDIA A10 GPU using openai/whisper-large-v3 as the main model and distil-whisper/distil-large-v3 as the assistant model demonstrate that speculative decoding is highly effective for short audio clips but less efficient for long ones:

  • Short Audio (8s): Assisted generation averaged 326.96ms, compared to 784.35ms for non-assisted generation.
  • Long Audio (60s): Non-assisted generation averaged 3.48s, while assisted generation averaged 4.15s.

This performance difference occurs because long audio is automatically chunked into batches, which conflicts with the batch size 1 limitation of speculative decoding.

Deployment and Configuration

Deployment is managed through a custom handler consisting of three primary files: handler.py (initialization and inference), diarization_utils.py (pre- and post-processing), and config.py (settings).

Configuration Parameters

Model settings are managed via ModelSettings and InferenceConfig classes. Users can adjust parameters using environment variables or API calls during endpoint creation. Key parameters include:

  • Task: Choice between transcribe and translate.
  • Batch Size: Default is 24 (must be 1 for assisted generation).
  • Assisted Flag: A boolean determining whether to use speculative decoding.
  • Speaker Constraints: Options to define num_speakers, min_speakers, and max_speakers for the diarization pipeline.

API Integration

Once deployed, the endpoint accepts base64-encoded audio files. Requests are sent via a JSON payload containing the audio in the inputs field and a parameters dictionary for adjusting InferenceConfig settings. This can be implemented using standard Python requests or the Hugging Face InferenceClient.

Sources