| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | from dataclasses import dataclass |
| 4 | |
| 5 | from transformers import ParakeetEncoderConfig, PreTrainedConfig |
| 6 | |
| 7 | |
| 8 | class ParakeetConfig(ParakeetEncoderConfig): |
| 9 | def __init__( |
| 10 | self, |
| 11 | llm_hidden_size: int, |
| 12 | projection_hidden_size: int, |
| 13 | projection_bias: bool, |
| 14 | sampling_rate: int, |
| 15 | projection_eps: float = 1e-5, |
| 16 | **kwargs, |
| 17 | ): |
| 18 | super().__init__(**kwargs) |
| 19 | self.llm_hidden_size = llm_hidden_size |
| 20 | self.projection_hidden_size = projection_hidden_size |
| 21 | self.projection_bias = projection_bias |
| 22 | self.sampling_rate = sampling_rate |
| 23 | self.projection_eps = projection_eps |
| 24 | |
| 25 | @staticmethod |
| 26 | def from_hf_config( |
| 27 | config: PreTrainedConfig, *, llm_hidden_size: int, max_model_len: int |
| 28 | ) -> "ParakeetConfig": |
| 29 | assert isinstance(config, PreTrainedConfig) |
| 30 | return ParakeetConfig( |
| 31 | **config.to_dict(), |
| 32 | scale_input=False, |
| 33 | attention_bias=False, |
| 34 | llm_hidden_size=llm_hidden_size, |
| 35 | max_position_embeddings=max_model_len |
| 36 | + 1, # + 1 because it seems like max_model_len+1 can be passed |
| 37 | ) |
| 38 | |
| 39 | |
| 40 | @dataclass(kw_only=True, frozen=True) |
| 41 | class ExtractorConfig: |
| 42 | feature_size: int |
| 43 | sampling_rate: int |
| 44 | subsampling_factor: int |
| 45 | subsampling_conv_kernel_size: int |
| 46 | subsampling_conv_stride: int |
| 47 | hop_length: int = 160 |
| 48 | """Default `160`: Matches HF default""" |
| 49 | clip_duration_s: int = 30 |
| 50 | clip_min_duration_s: float = 0.1 |
| 51 | |
| 52 | win_length: int = 400 |
| 53 | preemphasis: float = 0.97 |
| 54 | n_fft: int = 512 |
| 55 | padding_value: float = 0.0 |
| 56 | |
| 57 | @classmethod |
| 58 | def from_hf_config(cls, config: PreTrainedConfig) -> "ExtractorConfig": |
| 59 | assert isinstance(config, PreTrainedConfig) |
| 60 | defaults = ("hop_length", "win_length", "preemphasis", "n_fft", "padding_value") |
| 61 | optional_kwargs = { |
| 62 | name: getattr(config, name) for name in defaults if hasattr(config, name) |
| 63 | } |
| 64 | |
| 65 | return cls( |
| 66 | feature_size=config.num_mel_bins, |
| 67 | sampling_rate=config.sampling_rate, |
| 68 | subsampling_factor=config.subsampling_factor, |
| 69 | subsampling_conv_kernel_size=config.subsampling_conv_kernel_size, |
| 70 | subsampling_conv_stride=config.subsampling_conv_stride, |
| 71 | **optional_kwargs, |
| 72 | ) |