| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | from typing import Any |
| 5 | |
| 6 | from transformers import PreTrainedConfig, Qwen3Config |
| 7 | from transformers.models.whisper.configuration_whisper import WhisperConfig |
| 8 | |
| 9 | |
| 10 | class MossTranscribeDiarizeConfig(PreTrainedConfig): |
| 11 | """Configuration for MOSS-Transcribe-Diarize.""" |
| 12 | |
| 13 | model_type = "moss_transcribe_diarize" |
| 14 | sub_configs = {"text_config": Qwen3Config, "audio_config": WhisperConfig} |
| 15 | keys_to_ignore_at_inference = ["past_key_values"] |
| 16 | |
| 17 | def __init__( |
| 18 | self, |
| 19 | text_config: dict[str, Any] | Qwen3Config | None = None, |
| 20 | audio_config: dict[str, Any] | WhisperConfig | None = None, |
| 21 | audio_token_id: int = 151671, |
| 22 | audio_merge_size: int = 4, |
| 23 | adaptor_input_dim: int | None = None, |
| 24 | tie_word_embeddings: bool = True, |
| 25 | **kwargs: Any, |
| 26 | ) -> None: |
| 27 | text_config_obj: Qwen3Config |
| 28 | if text_config is None: |
| 29 | text_config_obj = Qwen3Config( |
| 30 | vocab_size=151936, |
| 31 | hidden_size=1024, |
| 32 | intermediate_size=3072, |
| 33 | num_hidden_layers=28, |
| 34 | num_attention_heads=16, |
| 35 | num_key_value_heads=8, |
| 36 | head_dim=128, |
| 37 | max_position_embeddings=40960, |
| 38 | tie_word_embeddings=tie_word_embeddings, |
| 39 | rope_theta=1_000_000.0, |
| 40 | layer_types=["full_attention"] * 28, |
| 41 | ) |
| 42 | elif isinstance(text_config, dict): |
| 43 | text_config_obj = Qwen3Config(**text_config) |
| 44 | else: |
| 45 | text_config_obj = text_config |
| 46 | |
| 47 | audio_config_obj: WhisperConfig |
| 48 | if audio_config is None: |
| 49 | audio_config_obj = WhisperConfig( |
| 50 | num_mel_bins=80, |
| 51 | d_model=1024, |
| 52 | encoder_layers=24, |
| 53 | encoder_attention_heads=16, |
| 54 | encoder_ffn_dim=4096, |
| 55 | max_source_positions=1500, |
| 56 | dropout=0.0, |
| 57 | attention_dropout=0.0, |
| 58 | activation_dropout=0.0, |
| 59 | activation_function="gelu", |
| 60 | encoder_layerdrop=0.0, |
| 61 | scale_embedding=False, |
| 62 | ) |
| 63 | elif isinstance(audio_config, dict): |
| 64 | audio_config_obj = WhisperConfig(**audio_config) |
| 65 | else: |
| 66 | audio_config_obj = audio_config |
| 67 | |
| 68 | text_config_obj.tie_word_embeddings = tie_word_embeddings |
| 69 | if not getattr(text_config_obj, "layer_types", None): |
| 70 | text_config_obj.layer_types = [ |
| 71 | "full_attention" |
| 72 | ] * text_config_obj.num_hidden_layers |
| 73 | |
| 74 | super().__init__(tie_word_embeddings=tie_word_embeddings, **kwargs) |
| 75 | |
| 76 | self.text_config = text_config_obj |
| 77 | self.audio_config = audio_config_obj |
| 78 | self.audio_token_id = int(audio_token_id) |
| 79 | self.audio_merge_size = int(audio_merge_size) |
| 80 | self.adaptor_input_dim = ( |
| 81 | int(adaptor_input_dim) |
| 82 | if adaptor_input_dim is not None |
| 83 | else int(audio_config_obj.d_model) * int(audio_merge_size) |
| 84 | ) |
| 85 | |
| 86 | self.vocab_size = int(text_config_obj.vocab_size) |
| 87 | self.hidden_size = int(text_config_obj.hidden_size) |
| 88 | self.intermediate_size = int(text_config_obj.intermediate_size) |
| 89 | self.num_hidden_layers = int(text_config_obj.num_hidden_layers) |
| 90 | self.num_attention_heads = int(text_config_obj.num_attention_heads) |
| 91 | self.num_key_value_heads = int(text_config_obj.num_key_value_heads) |
| 92 | self.head_dim = int(text_config_obj.head_dim) |
| 93 | self.hidden_act = text_config_obj.hidden_act |
| 94 | self.max_position_embeddings = int(text_config_obj.max_position_embeddings) |
| 95 | self.rms_norm_eps = float(text_config_obj.rms_norm_eps) |
| 96 | rope_parameters = getattr(text_config_obj, "rope_parameters", None) |
| 97 | rope_theta = float(getattr(text_config_obj, "rope_theta", 1_000_000.0)) |
| 98 | if rope_parameters is None: |
| 99 | rope_parameters = { |
| 100 | "rope_type": "default", |
| 101 | "rope_theta": rope_theta, |
| 102 | } |
| 103 | text_config_obj.rope_parameters = rope_parameters |
| 104 | self.rope_parameters = rope_parameters |
| 105 | self.rope_theta = float(rope_parameters.get("rope_theta", rope_theta)) |
| 106 | self.attention_bias = bool(text_config_obj.attention_bias) |
| 107 | self.attention_dropout = float(text_config_obj.attention_dropout) |
| 108 | self.is_causal = True |