Upload folder using huggingface_hub

model-00001-of-00007.safetensors

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:324d47a69b306b736f9c1ed9c3ac6b2f08dd25f3238e0995ca03d1f628d14d3f
-size 4962055488
+oid sha256:f1580ce43afa6d023bb89c1d68ed86e155119885fbf2f6dafe98e6696593f7b1
+size 4961987424

model.safetensors.index.json

@@ -1,6 +1,6 @@
 {
   "metadata": {
-    "total_size": 33127027980
+    "total_size": 33126960136
   },
   "weight_map": {
     "audio_encoder.blocks.0.attn.proj.bias": "model-00001-of-00007.safetensors",
@@ -388,8 +388,6 @@
     "audio_encoder.blocks.9.norm2.bias": "model-00001-of-00007.safetensors",
     "audio_encoder.blocks.9.norm2.weight": "model-00001-of-00007.safetensors",
     "audio_encoder.freq_pos_embed": "model-00001-of-00007.safetensors",
-    "audio_encoder.front_end.0.mel_scale.fb": "model-00001-of-00007.safetensors",
-    "audio_encoder.front_end.0.spectrogram.window": "model-00001-of-00007.safetensors",
     "audio_encoder.init_bn.bias": "model-00001-of-00007.safetensors",
     "audio_encoder.init_bn.num_batches_tracked": "model-00001-of-00007.safetensors",
     "audio_encoder.init_bn.running_mean": "model-00001-of-00007.safetensors",

modeling_midashenglm.py

@@ -5,7 +5,7 @@ from typing import Any, Callable, Iterable, List, Optional, Sequence, Tuple, Uni
 
 import torch
 import torch.nn as nn
-import torchaudio.transforms as audio_transforms
+import torchaudio.functional as F
 from torch import Tensor
 from transformers import GenerationMixin, PreTrainedModel
 from transformers.cache_utils import Cache
@@ -217,6 +217,59 @@ class DashengBlock(nn.Module):
         return x
 
 
+class DashengFrontend(nn.Module):
+    def __init__(self, config: DashengConfig):
+        super().__init__()
+        self.config = config
+
+        spectrogram_window = torch.hann_window(self.config.win_length)
+        self.register_buffer(
+            "spectrogram_window",
+            spectrogram_window,
+            persistent=False,
+        )
+        self.spectrogram_window: torch.Tensor
+
+        melscale_fbanks = F.melscale_fbanks(
+            n_freqs=self.config.n_fft // 2 + 1,
+            f_min=self.config.f_min,
+            f_max=self.config.f_max,
+            n_mels=self.config.n_mels,
+            sample_rate=self.config.sample_rate,
+        )
+        self.register_buffer("melscale_fbanks", melscale_fbanks, persistent=False)
+        self.melscale_fbanks: torch.Tensor
+
+    def forward(self, waveform: torch.Tensor) -> torch.Tensor:
+        spectrogram = F.spectrogram(
+            waveform=waveform.to(torch.float32),
+            pad=0,
+            window=self.spectrogram_window,
+            n_fft=self.config.n_fft,
+            hop_length=self.config.hop_length,
+            win_length=self.config.win_length,
+            power=2,
+            normalized=False,
+            center=self.config.center,
+        )
+        mel_spectrogram = (spectrogram.mT @ self.melscale_fbanks.to(torch.float32)).mT
+        # x has shape [batch, freq, time].
+        # F.amplitude_to_DB accepts inputs shaped as:
+        #   - [freq, time]
+        #   - [channel, freq, time]
+        #   - [..., channel, freq, time]
+        # Here we insert a channel dimension of size 1 before calling it,
+        # then remove that extra dimension afterward.
+        log_mel_spectrogram = F.amplitude_to_DB(
+            mel_spectrogram.unsqueeze(1),
+            multiplier=10,
+            amin=1e-10,
+            db_multiplier=0,
+            top_db=120,
+        ).squeeze(1)
+        return log_mel_spectrogram.to(waveform.dtype)
+
+
 class DashengAudioTransformer(PreTrainedModel):
     config_class = DashengConfig
     supports_gradient_checkpointing = True
@@ -229,19 +282,7 @@ class DashengAudioTransformer(PreTrainedModel):
         self.hop_length = config.hop_length
         self.gradient_checkpointing = False
 
-        self.front_end = nn.Sequential(
-            audio_transforms.MelSpectrogram(
-                f_min=config.f_min,
-                f_max=config.f_max,
-                center=config.center,
-                win_length=config.win_length,
-                hop_length=config.hop_length,
-                sample_rate=config.sample_rate,
-                n_fft=config.n_fft,
-                n_mels=config.n_mels,
-            ),
-            audio_transforms.AmplitudeToDB(top_db=120),
-        )
+        self.front_end = DashengFrontend(config)
 
         self.init_bn = nn.BatchNorm2d(config.n_mels, momentum=0.01)
 
@@ -272,7 +313,7 @@ class DashengAudioTransformer(PreTrainedModel):
                 drop=config.drop_rate,
                 attn_drop=config.attn_drop_rate,
             )
-            for i in range(config.depth)
+            for _ in range(config.depth)
         )
         self.norm = nn.LayerNorm(config.embed_dim, eps=1e-6)