Whisper架构拆解：Transformer做多任务ASR

import numpy as np def mel_filter_bank(n_fft=400, n_mels=80, sample_rate=16000):    mel_min = 2595.0 * np.log10(1.0 + 0 / 700.0)    mel_max = 2595.0 * np.log10(1.0 + (sample_rate / 2) / 700.0)    mel_points = np.linspace(mel_min, mel_max, n_mels + 2)    hz_points = 700.0 * (10.0 ** (mel_points / 2595.0) - 1.0)    bins = np.floor((n_fft + 1) * hz_points / sample_rate).astype(int)    filters = np.zeros((n_mels, n_fft // 2 + 1))    for i in range(n_mels):        left, center, right = bins[i], bins[i+1], bins[i+2]        for j in range(left, center):            filters[i, j] = (j - left) / max(center - left, 1)        for j in range(center, right):            filters[i, j] = (right - j) / max(right - center, 1)    return filters def compute_log_mel_spectrogram(audio, n_fft=400, hop_length=160, n_mels=80):    window = np.hanning(n_fft)    n_frames = 1 + (len(audio) - n_fft) // hop_length    stft = np.zeros((n_fft // 2 + 1, n_frames))    for i in range(n_frames):        frame = audio[i*hop_length : i*hop_length + n_fft] * window        spectrum = np.fft.rfft(frame)        stft[:, i] = np.abs(spectrum) ** 2    mel_filters = mel_filter_bank(n_fft=n_fft, n_mels=n_mels)    mel_spec = mel_filters @ stft    log_mel = np.log10(np.clip(mel_spec, 1e-10, None))    log_mel = np.maximum(log_mel, log_mel.max() - 8.0)    log_mel = (log_mel + 4.0) / 4.0    return log_mel class SimpleAttention:    def __init__(self, d_model, n_heads):        self.d_model = d_model        self.n_heads = n_heads        self.d_k = d_model // n_heads        self.W_q = np.random.randn(d_model, d_model) * 0.01        self.W_k = np.random.randn(d_model, d_model) * 0.01        self.W_v = np.random.randn(d_model, d_model) * 0.01        self.W_o = np.random.randn(d_model, d_model) * 0.01     def forward(self, x):        Q = x @ self.W_q        K = x @ self.W_k        V = x @ self.W_v        scores = Q @ K.T / np.sqrt(self.d_k)        attn = np.exp(scores) / np.exp(scores).sum(axis=-1, keepdims=True)        output = attn @ V        return output @ self.W_o class WhisperLikeModel:    def __init__(self, d_model=64, n_heads=4, n_layers=2, n_tasks=4):        self.d_model = d_model        self.encoder_layers = [SimpleAttention(d_model, n_heads) for _ in range(n_layers)]        self.decoder_layers = [SimpleAttention(d_model, n_heads) for _ in range(n_layers)]        self.task_embeddings = np.random.randn(n_tasks, d_model) * 0.01        self.output_proj = np.random.randn(d_model, 100) * 0.01        self.task_names = {0: "transcribe", 1: "translate", 2: "language_detect", 3: "lm"}     def encode(self, mel_input):        x = mel_input.T        x = x @ np.random.randn(x.shape[-1], self.d_model) * 0.01        for layer in self.encoder_layers:            x = x + layer.forward(x)        return x     def decode(self, encoder_output, task_id):        task_token = self.task_embeddings[task_id]        x = np.tile(task_token, (encoder_output.shape[0], 1))        x = x + encoder_output        for layer in self.decoder_layers:            x = x + layer.forward(x)        logits = x @ self.output_proj        return logits     def run(self, audio, task="transcribe"):        task_map = {"transcribe": 0, "translate": 1, "language_detect": 2, "lm": 3}        task_id = task_map[task]        mel = compute_log_mel_spectrogram(audio)        mel_padded = np.zeros((mel.shape[0], 3000))        mel_padded[:, :mel.shape[1]] = mel[:, :min(mel.shape[1], 3000)]        enc_out = self.encode(mel_padded)        logits = self.decode(enc_out, task_id)        print(f"Task: {task} | Output shape: {logits.shape}")        return logits if __name__ == "__main__":    audio = np.random.randn(16000 * 5).astype(np.float32)    model = WhisperLikeModel(d_model=64, n_heads=4, n_layers=2)    for task in ["transcribe", "translate", "language_detect"]:        model.run(audio, task=task)    print("Multi-task Whisper-like model demo completed!")