giangndm/moss-transcribe-diarize-encoder

MOSS-Transcribe-Diarize Standalone Audio Encoder (Whisper + 4x Merge + VQAdaptor)

This repository contains the standalone Audio Encoder extracted from OpenMOSS-Team/MOSS-Transcribe-Diarize.

Complete Acoustic Backbone with 4x Temporal Merge & VQAdaptor: This model includes:

1. Whisper-Medium Encoder Backbone (24 layers, 1024 hidden dimension, 80 mel bins)
2. 4x Temporal Merge Layer (compresses 50 Hz acoustic frames by $4\times$ down to 12.5 Hz / 80ms per token)
3. VQAdaptor (Linear $4096 \to 1024$ $\to$ SiLU $\to$ Linear $1024 \to 1024$ $\to$ LayerNorm)

It produces 1024-dimensional continuous representations at 12.5 Hz, saving 75% sequence length in Large Language Model (LLM) KV cache compared to raw Whisper.

---

Model Specifications

Parameter	Value	Description
Backbone	Whisper-Medium	24 Transformer layers, 16 Attention Heads, GELU activation
Compression	4x Temporal Merge	Combines 4 consecutive 50 Hz frames into 1 frame
Adaptor	VQAdaptor	2-layer MLP projection ($4096 \to 1024 \to 1024$) + LayerNorm
Total Parameters	312.5M	307.2M Whisper + 5.2M VQAdaptor
Input Audio	16 kHz Mono	80-bin log-mel filterbank features
Output Frame Rate	12.5 Hz	80ms per frame (drastically reduces LLM prompt length)
Output Dimension	1024	Continuous acoustic embeddings (last_hidden_state)
Max Audio Length	30 seconds	Dynamic length without mandatory 30s padding
Multilingual	50+ languages	Pretrained and fine-tuned for speech transcription and diarization

---

Installation

pip install torch torchaudio soundfile transformers huggingface_hub safetensors

---

Quickstart & Usage

Method 1: Load via AutoModel with Remote Code (Recommended)

This repository includes custom remote code (trust_remote_code=True) that runs the full pipeline: Log-Mel $\to$ Whisper Encoder $\to$ 4x Temporal Merge $\to$ VQAdaptor.

import soundfile as sf
import torch
from transformers import AutoFeatureExtractor, AutoModel

model_id = "giangndm/moss-transcribe-diarize-encoder"

# 1. Load feature extractor and encoder
feature_extractor = AutoFeatureExtractor.from_pretrained(model_id)
encoder = AutoModel.from_pretrained(
    model_id,
    trust_remote_code=True,
    dtype=torch.bfloat16,
).eval()

device = "cuda" if torch.cuda.is_available() else "cpu"
encoder = encoder.to(device)

# 2. Extract 80-bin log-mel features from raw audio
audio, sr = sf.read("example.wav")
inputs = feature_extractor(audio, sampling_rate=16000, return_tensors="pt")
input_features = inputs.input_features.to(device=device, dtype=torch.bfloat16)

# 3. Forward pass
with torch.no_grad():
    outputs = encoder(input_features)
    acoustic_embeddings = outputs.last_hidden_state

# Shape: [batch_size, time_frames, 1024] at 12.5 Hz (80ms/frame)
print("12.5 Hz Acoustic embeddings shape:", acoustic_embeddings.shape)

---

Downstream Application Patterns

1. Audio-LLM Linear Projector (e.g. Qwen / LLaMA)

Since the encoder already outputs at 12.5 Hz (80ms/frame), downstream LLM projectors do NOT need convolutional downsampling:

import torch
import torch.nn as nn

class AudioLLMProjector(nn.Module):
    def __init__(self, encoder_dim=1024, llm_dim=4096):
        super().__init__()
        # Direct projection from 1024 -> 4096 at 12.5 Hz
        self.proj = nn.Linear(encoder_dim, llm_dim)

    def forward(self, x):
        # x: [B, T, 1024] -> [B, T, 4096]
        return self.proj(x)

2. Connectionist Temporal Classification (CTC) Head

import torch.nn as nn

class AudioCTCModel(nn.Module):
    def __init__(self, encoder, vocab_size=151936, hidden_dim=1024):
        super().__init__()
        self.encoder = encoder
        self.head = nn.Linear(hidden_dim, vocab_size)

    def forward(self, input_features):
        hidden = self.encoder(input_features).last_hidden_state
        return self.head(hidden)

---

Citation & Acknowledgements

• Original Model: OpenMOSS-Team/MOSS-Transcribe-Diarize
• Paper: MOSS-Transcribe-Diarize 0.9B: An End-to-End Audio Understanding Model for Long-Form Multi-Speaker Transcription and Diarization
• License: Apache 2.0