Upload Main Files

configuration_mini_whisper.py

@@ -0,0 +1,167 @@
+# configuration_mini_whisper.py
+from transformers import PretrainedConfig
+
+NON_SPEECH_TOKENS = [
+    1, 2, 7, 8, 9, 10, 14, 25,
+    26, 27, 28, 29, 31, 58, 59, 60, 61, 62,
+    63, 90, 91, 92, 93, 357, 366, 438, 532, 685,
+    705, 796, 930, 1058, 1220, 1267, 1279, 1303, 1343, 1377,
+    1391, 1635, 1782, 1875, 2162, 2361, 2488, 3467, 4008, 4211,
+    4600, 4808, 5299, 5855, 6329, 7203, 9609, 9959, 10563, 10786,
+    11420, 11709, 11907, 13163, 13697, 13700, 14808, 15306, 16410, 16791,
+    17992, 19203, 19510, 20724, 22305, 22935, 27007, 30109, 30420, 33409,
+    34949, 40283, 40493, 40549, 47282, 49146, 50257, 50359, 50360, 50361
+]
+
+NON_SPEECH_TOKENS_MULTI = [
+    1, 2, 7, 8, 9, 10, 14, 25,
+    26, 27, 28, 29, 31, 58, 59, 60, 61, 62,
+    63, 90, 91, 92, 93, 359, 503, 522, 542, 873,
+    893, 902, 918, 922, 931, 1350, 1853, 1982, 2460, 2627,
+    3246, 3253, 3268, 3536, 3846, 3961, 4183, 4667, 6585, 6647,
+    7273, 9061, 9383, 10428, 10929, 11938, 12033, 12331, 12562, 13793,
+    14157, 14635, 15265, 15618, 16553, 16604, 18362, 18956, 20075, 21675,
+    22520, 26130, 26161, 26435, 28279, 29464, 31650, 32302, 32470, 36865,
+    42863, 47425, 49870, 50254, 50258, 50360, 50361, 50362
+]
+
+class MiniWhisperConfig(PretrainedConfig):
+    model_type = "mini_whisper"
+    keys_to_ignore_at_inference = ["past_key_values"]
+    attribute_map = {
+        "num_key_value_heads": "encoder_attention_heads",
+        "num_attention_heads": "encoder_attention_heads",
+        "hidden_size": "d_model",
+        "num_hidden_layers": "encoder_layers",
+    }
+
+    vocab_size: int = 51865
+    num_mel_bins: int = 80
+    encoder_layers: int = 4
+    encoder_attention_heads: int = 6
+    decoder_layers: int = 4
+    decoder_attention_heads: int = 6
+    decoder_ffn_dim: int = 1536
+    encoder_ffn_dim: int = 1536
+    encoder_layerdrop: float | int = 0.0
+    decoder_layerdrop: float | int = 0.0
+    decoder_start_token_id: int = 50257
+    use_cache: bool = True
+    is_encoder_decoder: bool = True
+    activation_function: str = "gelu"
+    d_model: int = 384
+    dropout: float | int = 0.0
+    attention_dropout: float | int = 0.0
+    activation_dropout: float | int = 0.0
+    init_std: float = 0.02
+    scale_embedding: bool = False
+    max_source_positions: int = 1500
+    max_target_positions: int = 448
+    pad_token_id: int | None = 50256
+    bos_token_id: int | None = 50256
+    eos_token_id: int | list[int] | None = 50256
+    suppress_tokens: list | None = None
+    begin_suppress_tokens: list[int] | tuple[int, ...] | None = (220, 50256)
+    use_weighted_layer_sum: bool = False
+    classifier_proj_size: int = 256
+    apply_spec_augment: bool = False
+    mask_time_prob: float | int = 0.05
+    mask_time_length: int = 10
+    mask_time_min_masks: int = 2
+    mask_feature_prob: float | int = 0.0
+    mask_feature_length: int = 10
+    mask_feature_min_masks: int = 0
+    median_filter_width: int = 7
+    tie_word_embeddings: bool = True
+
+    def __init__(
+        self,
+        vocab_size=51865,
+        num_mel_bins=80,
+        d_model=384,
+        encoder_layers=4,
+        decoder_layers=4,
+        encoder_attention_heads=6,
+        decoder_attention_heads=6,
+        encoder_ffn_dim=1536,
+        decoder_ffn_dim=1536,
+        encoder_layerdrop=0.0,
+        decoder_layerdrop=0.0,
+        decoder_start_token_id=50257,
+        use_cache=True,
+        is_encoder_decoder=True,
+        activation_function="gelu",
+        dropout=0.0,
+        attention_dropout=0.0,
+        activation_dropout=0.0,
+        init_std=0.02,
+        scale_embedding=False,
+        max_source_positions=1500,
+        max_target_positions=448,
+        pad_token_id=50256,
+        bos_token_id=50256,
+        eos_token_id=50256,
+        suppress_tokens=None,
+        begin_suppress_tokens=(220, 50256),
+        use_weighted_layer_sum=False,
+        classifier_proj_size=256,
+        apply_spec_augment=False,
+        mask_time_prob=0.05,
+        mask_time_length=10,
+        mask_time_min_masks=2,
+        mask_feature_prob=0.0,
+        mask_feature_length=10,
+        mask_feature_min_masks=0,
+        median_filter_width=7,
+        tie_word_embeddings=True,
+        **kwargs
+    ):
+        self.vocab_size = vocab_size
+        self.num_mel_bins = num_mel_bins
+        self.d_model = d_model
+        self.encoder_layers = encoder_layers
+        self.decoder_layers = decoder_layers
+        self.encoder_attention_heads = encoder_attention_heads
+        self.decoder_attention_heads = decoder_attention_heads
+        self.encoder_ffn_dim = encoder_ffn_dim
+        self.decoder_ffn_dim = decoder_ffn_dim
+        self.encoder_layerdrop = encoder_layerdrop
+        self.decoder_layerdrop = decoder_layerdrop
+        self.decoder_start_token_id = decoder_start_token_id
+        self.use_cache = use_cache
+        self.is_encoder_decoder = is_encoder_decoder
+        self.activation_function = activation_function
+        self.dropout = dropout
+        self.attention_dropout = attention_dropout
+        self.activation_dropout = activation_dropout
+        self.init_std = init_std
+        self.scale_embedding = scale_embedding
+        self.max_source_positions = max_source_positions
+        self.max_target_positions = max_target_positions
+        self.pad_token_id = pad_token_id
+        self.bos_token_id = bos_token_id
+        self.eos_token_id = eos_token_id
+        self.suppress_tokens = suppress_tokens if suppress_tokens is not None else NON_SPEECH_TOKENS
+        self.begin_suppress_tokens = list(begin_suppress_tokens)
+        self.use_weighted_layer_sum = use_weighted_layer_sum
+        self.classifier_proj_size = classifier_proj_size
+        self.apply_spec_augment = apply_spec_augment
+        self.mask_time_prob = mask_time_prob
+        self.mask_time_length = mask_time_length
+        self.mask_time_min_masks = mask_time_min_masks
+        self.mask_feature_prob = mask_feature_prob
+        self.mask_feature_length = mask_feature_length
+        self.mask_feature_min_masks = mask_feature_min_masks
+        self.median_filter_width = median_filter_width
+        self.tie_word_embeddings = tie_word_embeddings
+
+        super().__init__(
+            pad_token_id=pad_token_id,
+            bos_token_id=bos_token_id,
+            eos_token_id=eos_token_id,
+            decoder_start_token_id=decoder_start_token_id,
+            tie_word_embeddings=tie_word_embeddings,
+            **kwargs,
+        )
+
+__all__ = ["MiniWhisperConfig"]

feature_extraction_mini_whisper.py

@@ -0,0 +1,253 @@
+# feature_extraction_mini_whisper.py
+import numpy as np
+import torch
+
+
+class MiniWhisperFeatureExtractor:
+    """Feature extractor matching Whisper exactly."""
+    
+    model_input_names = ["input_features"]
+    
+    def __init__(
+        self,
+        feature_size=80,
+        sampling_rate=16000,
+        hop_length=160,
+        chunk_length=30,
+        n_fft=400,
+        padding_value=0.0,
+        dither=0.0,
+        return_attention_mask=False,
+        **kwargs,
+    ):
+        self.feature_size = feature_size
+        self.sampling_rate = sampling_rate
+        self.hop_length = hop_length
+        self.chunk_length = chunk_length
+        self.n_fft = n_fft
+        self.n_samples = chunk_length * sampling_rate
+        self.nb_max_frames = self.n_samples // hop_length
+        self.padding_value = padding_value
+        self.dither = dither
+        self.return_attention_mask = return_attention_mask
+        
+        self.mel_filters = self._mel_filter_bank()
+    
+    def _mel_filter_bank(self):
+        """Create mel filter bank matching Whisper's parameters."""
+        try:
+            from torchaudio.functional import melscale_fbanks
+            n_freq_bins = 1 + self.n_fft // 2
+            mel_filters = melscale_fbanks(
+                n_freqs=n_freq_bins,
+                f_min=0.0,
+                f_max=8000.0,
+                n_mels=self.feature_size,
+                sample_rate=self.sampling_rate,
+                norm="slaney",
+                mel_scale="slaney",
+            )
+            return mel_filters.numpy()
+        except ImportError:
+            return self._numpy_mel_filter_bank()
+    
+    def _numpy_mel_filter_bank(self):
+        """NumPy fallback for mel filter bank."""
+        def hz_to_mel(hz):
+            return 2595 * np.log10(1 + hz / 700)
+        
+        def mel_to_hz(mel):
+            return 700 * (10 ** (mel / 2595) - 1)
+        
+        n_freq_bins = 1 + self.n_fft // 2
+        fmin, fmax = 0.0, 8000.0
+        
+        mel_min = hz_to_mel(fmin)
+        mel_max = hz_to_mel(fmax)
+        mel_points = np.linspace(mel_min, mel_max, self.feature_size + 2)
+        hz_points = mel_to_hz(mel_points)
+        
+        bin_points = np.floor((self.n_fft + 1) * hz_points / self.sampling_rate).astype(int)
+        
+        filters = np.zeros((self.feature_size, n_freq_bins))
+        for i in range(self.feature_size):
+            left, center, right = bin_points[i], bin_points[i + 1], bin_points[i + 2]
+            for j in range(left, center):
+                filters[i, j] = (j - left) / (center - left)
+            for j in range(center, right):
+                filters[i, j] = (right - j) / (right - center)
+        
+        filters = filters / (filters.sum(axis=1, keepdims=True) + 1e-10)
+        return filters
+    
+    def _np_extract_fbank_features(self, waveform_batch, device="cpu"):
+        """NumPy implementation matching Whisper exactly."""
+        log_spec_batch = []
+        for waveform in waveform_batch:
+            # Apply dither
+            if self.dither != 0.0:
+                waveform = waveform + self.dither * np.random.randn(*waveform.shape).astype(np.float32)
+            
+            # Compute STFT
+            n_frames = 1 + (len(waveform) - self.n_fft) // self.hop_length
+            spec = np.zeros((self.n_fft // 2 + 1, n_frames), dtype=np.float32)
+            
+            window = np.hanning(self.n_fft).astype(np.float32)
+            for i in range(n_frames):
+                start = i * self.hop_length
+                frame = waveform[start:start + self.n_fft]
+                if len(frame) < self.n_fft:
+                    frame = np.pad(frame, (0, self.n_fft - len(frame)))
+                spec[:, i] = np.abs(np.fft.rfft(frame * window)) ** 2
+            
+            # Apply mel filters: mel_filters is [n_mels, n_freqs]
+            mel_spec = self.mel_filters @ spec
+            
+            # Log10 with floor
+            log_spec = np.log10(np.maximum(mel_spec, 1e-10))
+            
+            # Whisper uses: np.maximum(log_spec, log_spec.max() - 8.0)
+            log_spec = np.maximum(log_spec, log_spec.max() - 8.0)
+            
+            # Normalize: (log_spec + 4.0) / 4.0
+            log_spec = (log_spec + 4.0) / 4.0
+            
+            # Remove last frame to match Whisper
+            log_spec = log_spec[:, :-1]
+            
+            log_spec_batch.append(log_spec)
+        
+        return np.array(log_spec_batch)
+    
+    def _torch_extract_fbank_features(self, waveform, device="cpu"):
+        """PyTorch implementation matching Whisper exactly."""
+        waveform = torch.from_numpy(waveform).to(device, torch.float32)
+        window = torch.hann_window(self.n_fft, device=device)
+        
+        if self.dither != 0.0:
+            waveform += self.dither * torch.randn(waveform.shape, dtype=waveform.dtype, device=waveform.device)
+        
+        stft = torch.stft(waveform, self.n_fft, self.hop_length, window=window, return_complex=True)
+        magnitudes = stft[..., :-1].abs() ** 2  # Remove last frame
+        
+        # mel_filters shape: [n_mels, n_freqs] from mel_filter_bank
+        mel_filters = torch.from_numpy(self.mel_filters).to(device, torch.float32)
+        # Transpose and matmul: mel_filters.T @ magnitudes
+        mel_spec = mel_filters.T @ magnitudes
+        
+        log_spec = torch.clamp(mel_spec, min=1e-10).log10()
+        
+        if waveform.dim() == 2:
+            max_val = log_spec.max(dim=2, keepdim=True)[0].max(dim=1, keepdim=True)[0]
+            log_spec = torch.maximum(log_spec, max_val - 8.0)
+        else:
+            log_spec = torch.maximum(log_spec, log_spec.max() - 8.0)
+        
+        log_spec = (log_spec + 4.0) / 4.0
+        
+        if device != "cpu":
+            log_spec = log_spec.detach().cpu()
+        
+        return log_spec.numpy()
+    
+    @staticmethod
+    def zero_mean_unit_var_norm(input_values, attention_mask, padding_value=0.0):
+        if attention_mask is not None:
+            attention_mask = np.array(attention_mask, np.int32)
+            normed = []
+            for vector, length in zip(input_values, attention_mask.sum(-1)):
+                normed_slice = (vector - vector[:length].mean()) / (np.sqrt(vector[:length].var() + 1e-7))
+                if length < normed_slice.shape[0]:
+                    normed_slice[length:] = padding_value
+                normed.append(normed_slice)
+        else:
+            normed = [(x - x.mean()) / (np.sqrt(x.var() + 1e-7)) for x in input_values]
+        return normed
+    
+    def __call__(
+        self,
+        raw_speech,
+        sampling_rate=None,
+        truncation=True,
+        return_tensors="pt",
+        return_attention_mask=None,
+        do_normalize=False,
+        **kwargs,
+    ):
+        if sampling_rate is not None and sampling_rate != self.sampling_rate:
+            raise ValueError(f"Expected {self.sampling_rate}, got {sampling_rate}")
+        
+        is_batched = isinstance(raw_speech, list) or (isinstance(raw_speech, np.ndarray) and raw_speech.ndim > 1)
+        
+        if not is_batched:
+            raw_speech = [raw_speech]
+        
+        # Convert to numpy float32
+        processed = []
+        for audio in raw_speech:
+            if isinstance(audio, np.ndarray) and audio.dtype == np.float64:
+                audio = audio.astype(np.float32)
+            elif not isinstance(audio, np.ndarray):
+                audio = np.array(audio, dtype=np.float32)
+            
+            # Pad or truncate
+            if len(audio) > self.n_samples:
+                if truncation:
+                    audio = audio[:self.n_samples]
+            elif len(audio) < self.n_samples:
+                audio = np.pad(audio, (0, self.n_samples - len(audio)), constant_values=self.padding_value)
+            
+            processed.append(audio)
+        
+        waveform_batch = np.stack(processed, axis=0)
+        
+        # Use torch if available
+        if torch.cuda.is_available():
+            input_features = self._torch_extract_fbank_features(waveform_batch, device="cuda")
+            # Result is [batch, n_mels, n_frames]
+        else:
+            input_features = self._np_extract_fbank_features(waveform_batch)
+        
+        result = {"input_features": input_features}
+        
+        if return_attention_mask is None:
+            return_attention_mask = self.return_attention_mask
+        
+        if return_attention_mask:
+            result["attention_mask"] = np.ones(input_features.shape[0], dtype=np.int64)
+        
+        if return_tensors == "pt":
+            result["input_features"] = torch.from_numpy(input_features).float()
+            if "attention_mask" in result:
+                result["attention_mask"] = torch.from_numpy(result["attention_mask"]).long()
+        
+        return result
+    
+    def to_dict(self):
+        return {
+            "feature_size": self.feature_size,
+            "sampling_rate": self.sampling_rate,
+            "hop_length": self.hop_length,
+            "n_fft": self.n_fft,
+            "chunk_length": self.chunk_length,
+            "padding_value": self.padding_value,
+            "dither": self.dither,
+        }
+    
+    def save_pretrained(self, save_directory):
+        import json, os
+        os.makedirs(save_directory, exist_ok=True)
+        with open(os.path.join(save_directory, "preprocessor_config.json"), "w") as f:
+            json.dump(self.to_dict(), f, indent=2)
+    
+    @classmethod
+    def from_pretrained(cls, save_directory):
+        import json, os
+        path = os.path.join(save_directory, "preprocessor_config.json")
+        if os.path.exists(path):
+            with open(path) as f:
+                return cls(**json.load(f))
+        return cls()
+
+
+__all__ = ["MiniWhisperFeatureExtractor"]

mini_whisper_model.py

@@ -0,0 +1,16 @@
+# mini_whisper_model.py - Backwards compatibility alias
+from configuration_mini_whisper import MiniWhisperConfig
+from modeling_mini_whisper import (
+    MiniWhisperForConditionalGeneration,
+    MiniWhisperModel,
+    MiniWhisperEncoder,
+    MiniWhisperDecoder,
+)
+
+__all__ = [
+    "MiniWhisperConfig",
+    "MiniWhisperForConditionalGeneration",
+    "MiniWhisperModel",
+    "MiniWhisperEncoder",
+    "MiniWhisperDecoder",
+]

modeling_mini_whisper.py

@@ -0,0 +1,738 @@
+# modeling_mini_whisper.py
+import math
+from collections.abc import Callable
+
+import torch
+import torch.nn as nn
+from torch.nn import functional as F
+from transformers import PreTrainedModel, GenerationMixin
+from transformers.modeling_outputs import (
+    BaseModelOutput,
+    BaseModelOutputWithPastAndCrossAttentions,
+    CausalLMOutputWithCrossAttentions,
+    Seq2SeqLMOutput,
+    Seq2SeqModelOutput,
+)
+from transformers.activations import ACT2FN
+from transformers.utils import logging
+from configuration_mini_whisper import MiniWhisperConfig
+
+logger = logging.get_logger(__name__)
+
+_HIDDEN_STATES_START_POSITION = 1
+
+def sinusoids(length: int, channels: int, max_timescale: float = 10000) -> torch.Tensor:
+    """Returns sinusoids for positional embedding"""
+    if channels % 2 != 0:
+        raise ValueError(f"Number of channels has to be divisible by 2, got {channels} channels.")
+    log_timescale_increment = math.log(max_timescale) / (channels // 2 - 1)
+    inv_timescales = torch.exp(-log_timescale_increment * torch.arange(channels // 2))
+    scaled_time = torch.arange(length).view(-1, 1) * inv_timescales.view(1, -1)
+    return torch.cat([scaled_time.sin(), scaled_time.cos()], dim=1)
+
+
+def shift_tokens_right(input_ids: torch.Tensor, pad_token_id: int, decoder_start_token_id: int):
+    """Shift input ids one token to the right."""
+    shifted_input_ids = input_ids.new_zeros(input_ids.shape)
+    shifted_input_ids[:, 1:] = input_ids[:, :-1].clone()
+    shifted_input_ids[:, 0] = decoder_start_token_id
+
+    if pad_token_id is None:
+        raise ValueError("pad_token_id has to be defined.")
+    shifted_input_ids.masked_fill_(shifted_input_ids == -100, pad_token_id)
+
+    return shifted_input_ids
+
+def _compute_mask_indices(
+    shape: tuple[int, int],
+    mask_prob: float,
+    mask_length: int,
+    attention_mask: torch.LongTensor | None = None,
+    min_masks: int = 0,
+):
+    """Computes random mask spans for SpecAugment"""
+    import numpy as np
+
+    batch_size, sequence_length = shape
+
+    if mask_length < 1:
+        raise ValueError(f"mask_length has to be bigger than 0.")
+    if mask_length > sequence_length:
+        raise ValueError(f"mask_length has to be smaller than sequence_length.")
+
+    epsilon = np.random.rand(1).item()
+
+    def compute_num_masked_span(input_length):
+        num_masked_span = int(mask_prob * input_length / mask_length + epsilon)
+        num_masked_span = max(num_masked_span, min_masks)
+        if num_masked_span * mask_length > sequence_length:
+            num_masked_span = sequence_length // mask_length
+        if input_length - (mask_length - 1) < num_masked_span:
+            num_masked_span = max(input_length - (mask_length - 1), 0)
+        return num_masked_span
+
+    input_lengths = (
+        attention_mask.detach().sum(-1).tolist()
+        if attention_mask is not None
+        else [sequence_length for _ in range(batch_size)]
+    )
+
+    spec_aug_mask = np.zeros((batch_size, sequence_length), dtype=bool)
+    spec_aug_mask_idxs = []
+    max_num_masked_span = compute_num_masked_span(sequence_length)
+
+    if max_num_masked_span == 0:
+        return spec_aug_mask
+
+    for input_length in input_lengths:
+        num_masked_span = compute_num_masked_span(input_length)
+        spec_aug_mask_idx = np.random.choice(
+            np.arange(input_length - (mask_length - 1)), num_masked_span, replace=False
+        )
+
+        if len(spec_aug_mask_idx) == 0:
+            dummy_mask_idx = sequence_length - 1
+        else:
+            dummy_mask_idx = spec_aug_mask_idx[0]
+
+        spec_aug_mask_idx = np.concatenate(
+            [spec_aug_mask_idx, np.ones(max_num_masked_span - num_masked_span, dtype=np.int32) * dummy_mask_idx]
+        )
+        spec_aug_mask_idxs.append(spec_aug_mask_idx)
+
+    spec_aug_mask_idxs = np.array(spec_aug_mask_idxs)
+    spec_aug_mask_idxs = np.broadcast_to(
+        spec_aug_mask_idxs[:, :, None], (batch_size, max_num_masked_span, mask_length)
+    )
+    spec_aug_mask_idxs = spec_aug_mask_idxs.reshape(batch_size, max_num_masked_span * mask_length)
+
+    offsets = np.arange(mask_length)[None, None, :]
+    offsets = np.broadcast_to(offsets, (batch_size, max_num_masked_span, mask_length)).reshape(
+        batch_size, max_num_masked_span * mask_length
+    )
+    spec_aug_mask_idxs = spec_aug_mask_idxs + offsets
+
+    if spec_aug_mask_idxs.max() > sequence_length - 1:
+        spec_aug_mask_idxs[spec_aug_mask_idxs > sequence_length - 1] = sequence_length - 1
+
+    np.put_along_axis(spec_aug_mask, spec_aug_mask_idxs, 1, -1)
+
+    return spec_aug_mask
+
+class WhisperPositionalEmbedding(nn.Module):
+    def __init__(self, num_positions: int, embedding_dim: int, padding_idx: int | None = None):
+        super().__init__()
+        self.num_positions = num_positions
+        self.embedding_dim = embedding_dim
+        self.weight = nn.Parameter(torch.zeros(num_positions, embedding_dim))
+        self._init_weights()
+
+    def _init_weights(self):
+        with torch.no_grad():
+            self.weight.copy_(sinusoids(self.num_positions, self.embedding_dim))
+
+    def forward(self, input_ids, past_key_values_length=0, position_ids=None):
+        if position_ids is None:
+            return self.weight[past_key_values_length : past_key_values_length + input_ids.shape[1]]
+        else:
+            return self.weight[position_ids]
+
+class MiniWhisperAttention(nn.Module):
+    """Multi-headed attention from 'Attention Is All You Need' paper"""
+
+    def __init__(
+        self,
+        embed_dim: int,
+        num_heads: int,
+        dropout: float = 0.0,
+        is_decoder: bool = False,
+        bias: bool = True,
+        is_causal: bool = False,
+        layer_idx: int | None = None,
+        config: MiniWhisperConfig | None = None,
+    ):
+        super().__init__()
+        self.embed_dim = embed_dim
+        self.num_heads = num_heads
+        self.dropout = dropout
+        self.head_dim = embed_dim // num_heads
+        self.config = config
+
+        if (self.head_dim * num_heads) != self.embed_dim:
+            raise ValueError(f"embed_dim must be divisible by num_heads")
+        self.scaling = self.head_dim ** -0.5
+        self.is_decoder = is_decoder
+        self.is_causal = is_causal
+        self.layer_idx = layer_idx
+
+        self.k_proj = nn.Linear(embed_dim, embed_dim, bias=False)
+        self.v_proj = nn.Linear(embed_dim, embed_dim, bias=bias)
+        self.q_proj = nn.Linear(embed_dim, embed_dim, bias=bias)
+        self.out_proj = nn.Linear(embed_dim, embed_dim, bias=bias)
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        key_value_states: torch.Tensor | None = None,
+        past_key_value: tuple[torch.Tensor, torch.Tensor] | None = None,
+        attention_mask: torch.Tensor | None = None,
+        output_attentions: bool = False,
+    ):
+        """Input shape: Batch x Time x Channel"""
+        is_cross_attention = key_value_states is not None
+        input_shape = hidden_states.shape[:-1]
+        hidden_shape = (*input_shape, -1, self.head_dim)
+
+        query_states = (self.q_proj(hidden_states) * self.scaling).view(hidden_shape).transpose(1, 2).contiguous()
+
+        if past_key_value is not None and not is_cross_attention:
+            key_states = past_key_value[0]
+            value_states = past_key_value[1]
+        else:
+            current_states = key_value_states if key_value_states is not None else hidden_states
+            kv_shape = (input_shape[0], -1, self.num_heads, self.head_dim)
+            key_states = self.k_proj(current_states).view(kv_shape).transpose(1, 2).contiguous()
+            value_states = self.v_proj(current_states).view(kv_shape).transpose(1, 2).contiguous()
+
+        if past_key_value is not None and is_cross_attention:
+            key_states = torch.cat([past_key_value[0], key_states], dim=2)
+            value_states = torch.cat([past_key_value[1], value_states], dim=2)
+
+        attn_weights = torch.matmul(query_states, key_states.transpose(-2, -1)) * 1.0
+        if attention_mask is not None:
+            attn_weights = attn_weights + attention_mask
+
+        attn_weights = F.softmax(attn_weights, dim=-1)
+        attn_weights = F.dropout(attn_weights, p=self.dropout, training=self.training)
+
+        attn_output = torch.matmul(attn_weights, value_states)
+        attn_output = attn_output.reshape(*input_shape, -1).contiguous()
+        attn_output = self.out_proj(attn_output)
+
+        return attn_output, None
+
+class MiniWhisperEncoderLayer(nn.Module):
+    def __init__(self, config: MiniWhisperConfig):
+        super().__init__()
+        self.embed_dim = config.d_model
+
+        self.self_attn = MiniWhisperAttention(
+            embed_dim=self.embed_dim,
+            num_heads=config.encoder_attention_heads,
+            dropout=config.attention_dropout,
+            is_decoder=False,
+            config=config,
+        )
+        self.self_attn_layer_norm = nn.LayerNorm(self.embed_dim)
+        self.dropout = config.dropout
+        self.activation_fn = ACT2FN[config.activation_function]
+        self.activation_dropout = config.activation_dropout
+        self.fc1 = nn.Linear(self.embed_dim, config.encoder_ffn_dim)
+        self.fc2 = nn.Linear(config.encoder_ffn_dim, self.embed_dim)
+        self.final_layer_norm = nn.LayerNorm(self.embed_dim)
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        attention_mask: torch.Tensor,
+    ):
+        residual = hidden_states
+        hidden_states = self.self_attn_layer_norm(hidden_states)
+        hidden_states, _ = self.self_attn(
+            hidden_states=hidden_states,
+            attention_mask=attention_mask,
+        )
+        hidden_states = F.dropout(hidden_states, p=self.dropout, training=self.training)
+        hidden_states = residual + hidden_states
+
+        residual = hidden_states
+        hidden_states = self.final_layer_norm(hidden_states)
+        hidden_states = self.activation_fn(self.fc1(hidden_states))
+        hidden_states = F.dropout(hidden_states, p=self.activation_dropout, training=self.training)
+        hidden_states = self.fc2(hidden_states)
+        hidden_states = F.dropout(hidden_states, p=self.dropout, training=self.training)
+        hidden_states = residual + hidden_states
+
+        if hidden_states.dtype == torch.float16:
+            clamp_value = torch.finfo(hidden_states.dtype).max - 1000
+            hidden_states = torch.clamp(hidden_states, min=-clamp_value, max=clamp_value)
+
+        return hidden_states
+
+class MiniWhisperDecoderLayer(nn.Module):
+    def __init__(self, config: MiniWhisperConfig, layer_idx: int | None = None):
+        super().__init__()
+        self.embed_dim = config.d_model
+
+        self.self_attn = MiniWhisperAttention(
+            embed_dim=self.embed_dim,
+            num_heads=config.decoder_attention_heads,
+            dropout=config.attention_dropout,
+            is_decoder=True,
+            is_causal=True,
+            layer_idx=layer_idx,
+            config=config,
+        )
+        self.dropout = config.dropout
+        self.activation_fn = ACT2FN[config.activation_function]
+        self.activation_dropout = config.activation_dropout
+
+        self.self_attn_layer_norm = nn.LayerNorm(self.embed_dim)
+        self.encoder_attn = MiniWhisperAttention(
+            self.embed_dim,
+            config.decoder_attention_heads,
+            dropout=config.attention_dropout,
+            is_decoder=True,
+            layer_idx=layer_idx,
+            config=config,
+        )
+        self.encoder_attn_layer_norm = nn.LayerNorm(self.embed_dim)
+        self.fc1 = nn.Linear(self.embed_dim, config.decoder_ffn_dim)
+        self.fc2 = nn.Linear(config.decoder_ffn_dim, self.embed_dim)
+        self.final_layer_norm = nn.LayerNorm(self.embed_dim)
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        attention_mask: torch.Tensor | None = None,
+        encoder_hidden_states: torch.Tensor | None = None,
+        encoder_attention_mask: torch.Tensor | None = None,
+        past_key_value: tuple[tuple[torch.Tensor, torch.Tensor]] | None = None,
+        use_cache: bool | None = True,
+    ):
+        residual = hidden_states
+        hidden_states = self.self_attn_layer_norm(hidden_states)
+
+        self_attn_past = past_key_value[0] if past_key_value is not None else None
+        hidden_states, _ = self.self_attn(
+            hidden_states,
+            past_key_value=self_attn_past,
+            attention_mask=attention_mask,
+        )
+        hidden_states = F.dropout(hidden_states, p=self.dropout, training=self.training)
+        hidden_states = residual + hidden_states
+
+        if encoder_hidden_states is not None:
+            residual = hidden_states
+            hidden_states = self.encoder_attn_layer_norm(hidden_states)
+            cross_attn_past = past_key_value[1] if past_key_value is not None else None
+            hidden_states, _ = self.encoder_attn(
+                hidden_states,
+                key_value_states=encoder_hidden_states,
+                attention_mask=encoder_attention_mask,
+                past_key_value=cross_attn_past,
+            )
+            hidden_states = F.dropout(hidden_states, p=self.dropout, training=self.training)
+            hidden_states = residual + hidden_states
+
+        residual = hidden_states
+        hidden_states = self.final_layer_norm(hidden_states)
+        hidden_states = self.activation_fn(self.fc1(hidden_states))
+        hidden_states = F.dropout(hidden_states, p=self.activation_dropout, training=self.training)
+        hidden_states = self.fc2(hidden_states)
+        hidden_states = F.dropout(hidden_states, p=self.dropout, training=self.training)
+        hidden_states = residual + hidden_states
+
+        return hidden_states, None
+
+class MiniWhisperEncoder(PreTrainedModel):
+    config_class = MiniWhisperConfig
+    main_input_name = "input_features"
+
+    def __init__(self, config: MiniWhisperConfig):
+        super().__init__(config)
+        self.dropout = config.dropout
+        self.layerdrop = config.encoder_layerdrop
+
+        embed_dim = config.d_model
+        self.num_mel_bins = config.num_mel_bins
+        self.padding_idx = config.pad_token_id
+        self.max_source_positions = config.max_source_positions
+        self.embed_scale = math.sqrt(embed_dim) if config.scale_embedding else 1.0
+
+        self.conv1 = nn.Conv1d(self.num_mel_bins, embed_dim, kernel_size=3, padding=1)
+        self.conv2 = nn.Conv1d(embed_dim, embed_dim, kernel_size=3, stride=2, padding=1)
+
+        self.embed_positions = nn.Embedding(self.max_source_positions, embed_dim)
+        self.embed_positions.requires_grad_(False)
+
+        self.layers = nn.ModuleList([MiniWhisperEncoderLayer(config) for _ in range(config.encoder_layers)])
+        self.layer_norm = nn.LayerNorm(config.d_model)
+
+        self.gradient_checkpointing = False
+        self.post_init()
+
+    def _init_weights(self, module):
+        if isinstance(module, MiniWhisperEncoder):
+            with torch.no_grad():
+                module.embed_positions.weight.copy_(sinusoids(*module.embed_positions.weight.shape))
+
+    def forward(
+        self,
+        input_features,
+        attention_mask=None,
+        **kwargs,
+    ):
+        x = F.gelu(self.conv1(input_features))
+        x = F.gelu(self.conv2(x))
+
+        x = x.permute(0, 2, 1)
+
+        seq_len = x.shape[1]
+        positions = torch.arange(seq_len, device=x.device)
+        hidden_states = x + self.embed_positions(positions)
+        hidden_states = F.dropout(hidden_states, p=self.dropout, training=self.training)
+
+        for idx, encoder_layer in enumerate(self.layers):
+            if self.training:
+                dropout_probability = torch.rand([])
+                if dropout_probability < self.layerdrop:
+                    continue
+
+            hidden_states = encoder_layer(hidden_states, None)
+
+        hidden_states = self.layer_norm(hidden_states)
+
+        return BaseModelOutput(last_hidden_state=hidden_states)
+
+    def _get_feat_extract_output_lengths(self, input_lengths: torch.LongTensor):
+        input_lengths = (input_lengths - 1) // 2 + 1
+        return input_lengths
+
+class MiniWhisperDecoder(PreTrainedModel):
+    config_class = MiniWhisperConfig
+    main_input_name = "input_ids"
+
+    def __init__(self, config: MiniWhisperConfig):
+        super().__init__(config)
+        self.dropout = config.dropout
+        self.layerdrop = config.decoder_layerdrop
+        self.padding_idx = config.pad_token_id
+        self.max_target_positions = config.max_target_positions
+        self.max_source_positions = config.max_source_positions
+        self.embed_scale = math.sqrt(config.d_model) if config.scale_embedding else 1.0
+
+        self.embed_tokens = nn.Embedding(config.vocab_size, config.d_model, self.padding_idx)
+        self.embed_positions = WhisperPositionalEmbedding(self.max_target_positions, config.d_model)
+
+        self.layers = nn.ModuleList(
+            [MiniWhisperDecoderLayer(config, layer_idx) for layer_idx in range(config.decoder_layers)]
+        )
+
+        self.layer_norm = nn.LayerNorm(config.d_model)
+
+        self.gradient_checkpointing = False
+        self.post_init()
+
+    def forward(
+        self,
+        input_ids=None,
+        attention_mask=None,
+        encoder_hidden_states=None,
+        encoder_attention_mask=None,
+        past_key_values=None,
+        inputs_embeds=None,
+        position_ids=None,
+        use_cache=None,
+        **kwargs,
+    ):
+        if input_ids is not None and inputs_embeds is not None:
+            raise ValueError("You cannot specify both decoder_input_ids and decoder_inputs_embeds at the same time")
+        elif input_ids is not None:
+            inputs_embeds = self.embed_tokens(input_ids)
+        elif inputs_embeds is None:
+            raise ValueError("You have to specify either decoder_input_ids or decoder_inputs_embeds")
+
+        past_key_values_length = 0
+        if past_key_values is not None and len(past_key_values) > 0:
+            if past_key_values[0] is not None and past_key_values[0][0] is not None:
+                past_key_values_length = past_key_values[0][0].shape[2]
+
+        if position_ids is None:
+            position_ids = torch.arange(inputs_embeds.shape[1], device=inputs_embeds.device) + past_key_values_length
+            position_ids = position_ids.unsqueeze(0).repeat(inputs_embeds.shape[0], 1)
+
+        if input_ids is not None:
+            positions = self.embed_positions(input_ids, past_key_values_length=past_key_values_length, position_ids=position_ids)
+        else:
+            positions = self.embed_positions(inputs_embeds, past_key_values_length=past_key_values_length, position_ids=position_ids)
+
+        hidden_states = inputs_embeds + positions
+        hidden_states = F.dropout(hidden_states, p=self.dropout, training=self.training)
+
+        seq_len = inputs_embeds.shape[1]
+        causal_mask = torch.triu(
+            torch.ones(seq_len, seq_len, device=hidden_states.device) * float('-inf'),
+            diagonal=1
+        ).unsqueeze(0).unsqueeze(0)
+
+        if attention_mask is not None:
+            padding_mask = (1 - attention_mask) * float('-inf')
+            padding_mask = padding_mask.unsqueeze(1).unsqueeze(2)
+            attention_mask = causal_mask + padding_mask
+        else:
+            attention_mask = causal_mask
+
+        for idx, decoder_layer in enumerate(self.layers):
+            if self.training:
+                dropout_probability = torch.rand([])
+                if dropout_probability < self.layerdrop:
+                    continue
+
+            layer_past = past_key_values[idx] if past_key_values is not None and idx < len(past_key_values) else None
+            hidden_states, _ = decoder_layer(
+                hidden_states,
+                attention_mask=attention_mask,
+                encoder_hidden_states=encoder_hidden_states,
+                encoder_attention_mask=encoder_attention_mask,
+                past_key_value=layer_past,
+                use_cache=use_cache,
+            )
+
+        hidden_states = self.layer_norm(hidden_states)
+
+        return BaseModelOutputWithPastAndCrossAttentions(
+            last_hidden_state=hidden_states,
+            past_key_values=None,
+        )
+
+class MiniWhisperModel(PreTrainedModel):
+    config_class = MiniWhisperConfig
+
+    def __init__(self, config: MiniWhisperConfig):
+        super().__init__(config)
+
+        self.encoder = MiniWhisperEncoder(config)
+        self.decoder = MiniWhisperDecoder(config)
+        self.post_init()
+
+    def get_input_embeddings(self):
+        return self.decoder.embed_tokens
+
+    def set_input_embeddings(self, value):
+        self.decoder.embed_tokens = value
+
+    def _mask_input_features(
+        self,
+        input_features: torch.FloatTensor,
+        attention_mask: torch.LongTensor | None = None,
+    ):
+        """Apply SpecAugment"""
+        import numpy as np
+
+        if not getattr(self.config, "apply_spec_augment", True):
+            return input_features
+
+        batch_size, hidden_size, sequence_length = input_features.size()
+
+        if self.config.mask_time_prob > 0 and self.training:
+            mask_time_indices = _compute_mask_indices(
+                (batch_size, sequence_length),
+                mask_prob=self.config.mask_time_prob,
+                mask_length=self.config.mask_time_length,
+                attention_mask=attention_mask,
+                min_masks=self.config.mask_time_min_masks,
+            )
+            mask_time_indices = torch.tensor(mask_time_indices, device=input_features.device, dtype=torch.bool)
+            mask_time_indices = mask_time_indices[:, None].expand(-1, hidden_size, -1)
+            input_features[mask_time_indices] = 0
+
+        if self.config.mask_feature_prob > 0 and self.training:
+            mask_feature_indices = _compute_mask_indices(
+                (batch_size, hidden_size),
+                mask_prob=self.config.mask_feature_prob,
+                mask_length=self.config.mask_feature_length,
+                min_masks=self.config.mask_feature_min_masks,
+            )
+            mask_feature_indices = torch.tensor(mask_feature_indices, device=input_features.device, dtype=torch.bool)
+            input_features[mask_feature_indices] = 0
+
+        return input_features
+
+    def forward(
+        self,
+        input_features: torch.FloatTensor | None = None,
+        attention_mask: torch.LongTensor | None = None,
+        decoder_input_ids: torch.LongTensor | None = None,
+        decoder_attention_mask: torch.LongTensor | None = None,
+        encoder_outputs: tuple[torch.FloatTensor] | None = None,
+        past_key_values: tuple[tuple[torch.FloatTensor]] | None = None,
+        decoder_inputs_embeds: tuple[torch.FloatTensor] | None = None,
+        decoder_position_ids: tuple[torch.LongTensor] | None = None,
+        use_cache: bool | None = None,
+        **kwargs,
+    ):
+        if encoder_outputs is None:
+            input_features = self._mask_input_features(input_features, attention_mask=attention_mask)
+
+            encoder_outputs = self.encoder(input_features, **kwargs)
+        elif not isinstance(encoder_outputs, BaseModelOutput):
+            encoder_outputs = BaseModelOutput(
+                last_hidden_state=encoder_outputs[0],
+                hidden_states=encoder_outputs[1] if len(encoder_outputs) > 1 else None,
+                attentions=encoder_outputs[2] if len(encoder_outputs) > 2 else None,
+            )
+
+        decoder_outputs = self.decoder(
+            input_ids=decoder_input_ids,
+            attention_mask=decoder_attention_mask,
+            encoder_hidden_states=encoder_outputs.last_hidden_state,
+            past_key_values=past_key_values,
+            inputs_embeds=decoder_inputs_embeds,
+            position_ids=decoder_position_ids,
+            use_cache=use_cache,
+            **kwargs,
+        )
+
+        return Seq2SeqModelOutput(
+            last_hidden_state=decoder_outputs.last_hidden_state,
+            past_key_values=decoder_outputs.past_key_values,
+            decoder_hidden_states=decoder_outputs.hidden_states,
+            decoder_attentions=decoder_outputs.attentions,
+            cross_attentions=decoder_outputs.cross_attentions,
+            encoder_last_hidden_state=encoder_outputs.last_hidden_state,
+            encoder_hidden_states=encoder_outputs.hidden_states,
+            encoder_attentions=encoder_outputs.attentions,
+        )
+
+class MiniWhisperForConditionalGeneration(PreTrainedModel, GenerationMixin):
+    config_class = MiniWhisperConfig
+    base_model_prefix = "model"
+    _tied_weights_keys = ["proj_out.weight"]
+
+    def __init__(self, config: MiniWhisperConfig):
+        super().__init__(config)
+        self.model = MiniWhisperModel(config)
+        self.proj_out = nn.Linear(config.d_model, config.vocab_size, bias=False)
+        self.max_target_positions = config.max_target_positions
+
+        self.post_init()
+
+    def get_output_embeddings(self):
+        return self.proj_out
+
+    def set_output_embeddings(self, new_embeddings):
+        self.proj_out = new_embeddings
+
+    def get_input_embeddings(self):
+        return self.model.get_input_embeddings()
+
+    def forward(
+        self,
+        input_features: torch.FloatTensor | None = None,
+        attention_mask: torch.LongTensor | None = None,
+        decoder_input_ids: torch.LongTensor | None = None,
+        decoder_attention_mask: torch.LongTensor | None = None,
+        encoder_outputs: tuple[torch.FloatTensor] | None = None,
+        past_key_values: tuple[tuple[torch.FloatTensor]] | None = None,
+        decoder_inputs_embeds: tuple[torch.FloatTensor] | None = None,
+        decoder_position_ids: tuple[torch.LongTensor] | None = None,
+        labels: torch.LongTensor | None = None,
+        use_cache: bool | None = None,
+        **kwargs,
+    ):
+        if labels is not None:
+            if labels.shape[1] > self.max_target_positions:
+                raise ValueError(f"Labels' sequence length {labels.shape[1]} cannot exceed max_target_positions")
+            if decoder_input_ids is None and decoder_inputs_embeds is None:
+                decoder_input_ids = shift_tokens_right(
+                    labels, self.config.pad_token_id, self.config.decoder_start_token_id
+                )
+
+        outputs = self.model(
+            input_features,
+            attention_mask=attention_mask,
+            decoder_input_ids=decoder_input_ids,
+            encoder_outputs=encoder_outputs,
+            decoder_attention_mask=decoder_attention_mask,
+            past_key_values=past_key_values,
+            decoder_inputs_embeds=decoder_inputs_embeds,
+            decoder_position_ids=decoder_position_ids,
+            use_cache=use_cache,
+            **kwargs,
+        )
+
+        lm_logits = self.proj_out(outputs.last_hidden_state)
+
+        loss = None
+        if labels is not None:
+            loss_fct = nn.CrossEntropyLoss()
+            labels = labels.to(lm_logits.device)
+            loss = loss_fct(lm_logits.view(-1, self.config.vocab_size), labels.view(-1))
+
+        return Seq2SeqLMOutput(
+            loss=loss,
+            logits=lm_logits,
+            past_key_values=outputs.past_key_values,
+            decoder_hidden_states=outputs.decoder_hidden_states,
+            decoder_attentions=outputs.decoder_attentions,
+            cross_attentions=outputs.cross_attentions,
+            encoder_last_hidden_state=outputs.encoder_last_hidden_state,
+            encoder_hidden_states=outputs.encoder_hidden_states,
+            encoder_attentions=outputs.encoder_attentions,
+        )
+
+    def generate(
+        self,
+        input_features,
+        generation_config=None,
+        max_new_tokens=50,
+        **kwargs,
+    ):
+        self.eval()
+        batch_size = input_features.shape[0]
+        device = input_features.device
+
+        decoder_input_ids = torch.full(
+            (batch_size, 1), self.config.decoder_start_token_id, dtype=torch.long, device=device
+        )
+
+        encoder_outputs = self.model.encoder(input_features)
+
+        for _ in range(max_new_tokens):
+            outputs = self.forward(
+                decoder_input_ids=decoder_input_ids,
+                encoder_outputs=encoder_outputs,
+                use_cache=True,
+            )
+
+            next_token_logits = outputs.logits[:, -1, :]
+            next_token = torch.argmax(next_token_logits, dim=-1, keepdim=True)
+            decoder_input_ids = torch.cat([decoder_input_ids, next_token], dim=1)
+
+            if (next_token == self.config.eos_token_id).all():
+                break
+
+        return decoder_input_ids
+
+    def prepare_inputs_for_generation(
+        self,
+        decoder_input_ids,
+        past_key_values=None,
+        encoder_outputs=None,
+        attention_mask=None,
+        use_cache=None,
+        **kwargs,
+    ):
+        return {
+            "decoder_input_ids": decoder_input_ids,
+            "past_key_values": past_key_values,
+            "encoder_outputs": encoder_outputs,
+            "attention_mask": attention_mask,
+            "use_cache": use_cache,
+        }
+
+    def get_encoder(self):
+        return self.model.get_encoder()
+
+__all__ = [
+    "MiniWhisperConfig",
+    "MiniWhisperModel",
+    "MiniWhisperForConditionalGeneration",
+    "MiniWhisperEncoder",
+    "MiniWhisperDecoder",
+    "MiniWhisperEncoderLayer",
+    "MiniWhisperDecoderLayer",
+    "MiniWhisperAttention",
+    "WhisperPositionalEmbedding",
+]

processing_mini_whisper.py

@@ -0,0 +1,73 @@
+# processing_mini_whisper.py
+from typing import Union, List
+from transformers import ProcessorMixin
+from .feature_extraction_mini_whisper import MiniWhisperFeatureExtractor
+from .tokenization_mini_whisper import MiniWhisperTokenizer
+
+class MiniWhisperProcessor(ProcessorMixin):
+    """Combines feature extractor and tokenizer."""
+    
+    attributes = ["feature_extractor", "tokenizer"]
+    feature_extractor_class = "MiniWhisperFeatureExtractor"
+    tokenizer_class = "MiniWhisperTokenizer"
+
+    def __init__(self, feature_extractor, tokenizer):
+        super().__init__(feature_extractor, tokenizer)
+
+    def __call__(
+        self,
+        audio=None,
+        text=None,
+        sampling_rate=None,
+        return_tensors="pt",
+        **kwargs,
+    ):
+        if audio is not None and text is not None:
+            raise ValueError("Cannot specify both audio and text")
+
+        if audio is not None:
+            inputs = self.feature_extractor(
+                audio,
+                sampling_rate=sampling_rate,
+                return_tensors=return_tensors,
+                **kwargs,
+            )
+            return inputs
+
+        if text is not None:
+            encodings = self.tokenizer(
+                text,
+                return_tensors=return_tensors,
+                **kwargs,
+            )
+            return encodings
+
+    def get_decoder_prompt_ids(self, task=None, language=None, no_timestamps=True):
+        return self.tokenizer.get_decoder_prompt_ids(task=task, language=language, no_timestamps=no_timestamps)
+
+    def batch_decode(self, sequences, **kwargs):
+        return self.tokenizer.batch_decode(sequences, **kwargs)
+
+    def decode(self, token_ids, **kwargs):
+        return self.tokenizer.decode(token_ids, **kwargs)
+
+    def save_pretrained(self, save_directory):
+        self.feature_extractor.save_pretrained(save_directory)
+        self.tokenizer.save_pretrained(save_directory)
+        
+        import json
+        import os
+        config = {
+            "feature_extractor_type": "MiniWhisperFeatureExtractor",
+            "tokenizer_type": "MiniWhisperTokenizer",
+        }
+        with open(os.path.join(save_directory, "preprocessor_config.json"), "w") as f:
+            json.dump(config, f, indent=2)
+
+    @classmethod
+    def from_pretrained(cls, pretrained_model_name_or_path, **kwargs):
+        feature_extractor = MiniWhisperFeatureExtractor.from_pretrained(pretrained_model_name_or_path)
+        tokenizer = MiniWhisperTokenizer.from_pretrained(pretrained_model_name_or_path, **kwargs)
+        return cls(feature_extractor=feature_extractor, tokenizer=tokenizer)
+
+__all__ = ["MiniWhisperProcessor"]

tokenization_mini_whisper.py

@@ -0,0 +1,396 @@
+# tokenization_mini_whisper.py
+"""Tokenization classes for MiniWhisper"""
+
+import json
+import os
+import re
+from functools import lru_cache
+from typing import List, Optional, Tuple, Union
+
+import numpy as np
+from transformers import PreTrainedTokenizer
+from transformers.tokenization_utils_base import AddedToken
+from transformers.utils import logging
+
+logger = logging.get_logger(__name__)
+
+VOCAB_FILES_NAMES = {
+    "vocab_file": "vocab.json",
+    "tokenizer_file": "tokenizer.json",
+    "merges_file": "merges.txt",
+    "normalizer_file": "normalizer.json",
+}
+
+LANGUAGES = {
+    "en": "english",
+    "zh": "chinese",
+    "de": "german",
+    "es": "spanish",
+    "ru": "russian",
+    "ko": "korean",
+    "fr": "french",
+    "ja": "japanese",
+    "pt": "portuguese",
+    "tr": "turkish",
+    "pl": "polish",
+    "ca": "catalan",
+    "nl": "dutch",
+    "ar": "arabic",
+    "sv": "swedish",
+    "it": "italian",
+    "id": "indonesian",
+    "hi": "hindi",
+    "fi": "finnish",
+    "vi": "vietnamese",
+    "he": "hebrew",
+    "uk": "ukrainian",
+    "el": "greek",
+    "ms": "malay",
+    "cs": "czech",
+    "ro": "romanian",
+    "da": "danish",
+    "hu": "hungarian",
+    "ta": "tamil",
+    "no": "norwegian",
+    "th": "thai",
+    "ur": "urdu",
+    "hr": "croatian",
+    "bg": "bulgarian",
+    "lt": "lithuanian",
+    "la": "latin",
+    "mi": "maori",
+    "ml": "malayalam",
+    "cy": "welsh",
+    "sk": "slovak",
+    "te": "telugu",
+    "fa": "persian",
+    "lv": "latvian",
+    "bn": "bengali",
+    "sr": "serbian",
+    "az": "azerbaijani",
+    "sl": "slovenian",
+    "kn": "kannada",
+    "et": "estonian",
+    "mk": "macedonian",
+    "br": "breton",
+    "eu": "basque",
+    "is": "icelandic",
+    "hy": "armenian",
+    "ne": "nepali",
+    "mn": "mongolian",
+    "bs": "bosnian",
+    "kk": "kazakh",
+    "sq": "albanian",
+    "sw": "swahili",
+    "gl": "galician",
+    "mr": "marathi",
+    "pa": "punjabi",
+    "si": "sinhala",
+    "km": "khmer",
+    "sn": "shona",
+    "yo": "yoruba",
+    "so": "somali",
+    "af": "afrikaans",
+    "oc": "occitan",
+    "ka": "georgian",
+    "be": "belarusian",
+    "tg": "tajik",
+    "sd": "sindhi",
+    "gu": "gujarati",
+    "am": "amharic",
+    "yi": "yiddish",
+    "lo": "lao",
+    "uz": "uzbek",
+    "fo": "faroese",
+    "ht": "haitian creole",
+    "ps": "pashto",
+    "tk": "turkmen",
+    "nn": "nynorsk",
+    "mt": "maltese",
+    "sa": "sanskrit",
+    "lb": "luxembourgish",
+    "my": "myanmar",
+    "bo": "tibetan",
+    "tl": "tagalog",
+    "mg": "malagasy",
+    "as": "assamese",
+    "tt": "tatar",
+    "haw": "hawaiian",
+    "ln": "lingala",
+    "ha": "hausa",
+    "ba": "bashkir",
+    "jw": "javanese",
+    "su": "sundanese",
+    "yue": "cantonese",
+}
+
+TO_LANGUAGE_CODE = {
+    **{language: code for code, language in LANGUAGES.items()},
+    "burmese": "my",
+    "valencian": "ca",
+    "flemish": "nl",
+    "haitian": "ht",
+    "letzeburgesch": "lb",
+    "pushto": "ps",
+    "panjabi": "pa",
+    "moldavian": "ro",
+    "moldovan": "ro",
+    "sinhalese": "si",
+    "castilian": "es",
+    "mandarin": "zh",
+}
+
+TASK_IDS = ["translate", "transcribe"]
+
+class MiniWhisperTokenizer(PreTrainedTokenizer):
+    """
+    Construct a MiniWhisper tokenizer.
+    
+    For actual use, load pretrained:
+        tokenizer = WhisperTokenizer.from_pretrained("openai/whisper-tiny")
+    
+    Or provide path to trained vocab/merges files:
+        tokenizer = MiniWhisperTokenizer(vocab_file="vocab.json", merges_file="merges.txt")
+    """
+    
+    vocab_files_names = VOCAB_FILES_NAMES
+    model_input_names = ["input_ids", "attention_mask"]
+    
+    def __init__(
+        self,
+        vocab_file=None,
+        merges_file=None,
+        normalizer_file=None,
+        tokenizer_file=None,
+        unk_token="<|endoftext|>",
+        bos_token="<|endoftext|>",
+        eos_token="<|endoftext|>",
+        pad_token="<|endoftext|>",
+        add_prefix_space=False,
+        language=None,
+        task=None,
+        predict_timestamps=False,
+        **kwargs,
+    ):
+        bos_token = AddedToken(bos_token, lstrip=False, rstrip=False, normalized=False, special=True) if isinstance(bos_token, str) else bos_token
+        eos_token = AddedToken(eos_token, lstrip=False, rstrip=False, normalized=False, special=True) if isinstance(eos_token, str) else eos_token
+        unk_token = AddedToken(unk_token, lstrip=False, rstrip=False, normalized=False, special=True) if isinstance(unk_token, str) else unk_token
+        pad_token = AddedToken(pad_token, lstrip=False, rstrip=False, normalized=False, special=True) if isinstance(pad_token, str) else pad_token
+
+        self._vocab = {}
+        self._merges = []
+
+        if vocab_file is not None and os.path.isfile(vocab_file):
+            with open(vocab_file, encoding="utf-8") as f:
+                self._vocab = json.load(f)
+        
+        if merges_file is not None and os.path.isfile(merges_file):
+            with open(merges_file, encoding="utf-8") as f:
+                lines = f.read().splitlines()
+                if lines and lines[0] == "#version: 0.2":
+                    self._merges = lines[1:]
+                else:
+                    self._merges = lines
+
+        self.english_spelling_normalizer = None
+        if normalizer_file is not None and os.path.isfile(normalizer_file):
+            with open(normalizer_file, encoding="utf-8") as f:
+                self.english_spelling_normalizer = json.load(f)
+
+        self.add_prefix_space = add_prefix_space
+        self.language = language
+        self.task = task
+        self.predict_timestamps = predict_timestamps
+        self.timestamp_pat = re.compile(r"<\|(\d+\.\d+)\|>")
+
+        super().__init__(
+            unk_token=unk_token,
+            bos_token=bos_token,
+            eos_token=eos_token,
+            pad_token=pad_token,
+            add_prefix_space=add_prefix_space,
+            **kwargs,
+        )
+
+        self.set_prefix_tokens()
+
+    @property
+    def vocab_size(self) -> int:
+        if self._vocab:
+            return len(self._vocab)
+        return 51865
+
+    def get_vocab(self):
+        if self._vocab:
+            return dict(self._vocab, **self.added_tokens_encoder)
+        return {self.unk_token: 0, self.bos_token: 1, self.eos_token: 2, self.pad_token: 3}
+
+    def _tokenize(self, text, **kwargs):
+        """Tokenize a string using BPE"""
+        if not self._vocab or not self._merges:
+            return list(text)
+        
+        tokens = []
+        i = 0
+        text = text.lower()
+        
+        while i < len(text):
+            if i < len(text) - 1:
+                bigram = text[i:i+2]
+                if bigram in self._merges:
+                    tokens.append(bigram)
+                    i += 2
+                    continue
+            tokens.append(text[i])
+            i += 1
+        
+        return tokens
+
+    def _convert_token_to_id(self, token):
+        if self._vocab:
+            return self._vocab.get(token, self._vocab.get(self.unk_token))
+        index = len(self._vocab) if self._vocab else 0
+        if token == self.unk_token:
+            return index
+        if token == self.bos_token:
+            return index + 1
+        if token == self.eos_token:
+            return index + 2
+        if token == self.pad_token:
+            return index + 3
+        return index
+
+    def _convert_id_to_token(self, index):
+        if hasattr(self, 'ids_to_tokens') and self.ids_to_tokens:
+            return self.ids_to_tokens.get(index, self.unk_token)
+        return super()._convert_id_to_token(index)
+
+    def convert_tokens_to_string(self, tokens):
+        return "".join(tokens)
+
+    @property
+    def prefix_tokens(self) -> List[int]:
+        bos_token_id = self.convert_tokens_to_ids("<|startoftranscript|>")
+        transcribe_token_id = self.convert_tokens_to_ids("<|transcribe|>")
+        translate_token_id = self.convert_tokens_to_ids("<|translate|>")
+        notimestamps_token_id = self.convert_tokens_to_ids("<|notimestamps|>")
+        langs = tuple(LANGUAGES.keys())
+
+        if self.language is not None:
+            self.language = self.language.lower()
+            if self.language in TO_LANGUAGE_CODE:
+                language_id = TO_LANGUAGE_CODE[self.language]
+            elif self.language in TO_LANGUAGE_CODE.values():
+                language_id = self.language
+            else:
+                language_id = "en"
+
+        bos_sequence = [bos_token_id]
+        
+        if self.language is not None:
+            bos_sequence.append(bos_token_id + 1 + langs.index(language_id))
+        
+        if self.task is not None:
+            if self.task == "transcribe":
+                bos_sequence.append(transcribe_token_id)
+            elif self.task == "translate":
+                bos_sequence.append(translate_token_id)
+        
+        if not self.predict_timestamps:
+            bos_sequence.append(notimestamps_token_id)
+        
+        return bos_sequence
+
+    def set_prefix_tokens(self, language=None, task=None, predict_timestamps=None):
+        if language is not None:
+            self.language = language
+        if task is not None:
+            self.task = task
+        if predict_timestamps is not None:
+            self.predict_timestamps = predict_timestamps
+
+    def build_inputs_with_special_tokens(self, token_ids_0, token_ids_1=None) -> List[int]:
+        prefix = self.prefix_tokens
+        eos = [self.eos_token_id]
+
+        if token_ids_1 is None:
+            return prefix + token_ids_0 + eos
+        return prefix + token_ids_0 + token_ids_1 + eos
+
+    def get_special_tokens_mask(
+        self, token_ids_0: List[int], token_ids_1: Optional[List[int]] = None, 
+        already_has_special_tokens: bool = False
+    ) -> List[int]:
+        if already_has_special_tokens:
+            return super().get_special_tokens_mask(
+                token_ids_0=token_ids_0, token_ids_1=token_ids_1, 
+                already_has_special_tokens=True
+            )
+
+        prefix_ones = [1] * len(self.prefix_tokens)
+        suffix_ones = [1]
+
+        if token_ids_1 is None:
+            return prefix_ones + ([0] * len(token_ids_0)) + suffix_ones
+        return prefix_ones + ([0] * len(token_ids_0)) + ([0] * len(token_ids_1)) + suffix_ones
+
+    def get_decoder_prompt_ids(self, task=None, language=None, no_timestamps=True):
+        self.set_prefix_tokens(task=task, language=language, predict_timestamps=not no_timestamps)
+        forced_tokens = self.prefix_tokens[1:]
+        forced_decoder_ids = [(rank + 1, token_id) for rank, token_id in enumerate(forced_tokens)]
+        return forced_decoder_ids
+
+    def _decode(
+        self, token_ids, skip_special_tokens=False, clean_up_tokenization_spaces=True, 
+        normalize=False, basic_normalize=False, remove_diacritics=False, **kwargs
+    ) -> str:
+        filtered_ids = token_ids
+        if skip_special_tokens:
+            filtered_ids = [t for t in token_ids if t not in self.all_special_ids]
+
+        tokens = self.convert_ids_to_tokens(filtered_ids)
+        text = self.convert_tokens_to_string(tokens)
+        
+        if clean_up_tokenization_spaces:
+            text = re.sub(r"\s+", " ", text).strip()
+        
+        text = re.sub(self.timestamp_pat, "", text)
+        
+        return text
+
+    def decode(
+        self, token_ids, skip_special_tokens=False, clean_up_tokenization_spaces=True, 
+        output_offsets=False, time_precision=0.02, decode_with_timestamps=False, **kwargs
+    ) -> str:
+        return self._decode(token_ids, skip_special_tokens, clean_up_tokenization_spaces, **kwargs)
+
+    def save_vocabulary(self, save_directory: str, filename_prefix: Optional[str] = None) -> Tuple[str]:
+        if not os.path.isdir(save_directory):
+            logger.error(f"Vocabulary path ({save_directory}) should be a directory")
+            return
+
+        vocab_file = os.path.join(
+            save_directory, 
+            (filename_prefix + "-" if filename_prefix else "") + VOCAB_FILES_NAMES["vocab_file"]
+        )
+        merges_file = os.path.join(
+            save_directory, 
+            (filename_prefix + "-" if filename_prefix else "") + VOCAB_FILES_NAMES["merges_file"]
+        )
+
+        if self._vocab:
+            with open(vocab_file, "w", encoding="utf-8") as f:
+                json.dump(self._vocab, f, indent=2, ensure_ascii=False)
+
+        if self._merges:
+            with open(merges_file, "w", encoding="utf-8") as f:
+                f.write("#version: 0.2\n")
+                f.writelines(merge + "\n" for merge in self._merges)
+
+        return (vocab_file, merges_file)
+
+    @lru_cache
+    def timestamp_ids(self, time_precision=0.02):
+        return self.convert_tokens_to_ids([("<|%.2f|>" % (i * time_precision)) for i in range(1500 + 1)])
+
+__all__ = ["MiniWhisperTokenizer"]