Architectures / Audio

🎙️ Whisper Small

Whisper speech encoder-decoder — conv1d audio stem + transformer encoder/decoder (384D)

Every number on this page is computed from the graph by the same functions the app runs, not written by hand.

When to pick it

Structure

18 layers. Output shapes are propagated from the input shape, batch dimension excluded.

What the verifier says

The same 41 structural checks that run on every edit in the app, on this graph.

The PyTorch it exports

Generated from the graph above. First 46 lines; the app exports the whole file, plus the training loop, the data contract and a deploy bundle.

# Architecture designed with Neurarch: https://neurarch.com
# PyTorch: compatible with Python 3.8+ and torch>=1.12
# Colab: pip install torch torchvision  (usually pre-installed)
# Audio: pip install torchaudio
#
# WARNING: 1 layer(s) below are not yet supported by the PyTorch
# exporter and pass their input through UNCHANGED in forward():
#   - to_tokens (permute)

import torch
import torch.nn as nn
import torch.nn.functional as F
from typing import Tuple
import torchaudio

class WhisperSmall(nn.Module):
    def __init__(self):
        super().__init__()

        self.audioConv_1 = nn.Conv1d(1, 384, kernel_size=3, stride=1, padding=1)
        self.gelu_1 = nn.GELU()
        self.audioConv_2 = nn.Conv1d(1, 384, kernel_size=3, stride=2, padding=1)
        self.gelu_2 = nn.GELU()
        self.transformerBlock_1 = nn.TransformerEncoderLayer(d_model=384, nhead=6, dim_feedforward=1536, batch_first=True)
        self.transformerBlock_2 = nn.TransformerEncoderLayer(d_model=384, nhead=6, dim_feedforward=1536, batch_first=True)
        self.layerNorm_1 = nn.LayerNorm(384)
        self.embedding_1 = nn.Embedding(51865, 384)
        self.transformerBlock_3 = nn.TransformerDecoderLayer(d_model=384, nhead=6, dim_feedforward=1536, batch_first=True)
        self.transformerBlock_4 = nn.TransformerDecoderLayer(d_model=384, nhead=6, dim_feedforward=1536, batch_first=True)
        self.layerNorm_2 = nn.LayerNorm(384)
        self.linear_1 = nn.Linear(384, 51865)

    def forward(self, src, tgt=None):
        # mel_features shape: [1,80,3000]
        # decoder_tokens shape: [1,448]
        audio_conv_Conv_1 = self.audioConv_1(src)
        gelu_gelu_1 = self.gelu_1(audio_conv_Conv_1)
        audio_conv_Conv_2 = self.audioConv_2(gelu_gelu_1)
        gelu_gelu_2 = self.gelu_2(audio_conv_Conv_2)
        # TODO: layer 'to_tokens' (permute) is not yet supported by the exporter; passing through unchanged
        # positionalEncoding: add positional encoding externally (e.g. sinusoidal or learned PE)
        transformer_block__enc_1 = self.transformerBlock_1(gelu_gelu_2)
        transformer_block__enc_2 = self.transformerBlock_2(transformer_block__enc_1)
        layer_norm_rm_enc = self.layerNorm_1(transformer_block__enc_2)
        embedding_ng_dec = self.embedding_1(tgt)
        # positionalEncoding: add positional encoding externally (e.g. sinusoidal or learned PE)

For agents