Upload folder using huggingface_hub

config.json

@@ -67,7 +67,7 @@
     "torch_dtype": "bfloat16",
     "use_cache": true,
     "use_sliding_window": false,
-    "vocab_size": 152064
+    "vocab_size": 151936
   },
   "torch_dtype": "float32",
   "transformers_version": "4.52.4"

configuration_midashenglm.py

@@ -1,5 +1,4 @@
-from ast import Dict
-from typing import Optional, Tuple, Union
+from typing import Dict, Optional, Tuple, Union
 
 from transformers import PretrainedConfig
 from transformers.models.qwen2_5_omni.configuration_qwen2_5_omni import (
@@ -66,7 +65,7 @@ class MiDashengLMConfig(PretrainedConfig):
         self,
         audio_encoder_config: Dict = {},
         subsample_factor: int = 5,
-        text_config: Dict = None,
+        text_config: Dict = {},
         **kwargs,
     ):
         self.audio_encoder_config = DashengConfig(**audio_encoder_config)

model.safetensors.index.json

@@ -1,6 +1,6 @@
 {
   "metadata": {
-    "total_size": 9385880844
+    "total_size": 9384832268
   },
   "weight_map": {
     "audio_encoder.blocks.0.attn.proj.bias": "model-00001-of-00002.safetensors",

modeling_midashenglm.py

@@ -1,13 +1,14 @@
+import collections
 import collections.abc
 from dataclasses import dataclass
-from functools import partial
-from typing import Any, Callable, Iterable, List, Optional, Tuple, Type, Union
+from typing import Any, Callable, Iterable, Optional, Sequence, Tuple, Union, cast
 
 import torch
 import torch.nn as nn
 import torchaudio.transforms as audio_transforms
 from torch import Tensor
 from transformers import GenerationMixin, PreTrainedModel
+from transformers.cache_utils import Cache
 from transformers.modeling_outputs import BaseModelOutputWithPast, ModelOutput
 from transformers.models.qwen2_5_omni.configuration_qwen2_5_omni import (
     Qwen2_5OmniTextConfig,
@@ -18,28 +19,33 @@ from transformers.models.qwen2_5_omni.modeling_qwen2_5_omni import (
 
 from .configuration_midashenglm import DashengConfig, MiDashengLMConfig
 
+_Tuple2 = Union[int, Tuple[int, int], Sequence[int]]
 
-def to_2tuple(x: Any) -> Tuple[Any, Any]:
-    if isinstance(x, collections.abc.Iterable):
-        return x
+
+def _resolve_tuple2(x: _Tuple2) -> Tuple[int, int]:
+    if isinstance(x, collections.abc.Sequence):
+        assert len(x) == 2, (
+            f"Expected a sequence of length 2, got {x} with length {len(x)}"
+        )
+        return cast(Tuple[int, int], tuple(x))
     return (x, x)
 
 
 class AudioPatchEmbed(nn.Module):
     def __init__(
         self,
-        input_size: Union[int, Tuple[int, int]] = 64,
-        patch_size: Union[int, Tuple[int, int]] = 16,
-        patch_stride: Union[int, Tuple[int, int]] = 16,
+        input_size: _Tuple2 = 64,
+        patch_size: _Tuple2 = 16,
+        patch_stride: _Tuple2 = 16,
         in_chans: int = 1,
         embed_dim: int = 768,
         norm_layer: Optional[Callable] = None,
         flatten: bool = False,
     ):
         super().__init__()
-        self.input_size = to_2tuple(input_size)
-        self.patch_size = to_2tuple(patch_size)
-        self.patch_stride = to_2tuple(patch_stride)
+        self.input_size = _resolve_tuple2(input_size)
+        self.patch_size = _resolve_tuple2(patch_size)
+        self.patch_stride = _resolve_tuple2(patch_stride)
         self.grid_size = (
             self.input_size[0] // self.patch_stride[0],
             self.input_size[1] // self.patch_stride[1],
@@ -48,7 +54,10 @@ class AudioPatchEmbed(nn.Module):
         self.flatten = flatten
 
         self.proj = nn.Conv2d(
-            in_chans, embed_dim, kernel_size=patch_size, stride=patch_stride
+            in_chans,
+            embed_dim,
+            kernel_size=self.patch_size,
+            stride=self.patch_stride,
         )
         self.norm = norm_layer(embed_dim) if norm_layer else nn.Identity()
 
@@ -78,14 +87,13 @@ class DashengMlp(nn.Module):
         in_features: int,
         hidden_features: Optional[int] = None,
         out_features: Optional[int] = None,
-        act_layer: Type[nn.Module] = nn.GELU,
         drop: float = 0.0,
     ):
         super().__init__()
         out_features = out_features or in_features
         hidden_features = hidden_features or in_features
         self.fc1 = nn.Linear(in_features, hidden_features)
-        self.act = act_layer()
+        self.act = nn.GELU()
         self.fc2 = nn.Linear(hidden_features, out_features)
         self.drop = nn.Dropout(drop)
 
@@ -173,13 +181,10 @@ class DashengBlock(nn.Module):
         drop: float = 0.0,
         attn_drop: float = 0.0,
         init_values: Optional[float] = None,
-        act_layer: Type[nn.Module] = nn.GELU,
-        norm_layer: Type[nn.Module] = nn.LayerNorm,
-        attention_type: Type[nn.Module] = DashengAttention,
     ):
         super().__init__()
-        self.norm1 = norm_layer(dim)
-        self.attn = attention_type(
+        self.norm1 = nn.LayerNorm(dim, eps=1e-6)
+        self.attn = DashengAttention(
             dim,
             num_heads=num_heads,
             qkv_bias=qkv_bias,
@@ -190,11 +195,10 @@ class DashengBlock(nn.Module):
             LayerScale(dim, init_values=init_values) if init_values else nn.Identity()
         )
 
-        self.norm2 = norm_layer(dim)
+        self.norm2 = nn.LayerNorm(dim, eps=1e-6)
         self.mlp = DashengMlp(
             in_features=dim,
             hidden_features=int(dim * mlp_ratio),
-            act_layer=act_layer,
             drop=drop,
         )
         self.ls2 = (
@@ -250,7 +254,6 @@ class DashengAudioTransformer(PreTrainedModel):
             torch.randn(1, config.embed_dim, self.patch_embed.grid_size[0], 1) * 0.02
         )
 
-        norm_layer = partial(nn.LayerNorm, eps=1e-6)
         self.pos_drop = nn.Dropout(p=config.drop_rate)
         self.blocks = nn.ModuleList(
             DashengBlock(
@@ -261,11 +264,10 @@ class DashengAudioTransformer(PreTrainedModel):
                 init_values=config.init_values,
                 drop=config.drop_rate,
                 attn_drop=config.attn_drop_rate,
-                norm_layer=norm_layer,
             )
             for i in range(config.depth)
         )
-        self.norm = norm_layer(config.embed_dim)
+        self.norm = nn.LayerNorm(config.embed_dim, eps=1e-6)
 
         self.post_init()
 
@@ -295,7 +297,7 @@ class DashengAudioTransformer(PreTrainedModel):
         self,
         x: torch.Tensor,
         x_length: Optional[torch.Tensor] = None,
-    ) -> torch.Tensor:
+    ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]:
         x = self.front_end(x)
         target_length_in_patches = self.target_length // 4
         x = x.unsqueeze(1)
@@ -363,10 +365,10 @@ class AudioProjectorSubsample(nn.Module):
 
 @dataclass
 class Qwen25OmniTextModelOutput(ModelOutput):
-    logits: torch.FloatTensor = None
-    past_key_values: Optional[List[torch.FloatTensor]] = None
-    hidden_states: Optional[Tuple[torch.FloatTensor]] = None
-    attentions: Optional[Tuple[torch.FloatTensor]] = None
+    logits: Optional[torch.FloatTensor] = None
+    past_key_values: Optional[Cache] = None
+    hidden_states: Optional[Tuple[torch.FloatTensor, ...]] = None
+    attentions: Optional[Tuple[torch.FloatTensor, ...]] = None
 
 
 class Qwen25OmniThinkerTextOnlyDecoder(PreTrainedModel, GenerationMixin):
@@ -390,10 +392,22 @@ class Qwen25OmniThinkerTextOnlyDecoder(PreTrainedModel, GenerationMixin):
 
     def forward(
         self,
+        attention_mask: Optional[Tensor] = None,
+        position_ids: Optional[torch.Tensor] = None,
         return_dict: Optional[bool] = None,
         **kwargs: Any,
-    ) -> Qwen25OmniTextModelOutput:
+    ) -> Union[Tuple, Qwen25OmniTextModelOutput]:
+        if attention_mask is not None and position_ids is None:
+            position_ids = (
+                attention_mask.long()
+                .cumsum(dim=-1)
+                .masked_fill_(attention_mask == 0, 1)
+                - 1
+            )
+
         outputs: BaseModelOutputWithPast = self.model(
+            attention_mask=attention_mask,
+            position_ids=position_ids,
             return_dict=True,
             **kwargs,
         )
@@ -463,23 +477,26 @@ class MiDashengLMModel(PreTrainedModel):
     def _prepare_with_input_ids(
         self,
         input_ids: torch.Tensor,
-        audio_embeddings: torch.Tensor,
-        audio_token_id: int,
+        audio_embeddings: Optional[torch.Tensor],
+        audio_token_id: Optional[int],
     ) -> torch.Tensor:
-        special_mask = input_ids == audio_token_id
-        assert audio_embeddings.shape[1] <= (special_mask.sum(-1)).max(), (
-            "Mask and audio embeddings seem to have different sizes: "
-            f"{audio_embeddings.shape=}, {special_mask=}, {input_ids=}, "
-            f"{audio_embeddings.shape[1]=} vs {(special_mask.sum(-1)).max()=}"
-        )
         input_embeddings = self.decoder.model.embed_tokens(input_ids)
-        audio_embeddings = audio_embeddings.to(input_embeddings.dtype)
+        if audio_embeddings is not None:
+            special_mask = input_ids == audio_token_id
+            assert audio_embeddings.shape[1] <= (special_mask.sum(-1)).max(), (
+                "Mask and audio embeddings seem to have different sizes: "
+                f"{audio_embeddings.shape=}, {special_mask=}, {input_ids=}, "
+                f"{audio_embeddings.shape[1]=} vs {(special_mask.sum(-1)).max()=}"
+            )
+            audio_embeddings = audio_embeddings.to(input_embeddings.dtype)
 
-        for i in range(len(special_mask)):
-            mask = special_mask[i]
-            number_of_tokens = mask.sum(-1)
-            input_embeddings[i, mask] = audio_embeddings[i, :number_of_tokens]
-        return input_embeddings
+            for i in range(len(special_mask)):
+                mask = special_mask[i]
+                number_of_tokens = mask.sum(-1)
+                input_embeddings[i, mask] = audio_embeddings[i, :number_of_tokens]
+            return input_embeddings
+        else:
+            return input_embeddings
 
     def forward(
         self,
@@ -487,7 +504,6 @@ class MiDashengLMModel(PreTrainedModel):
         input_values: Optional[Tensor] = None,
         inputs_embeds: Optional[Tensor] = None,
         audio_length: Optional[Iterable[int]] = None,
-        attention_mask: Optional[Tensor] = None,
         audio_token_id: Optional[int] = None,
         **kwargs: Any,
     ):
@@ -498,6 +514,11 @@ class MiDashengLMModel(PreTrainedModel):
                 )
 
             if input_values is not None:
+                if audio_token_id is None:
+                    raise ValueError(
+                        "If `input_values` is provided, `audio_token_id` must also be provided."
+                    )
+
                 input_values = input_values.to(self.device)
                 audio_encoder_hidden_states = self._forward_audio_encoder(
                     input_values, audio_length=audio_length
@@ -530,7 +551,6 @@ class MiDashengLMModel(PreTrainedModel):
         return self.decoder(
             input_ids=None,
             inputs_embeds=inputs_embeds,
-            attention_mask=attention_mask,
             **kwargs,
         )
 
@@ -548,6 +568,7 @@ class MiDashengLMModel(PreTrainedModel):
                 raise ValueError(
                     "Both `inputs_embeds` and `input_ids` are passed. Please pass only one of them."
                 )
+            input_ids = input_ids.to(self.device)
 
             if input_values is not None:
                 input_values = input_values.to(self.device)
@@ -555,15 +576,7 @@ class MiDashengLMModel(PreTrainedModel):
                     input_values, audio_length=audio_length
                 )
             else:
-                batch, _ = input_ids.shape
-                input_values = torch.zeros(
-                    batch,
-                    0,
-                    self.audio_encoder.embed_dim,
-                    device=input_ids.device,
-                )
-
-            input_ids = input_ids.to(self.device)
+                audio_encoder_hidden_states = None
             inputs_embeds = self._prepare_with_input_ids(
                 input_ids=input_ids,
                 audio_embeddings=audio_encoder_hidden_states,

preprocessor_config.json

@@ -1,13 +1,13 @@
 {
   "auto_map": {
-    "AutoProcessor": "processing_midashenglm.MiAudioLLMProcessor"
+    "AutoProcessor": "processing_midashenglm.MiDashengLMProcessor"
   },
   "do_normalize": false,
   "feature_extractor_type": "Wav2Vec2FeatureExtractor",
   "feature_size": 1,
   "padding_side": "right",
   "padding_value": 0.0,
-  "processor_class": "MiAudioLLMProcessor",
+  "processor_class": "MiDashengLMProcessor",
   "return_attention_mask": false,
   "sampling_rate": 16000
 }

processing_midashenglm.py

@@ -1,4 +1,4 @@
-from typing import List, Optional, Union
+from typing import Dict, List, Optional, Union, cast
 
 import numpy as np
 import torch
@@ -7,8 +7,8 @@ from transformers.feature_extraction_utils import BatchFeature
 from transformers.processing_utils import ProcessingKwargs, ProcessorMixin, Unpack
 
 
-class MiAudioLLMProcessorKwargs(ProcessingKwargs):
-    _defaults = {
+class MiDashengLMProcessorKwargs(ProcessingKwargs):
+    _defaults = {  # type: ignore
         "text_kwargs": {
             "padding": True,
             "padding_side": "left",
@@ -36,7 +36,7 @@ def calculate_mel_frames_dasheng(
     )
 
 
-class MiAudioLLMProcessor(ProcessorMixin):
+class MiDashengLMProcessor(ProcessorMixin):
     attributes = ["feature_extractor", "tokenizer"]
     valid_kwargs = [
         "chat_template",
@@ -49,15 +49,14 @@ class MiAudioLLMProcessor(ProcessorMixin):
 
     def __init__(
         self,
-        feature_extractor: Optional[Wav2Vec2FeatureExtractor] = None,
-        tokenizer: Optional[Union[Qwen2Tokenizer, Qwen2TokenizerFast]] = None,
+        feature_extractor: Wav2Vec2FeatureExtractor,
+        tokenizer: Union[Qwen2Tokenizer, Qwen2TokenizerFast],
         model_subsampling: int = 5,
-        chat_template: Optional[str] = None,
+        chat_template: Optional[Union[str, Dict[str, str]]] = None,
         audio_token: Optional[str] = None,
         audio_bos_token: Optional[str] = None,
         audio_eos_token: Optional[str] = None,
     ):
-        assert tokenizer is not None, "Tokenizer Needs to be passed"
         assert audio_token is not None or hasattr(tokenizer, "audio_token"), (
             "Either `audio_token` must be provided or tokenizer must have `audio_token` attribute."
         )
@@ -67,22 +66,62 @@ class MiAudioLLMProcessor(ProcessorMixin):
         assert audio_eos_token is not None or hasattr(tokenizer, "audio_eos_token"), (
             "Either `audio_eos_token` must be provided or tokenizer must have `audio_eos_token` attribute."
         )
+        assert not feature_extractor.do_normalize, (
+            "This model does not use normalization. Please set `do_normalize=False` in the feature extractor."
+        )
 
         if chat_template is None:
             chat_template = tokenizer.chat_template
 
-        self.audio_token: str = audio_token or tokenizer.audio_token
-        self.audio_bos_token = audio_bos_token or tokenizer.audio_bos_token
-        self.audio_eos_token = audio_eos_token or tokenizer.audio_eos_token
-        self.audio_token_id = tokenizer.convert_tokens_to_ids(self.audio_token)
-        self.model_subsampling = model_subsampling
+        def get_token(token_name: str) -> str:
+            if not hasattr(tokenizer, token_name):
+                raise ValueError(
+                    f"Tokenizer does not have attribute `{token_name}`. "
+                    "Please provide it as an argument to the processor."
+                )
+            token = getattr(tokenizer, token_name)
+            if not isinstance(token, str):
+                raise TypeError(
+                    f"Expected token {token_name} to be a string, but got {type(token)}."
+                )
+            return token
 
-        if feature_extractor is not None:
-            assert not feature_extractor.do_normalize, (
-                "This model does not use normalization. Please set `do_normalize=False` in the feature extractor."
-            )
+        self.audio_token = audio_token or get_token("audio_token")
+        self.audio_bos_token = audio_bos_token or get_token("audio_bos_token")
+        self.audio_eos_token = audio_eos_token or get_token("audio_eos_token")
+
+        self.audio_token_id = cast(
+            int, tokenizer.convert_tokens_to_ids(self.audio_token)
+        )
+        self.model_subsampling = model_subsampling
+        self.sampling_rate = feature_extractor.sampling_rate
 
         super().__init__(feature_extractor, tokenizer, chat_template=chat_template)
+        self.feature_extractor: Wav2Vec2FeatureExtractor
+        self.tokenizer: Union[Qwen2Tokenizer, Qwen2TokenizerFast]
+        self.chat_template: Optional[Union[str, Dict[str, str]]]
+
+    def _process_messages_for_chat_template(
+        self,
+        conversation,
+        batch_images,
+        batch_videos,
+        batch_video_metadata,
+        **mm_load_kwargs,
+    ):
+        if (sr := mm_load_kwargs.get("sampling_rate", None)) is not None:
+            if sr != self.sampling_rate:
+                raise ValueError(
+                    f"This model is trained with a sampling rate of {self.sampling_rate}, "
+                    f"but the sampling rate {sr} is used to load audio."
+                )
+        return super()._process_messages_for_chat_template(
+            conversation,
+            batch_images,
+            batch_videos,
+            batch_video_metadata,
+            **mm_load_kwargs,
+        )
 
     @classmethod
     def _validate_audio_sample(
@@ -117,7 +156,7 @@ class MiAudioLLMProcessor(ProcessorMixin):
         self,
         text: Optional[List[str]] = None,
         audio: Optional[Union[List[np.ndarray], List[torch.Tensor]]] = None,
-        **kwargs: Unpack[MiAudioLLMProcessorKwargs],
+        **kwargs: Unpack[MiDashengLMProcessorKwargs],
     ) -> BatchFeature:
         if text is None:
             raise ValueError("You need to specify `text` input to process.")
@@ -135,7 +174,7 @@ class MiAudioLLMProcessor(ProcessorMixin):
             raise ValueError("This model does not support images or videos.")
 
         output_kwargs = self._merge_kwargs(
-            MiAudioLLMProcessorKwargs,
+            MiDashengLMProcessorKwargs,  # type: ignore # Bad type hint in transformers
             tokenizer_init_kwargs=self.tokenizer.init_kwargs,
             **kwargs,
         )
@@ -157,7 +196,9 @@ class MiAudioLLMProcessor(ProcessorMixin):
 
             # + Padding
             audio_inputs = self.feature_extractor(
-                audio, **output_kwargs["audio_kwargs"]
+                audio,
+                sampling_rate=self.sampling_rate,
+                **output_kwargs["audio_kwargs"],
             )
 
             # remove attention mask, dasheng uses lengths
@@ -216,28 +257,17 @@ class MiAudioLLMProcessor(ProcessorMixin):
 
         return_tensors = output_kwargs["text_kwargs"].pop("return_tensors", "pt")
         inputs = self.tokenizer(text, **output_kwargs["text_kwargs"])
-        if hasattr(self, "_check_special_mm_tokens"):
-            self._check_special_mm_tokens(text, inputs, modalities=["audio"])
+        self._check_special_mm_tokens(
+            text,
+            BatchFeature(inputs),  # type: ignore
+            modalities=["audio"],
+        )
 
         if audio is not None:
             inputs.update(audio_inputs)
 
         return BatchFeature(data={**inputs}, tensor_type=return_tensors)
 
-    def batch_decode(self, *args, **kwargs):
-        """
-        This method forwards all its arguments to Qwen2TokenizerFast's [`~PreTrainedTokenizer.batch_decode`]. Please
-        refer to the docstring of this method for more information.
-        """
-        return self.tokenizer.batch_decode(*args, **kwargs)
-
-    def decode(self, *args, **kwargs):
-        """
-        This method forwards all its arguments to Qwen2TokenizerFast's [`~PreTrainedTokenizer.decode`]. Please refer to
-        the docstring of this method for more information.
-        """
-        return self.tokenizer.decode(*args, **kwargs)
-
     @property
     def model_input_names(self):
         tokenizer_input_names = self.tokenizer.model_input_names

processor_config.json

@@ -3,8 +3,8 @@
   "audio_eos_token": "<|audio_eos|>",
   "audio_token": "<|AUDIO|>",
   "auto_map": {
-    "AutoProcessor": "processing_midashenglm.MiAudioLLMProcessor"
+    "AutoProcessor": "processing_midashenglm.MiDashengLMProcessor"
   },
   "model_subsampling": 5,
-  "processor_class": "MiAudioLLMProcessor"
+  "processor_class": "MiDashengLMProcessor"
 }

tokenizer_config.json

@@ -337,7 +337,7 @@
   "audio_eos_token": "<|audio_eos|>",
   "audio_token": "<|AUDIO|>",
   "auto_map": {
-    "AutoProcessor": "processing_midashenglm.MiAudioLLMProcessor"
+    "AutoProcessor": "processing_midashenglm.MiDashengLMProcessor"
   },
   "bos_token": null,
   "clean_up_tokenization_spaces": false,
@@ -355,7 +355,7 @@
   "image_token": "<|IMAGE|>",
   "model_max_length": 32768,
   "pad_token": "<|endoftext|>",
-  "processor_class": "MiAudioLLMProcessor",
+  "processor_class": "MiDashengLMProcessor",
   "split_special_tokens": false,
   "tokenizer_class": "Qwen2Tokenizer",
   "unk_token": null,