Speech Recognition (Whisper)

Build an on-device speech recognition application using OpenAI's Whisper model with ZETIC Melange. This tutorial covers splitting Whisper into encoder and decoder components, deploying them to Melange, and running the full speech-to-text pipeline on Android and iOS.

What You Will Build

An on-device speech-to-text application that processes audio through Whisper's three-component architecture (Feature Extractor, Encoder, Decoder) with NPU acceleration for real-time transcription.

Prerequisites

A ZETIC Melange account with a Personal Key (sign up at melange.zetic.ai)
Python 3.8+ with torch, transformers, datasets, and numpy installed
Android Studio or Xcode for mobile deployment

What is Whisper?

Whisper is a state-of-the-art speech recognition model developed by OpenAI that offers:

Multilingual support: Recognizes speech in multiple languages
Multiple capabilities: Performs speech recognition, language detection, and translation
Open source: Available through Hugging Face

Architecture Overview

The Whisper implementation consists of three main components:

Feature Extractor: Processes raw audio into Mel Spectrogram features
Encoder: Processes Mel Spectrogram to generate audio embeddings
Decoder: Generates text tokens from the audio embeddings

Step 1: Prepare Sample Inputs for Tracing

To convert the model for deployment, we need to trace the PyTorch model with sample inputs that match the expected tensor shapes.

from datasets import load_dataset
import numpy as np
import torch
from transformers import WhisperFeatureExtractor, WhisperForConditionalGeneration

model_name = "openai/whisper-tiny"
model = WhisperForConditionalGeneration.from_pretrained(model_name)
feature_extractor = WhisperFeatureExtractor.from_pretrained(model_name, sampling_rate=16_000)

ds = load_dataset("hf-internal-testing/librispeech_asr_dummy", "clean", split="validation")
inputs = feature_extractor(ds[0]["audio"]["array"], return_tensors="pt")
input_features = inputs.input_features

MAX_TOKEN_LENGTH = model.config.max_target_positions
dummy_decoder_input_ids = torch.tensor([[0 for _ in range(MAX_TOKEN_LENGTH)]])
dummy_encoder_hidden_states = torch.randn(1, 1500, model.config.d_model).float()
dummy_decoder_attention_mask = torch.ones_like(dummy_decoder_input_ids)

Step 2: Convert Encoder to TorchScript

Wrap and trace the Whisper encoder:

from transformers import WhisperModel
import torch.nn as nn

class WhisperEncoderWrapper(nn.Module):
    def __init__(self, whisper_model):
        super().__init__()
        self.enc = whisper_model.model.encoder

    def forward(self, input_features):
        return self.enc(input_features=input_features, return_dict=False)[0]

with torch.no_grad():
    encoder = WhisperEncoderWrapper(model).eval()
    traced_encoder = torch.jit.trace(encoder, input_features)
    traced_encoder.save("whisper_encoder.pt")

Step 3: Convert Decoder to TorchScript

Wrap and trace the Whisper decoder:

class WhisperDecoderWrapper(nn.Module):
    def __init__(self, whisper_model):
        super().__init__()
        self.decoder = whisper_model.model.decoder
        self.proj_out = whisper_model.proj_out

    def forward(self, input_ids, encoder_hidden_states, decoder_attention_mask):
        hidden = self.decoder(
            input_ids=input_ids,
            encoder_hidden_states=encoder_hidden_states,
            attention_mask=decoder_attention_mask,
            use_cache=False,
            return_dict=False,
        )[0]
        return self.proj_out(hidden)

with torch.no_grad():
    decoder = WhisperDecoderWrapper(model).eval()
    traced_decoder = torch.jit.trace(
        decoder,
        (dummy_decoder_input_ids, dummy_encoder_hidden_states, dummy_decoder_attention_mask)
    )
    traced_decoder.save("whisper_decoder.pt")

Step 4: Save Input Samples

Save all input tensors as .npy files for model upload:

import numpy as np

# Save encoder inputs
np.save("whisper_input_features.npy", input_features.cpu().numpy())

# Save decoder inputs
np.save(
    "whisper_decoder_input_ids.npy",
    dummy_decoder_input_ids.cpu().numpy().astype(np.int64),
)
np.save(
    "whisper_encoder_hidden_states.npy",
    dummy_encoder_hidden_states.cpu().numpy().astype(np.float32),
)
np.save(
    "whisper_decoder_attention_mask.npy",
    dummy_decoder_attention_mask.cpu().numpy().astype(np.int64),
)

Step 5: Generate Melange Models

Upload both models and their inputs via the Melange Dashboard or use the CLI:

# Upload encoder model
zetic gen -p $PROJECT_NAME -i whisper_input_features.npy whisper_encoder.pt

# Upload decoder model
zetic gen -p $PROJECT_NAME \
 -i whisper_decoder_input_ids.npy \
 -i whisper_encoder_hidden_states.npy \
 -i whisper_decoder_attention_mask.npy \
 whisper_decoder.pt

The decoder model requires three input files. Make sure to provide them in the correct order as shown above. See Supported Formats for details on input ordering.

Step 6: Implement ZeticMLangeModel

We provide model keys for the demo app: OpenAI/whisper-tiny-encoder and OpenAI/whisper-tiny-decoder. You can use these model keys to try the Melange Application.

For detailed application setup, please follow the Android Integration Guide guide.

val encoderModel = ZeticMLangeModel(this, PERSONAL_KEY, "OpenAI/whisper-tiny-encoder")
val decoderModel = ZeticMLangeModel(this, PERSONAL_KEY, "OpenAI/whisper-tiny-decoder")

// Run encoder
val outputs = encoderModel.run(inputs)

// Run decoder
decoderModel.run(..., ...)

For detailed application setup, please follow the iOS Integration Guide guide.

let encoderModel = try ZeticMLangeModel(personalKey: PERSONAL_KEY, name: "OpenAI/whisper-tiny-encoder")
let decoderModel = try ZeticMLangeModel(personalKey: PERSONAL_KEY, name: "OpenAI/whisper-tiny-decoder")

// Run encoder
let outputs = try encoderModel.run(inputs)

// Run decoder
try decoderModel.run(..., ...)

Step 7: Use the Whisper Feature Wrapper

The WhisperFeatureWrapper handles audio-to-Mel-Spectrogram conversion and token decoding.

You can find WhisperDecoder and WhisperEncoder implementations in ZETIC Melange apps.

Complete Speech Recognition Implementation

// Initialize components
val whisper = WhisperFeatureWrapper()
val encoder = ZeticMLangeModel(this, PERSONAL_KEY, "OpenAI/whisper-tiny-encoder")
val decoder = ZeticMLangeModel(this, PERSONAL_KEY, "OpenAI/whisper-tiny-decoder")

// Process audio
val features = whisper.process(audioData)

// Run encoder
encoder.process(features)

// Generate tokens using decoder
val generatedIds = decoder.generateTokens(outputs)

// Convert tokens to text
val text = whisper.decodeToken(generatedIds.toIntArray(), true)

// Initialize components
let wrapper = WhisperFeatureWrapper()
let encoder = try ZeticMLangeModel(personalKey: PERSONAL_KEY, name: "OpenAI/whisper-tiny-encoder")
let decoder = try ZeticMLangeModel(personalKey: PERSONAL_KEY, name: "OpenAI/whisper-tiny-decoder")

// Process audio to features
let features = wrapper.process(input.audio)

// Run encoder
let outputs = encoder.process(features)

// Generate tokens using decoder
let generatedIds = decoder.process(outputs)

// Convert tokens to text
let text = wrapper.decodeToken(generatedIds, true)
return WhisperOutput(text: text)

Conclusion

With ZETIC Melange, implementing on-device speech recognition with NPU acceleration is straightforward and efficient. Whisper provides robust multilingual speech recognition and translation capabilities. The three-component pipeline (Feature Extractor, Encoder, Decoder) is cleanly abstracted through the Melange SDK.

We are continuously adding new models to our examples and HuggingFace page.

Stay tuned, and contact us for collaborations!

Speech Recognition (Whisper)

On this page