@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / c8aca0c9e1

Fork 0

A high-throughput and memory-efficient inference and serving engine for LLMs

Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git

Synced Nov 7, 14:00 UTC

Support parakeet as audio encoder for nemotron-nano-vl (#35100)

Signed-off-by: Netanel Haber <58652339+netanel-haber@users.noreply.github.com> Co-authored-by: Roger Wang <hey@rogerw.io>

Netanel Haber committed Feb 27, 2026 at 20:07 UTC c8aca0c9e1b35ee4a1683a01467e638b23076a37

3 files changed +448 -20

vllm/model_executor/models/nano_nemotron_vl.py

+254 -20

@@ -44,6 +44,7 @@ from vllm.model_executor.models.internvl import (

44 )

45 from vllm.model_executor.models.module_mapping import MultiModelKeys

46 from vllm.model_executor.models.nemotron_h import NemotronHForCausalLM

47 +from vllm.model_executor.models.parakeet import ParakeetExtractor, ProjectedParakeet

48 from vllm.model_executor.models.radio import RadioModel, calc_seq_lens

49 from vllm.model_executor.models.utils import (

50 init_vllm_registered_model,

@@ -55,12 +56,14 @@ from vllm.multimodal.evs import (

56 compute_retention_mask,

57 )

58 from vllm.multimodal.inputs import (

59 + AudioItem,

60 MultiModalDataDict,

61 MultiModalFieldConfig,

62 MultiModalKwargsItems,

63 VideoItem,

64 )

65 from vllm.multimodal.parse import (

66 + AudioProcessorItems,

67 ImageEmbeddingItems,

68 ImageProcessorItems,

69 ImageSize,

@@ -91,9 +94,29 @@ Image.MAX_IMAGE_PIXELS = None # Disable the limit entirely

94 # Alternative: Set a specific higher limit

95 # Image.MAX_IMAGE_PIXELS = 300000000 # ~300M pixels

97 +

98 +class NanoNemotronVLAudioFeatureInputs(TensorSchema):

99 + """

100 + Dimensions:

101 + - b: Number of audio clips

102 + - t: Audio feature length

103 + - f: Feature size (mel bins)

104 + """

105 +

106 + type: Literal["audio_features"] = "audio_features"

107 + input_audio_features: Annotated[torch.Tensor, TensorShape("b", "t", "f")]

108 + feature_attention_mask: Annotated[torch.Tensor, TensorShape("b", "t")]

109 + audio_feature_lengths: Annotated[torch.Tensor, TensorShape("b")]

110 +

111 +

112 +MAX_AUDIO_LEN_S = 10 * 60 # 10 minutes

113 +

114 IMG_START = "<img>"

115 IMG_END = "</img>"

116 IMG_CONTEXT = "<image>"

117 +AUDIO_START = "<so_start>"

118 +AUDIO_END = "<so_end>"

119 +AUDIO_CONTEXT = "<so_embedding>"

120

121 # Profiling

122 # MAX_FRAMES = 16

@@ -820,6 +843,11 @@ class NanoNemotronVLProcessor(BaseNanoNemotronVLProcessor):

843 self.video_token = video_token

844 self.video_pruning_rate = video_pruning_rate

845

846 + self.audio_extractor: ParakeetExtractor | None = None

847 + raw_sound_config = getattr(config, "sound_config", None)

848 + if raw_sound_config is not None:

849 + self.audio_extractor = ParakeetExtractor(raw_sound_config)

850 +

851 # Pre-tokenize special tokens for video processing

852 # to avoid repeated tokenization

853 self._img_start_token_ids = tokenizer.encode(

@@ -952,11 +980,53 @@ class NanoNemotronVLProcessor(BaseNanoNemotronVLProcessor):

980 text = [t.replace("<video>", video_repl_text, 1) for t in text]

981 return text, video_inputs

982

983 + def _preprocess_audio(

984 + self,

985 + text: list[str],

986 + audios: list[npt.NDArray],

987 + ):

988 + if len(audios) == 0:

989 + return text, {}

990 + assert self.audio_extractor is not None

991 +

992 + extractor = self.audio_extractor

993 +

994 + parts = [x for x in re.split(f"({re.escape(AUDIO_CONTEXT)})", text[0]) if x]

995 + token_count = parts.count(AUDIO_CONTEXT)

996 + if token_count != len(audios):

997 + raise ValueError(

998 + "Number of audio tokens in text does not match the number "

999 + f"of audios (tokens={token_count}, audios={len(audios)})."

1000 + )

1001 + audio_index = 0

1002 + for idx, part in enumerate(parts):

1003 + if part == AUDIO_CONTEXT:

1004 + audio_repl = self.get_audio_repl(audios[audio_index])

1005 + parts[idx] = audio_repl.full

1006 + audio_index += 1

1007 + text = ["".join(parts)]

1008 + audio_inputs = extractor(

1009 + audios,

1010 + sampling_rate=extractor.sampling_rate,

1011 + return_tensors="pt",

1012 + )

1013 + input_audio_features = audio_inputs.input_features

1014 + feature_attention_mask = audio_inputs.attention_mask

1015 + audio_feature_lengths = feature_attention_mask.sum(dim=1)

1016 + audio_inputs = {

1017 + "input_audio_features": input_audio_features,

1018 + "feature_attention_mask": feature_attention_mask,

1019 + "audio_feature_lengths": audio_feature_lengths,

1020 + }

1021 +

1022 + return text, audio_inputs

1023 +

1024 def __call__(

1025 self,

1026 text: str | list[str] | None = None,

1027 images: Image.Image | list[Image.Image] | None = None,

1028 videos: list[tuple[npt.NDArray, dict[str, Any]]] | None = None,

1029 + audios: AudioItem | list[AudioItem] | None = None,

1030 return_tensors: str | TensorType | None = None,

1031 max_num_tiles: int | None = None,

1032 ) -> BatchFeature:

@@ -964,8 +1034,8 @@ class NanoNemotronVLProcessor(BaseNanoNemotronVLProcessor):

1034 if max_num_tiles is None:

1035 max_num_tiles = self.max_num_tiles

1036

967 - text, images, videos = [

968 - self._make_batch_input(x) for x in (text, images, videos)

1037 + text, images, videos, audios = [

1038 + self._make_batch_input(x) for x in (text, images, videos, audios)

1039 ]

1040

1041 text, image_inputs = self._preprocess_image(

@@ -980,17 +1050,22 @@ class NanoNemotronVLProcessor(BaseNanoNemotronVLProcessor):

1050 max_num_tiles=1,

1051 )

1052

1053 + text, audio_inputs = self._preprocess_audio(

1054 + text=text,

1055 + audios=audios,

1056 + )

1057 +

1058 text_inputs = self.tokenizer(text, add_special_tokens=False)

1059

1060 + combined_inputs = {**text_inputs, **video_inputs, **audio_inputs}

1061 +

1062 if self.dynamic_tiler is None:

1063 batch = BatchFeature(

987 - {**text_inputs, **video_inputs, **image_inputs},

1064 + {**combined_inputs, **image_inputs},

1065 tensor_type=return_tensors,

1066 )

1067 else:

991 - batch = BatchFeature(

992 - {**text_inputs, **video_inputs}, tensor_type=return_tensors

993 - )

1068 + batch = BatchFeature(combined_inputs, tensor_type=return_tensors)

1069 # allow images to be exempt from the BatchFeature validation:

1070 # We will .stack() them in _parse_and_validate_image_input

1071 batch.update(image_inputs)

@@ -1006,6 +1081,15 @@ class NanoNemotronVLProcessor(BaseNanoNemotronVLProcessor):

1081

1082 return PromptUpdateDetails.select_text(repl_full, IMG_CONTEXT)

1083

1084 + def get_audio_repl(

1085 + self,

1086 + audio: npt.NDArray,

1087 + ) -> PromptUpdateDetails[str]:

1088 + assert self.audio_extractor is not None

1089 + num_tokens = self.audio_extractor.audio_token_count(len(audio))

1090 + repl_full = f"{AUDIO_START}{AUDIO_CONTEXT * num_tokens}{AUDIO_END}"

1091 + return PromptUpdateDetails.select_text(repl_full, AUDIO_CONTEXT)

1092 +

1093 @classmethod

1094 def get_video_repl(

1095 cls,

@@ -1147,15 +1231,28 @@ class NanoNemotronVLProcessingInfo(BaseNanoNemotronVLProcessingInfo):

1231 def supports_video(self):

1232 return self.get_hf_processor().supports_video

1233

1234 + @property

1235 + def audio_extractor(self) -> ParakeetExtractor | None:

1236 + return self.get_hf_processor().audio_extractor

1237 +

1238 def get_data_parser(self):

1239 + target_sr = None

1240 + target_channels = None

1241 + if extractor := self.audio_extractor:

1242 + target_sr = extractor.sampling_rate

1243 + target_channels = 1

1244 +

1245 return MultiModalDataParser(

1246 video_needs_metadata=True,

1247 + target_sr=target_sr,

1248 + target_channels=target_channels,

1249 expected_hidden_size=self._get_expected_hidden_size(),

1250 )

1251

1252 def get_supported_mm_limits(self):

1253 video_limit = {"video": None} if self.supports_video else {}

1158 - return {**super().get_supported_mm_limits(), **video_limit}

1254 + audio_limit = {"audio": None} if self.audio_extractor is not None else {}

1255 + return {**super().get_supported_mm_limits(), **video_limit, **audio_limit}

1256

1257 def get_video_token(self) -> str | None:

1258 return IMG_CONTEXT

@@ -1304,7 +1401,16 @@ class NanoNemotronVLMultiModalProcessor(

1401 else:

1402 video_fields = {}

1403

1307 - return image_fields | video_fields

1404 + if self.info.audio_extractor is not None:

1405 + audio_fields = dict(

1406 + input_audio_features=MultiModalFieldConfig.batched("audio"),

1407 + feature_attention_mask=MultiModalFieldConfig.batched("audio"),

1408 + audio_feature_lengths=MultiModalFieldConfig.batched("audio"),

1409 + )

1410 + else:

1411 + audio_fields = {}

1412 +

1413 + return image_fields | video_fields | audio_fields

1414

1415 def _get_prompt_updates(

1416 self,

@@ -1373,6 +1479,20 @@ class NanoNemotronVLMultiModalProcessor(

1479 ),

1480 ]

1481

1482 + def get_audio_replacement(item_idx: int):

1483 + audios = mm_items.get_items("audio", AudioProcessorItems)

1484 + return hf_processor.get_audio_repl(audios.get(item_idx))

1485 +

1486 + if self.info.audio_extractor is not None:

1487 + prompt_repl = [

1488 + *prompt_repl,

1489 + PromptReplacement(

1490 + modality="audio",

1491 + target=AUDIO_CONTEXT,

1492 + replacement=get_audio_replacement,

1493 + ),

1494 + ]

1495 +

1496 return prompt_repl

1497

1498

@@ -1422,8 +1542,13 @@ class NanoNemotronVLDummyInputsBuilder(

1542

1543 def get_dummy_text(self, mm_counts: Mapping[str, int]) -> str:

1544 num_videos = mm_counts.get("video", 0)

1545 + num_audios = mm_counts.get("audio", 0)

1546

1426 - return super().get_dummy_text(mm_counts) + "<video>" * num_videos

1547 + return (

1548 + super().get_dummy_text(mm_counts)

1549 + + "<video>" * num_videos

1550 + + AUDIO_CONTEXT * num_audios

1551 + )

1552

1553 def _get_dummy_videos(

1554 self,

@@ -1482,7 +1607,25 @@ class NanoNemotronVLDummyInputsBuilder(

1607 }

1608 else:

1609 dummy_video = {}

1485 - return {**dummy_image, **dummy_video}

1610 +

1611 + if extractor := self.info.audio_extractor:

1612 + num_audios = mm_counts.get("audio", 0)

1613 + audio_overrides = mm_options.get("audio") if mm_options else None

1614 + tokens_per_audio = max(1, seq_len // max(num_audios, 1))

1615 + max_audio_num_samples = MAX_AUDIO_LEN_S * extractor.sampling_rate

1616 + calculated_max_audio_num_samples = extractor.audio_length(tokens_per_audio)

1617 + audio_len = min(max_audio_num_samples, calculated_max_audio_num_samples)

1618 + dummy_audio = {

1619 + "audio": self._get_dummy_audios(

1620 + length=audio_len,

1621 + num_audios=num_audios,

1622 + overrides=audio_overrides,

1623 + )

1624 + }

1625 + else:

1626 + dummy_audio = {}

1627 +

1628 + return {**dummy_image, **dummy_video, **dummy_audio}

1629

1630

1631 @MULTIMODAL_REGISTRY.register_processor(

@@ -1499,12 +1642,15 @@ class NemotronH_Nano_VL_V2(

1642 return "<image>"

1643 if modality.startswith("video"):

1644 return "<video>"

1645 + if modality.startswith("audio"):

1646 + return AUDIO_CONTEXT

1647 return None

1648

1649 def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):

1650 super().__init__()

1506 - config = vllm_config.model_config.hf_config

1507 - multimodal_config = vllm_config.model_config.multimodal_config

1651 + model_config = vllm_config.model_config

1652 + config = model_config.hf_config

1653 + multimodal_config = model_config.multimodal_config

1654 image_size = config.force_image_size

1655 patch_size = config.patch_size

1656 self.patch_size = patch_size

@@ -1523,10 +1669,12 @@ class NemotronH_Nano_VL_V2(

1669 hf_config=config.text_config,

1670 prefix=maybe_prefix(prefix, "language_model"),

1671 )

1526 -

1527 - with self._mark_tower_model(vllm_config, {"image", "video"}):

1672 + llm_dtype = self.language_model.config.dtype

1673 + assert isinstance(llm_dtype, torch.dtype)

1674 + self.llm_dtype = llm_dtype

1675 + with self._mark_tower_model(vllm_config, {"image", "video", "audio"}):

1676 self.vision_model = self.get_vit_model_from_radio_config(config).to(

1529 - self.language_model.config.dtype

1677 + llm_dtype

1678 )

1679

1680 # Construct the vision projection.

@@ -1547,14 +1695,26 @@ class NemotronH_Nano_VL_V2(

1695 ReLUSquaredActivation(),

1696 nn.Linear(vision_projection_hidden_size, llm_hidden_size, bias=False),

1697 )

1550 - self.mlp1 = mlp1.to(self.language_model.config.dtype)

1698 + self.mlp1 = mlp1.to(llm_dtype)

1699 + self.sound_encoder: ProjectedParakeet | None = None

1700 + if getattr(config, "sound_config", None) is not None:

1701 + logger.info_once(

1702 + "Found sound config, initializing sound encoder for Nemotron AVLM",

1703 + scope="global",

1704 + )

1705 + self.sound_encoder = ProjectedParakeet(

1706 + config.sound_config,

1707 + dtype=llm_dtype,

1708 + llm_hidden_size=llm_hidden_size,

1709 + max_model_len=model_config.max_model_len,

1710 + )

1711

1712 self.config = config

1713 self.model_config = vllm_config.model_config

1714

1715 # Pre-tokenize special tokens for video processing

1716 # to avoid repeated tokenization

1557 - tokenizer = cached_tokenizer_from_config(vllm_config.model_config)

1717 + tokenizer = cached_tokenizer_from_config(model_config)

1718 self._img_start_token_ids = tokenizer.encode(

1719 IMG_START, add_special_tokens=False

1720 )

@@ -1566,7 +1726,10 @@ class NemotronH_Nano_VL_V2(

1726 config

1727 )

1728 if self.dynamic_resolution:

1569 - logger.info("Dynamic resolution is enabled for NanoNemotronVLProcessor")

1729 + logger.info_once(

1730 + "Dynamic resolution is enabled for NanoNemotronVLProcessor",

1731 + scope="global",

1732 + )

1733

1734 def pixel_shuffle(self, x, scale_factor=0.5):

1735 n, w, h, c = x.size()

@@ -1780,6 +1943,51 @@ class NemotronH_Nano_VL_V2(

1943

1944 return final_video_embeddings

1945

1946 + def _process_audio_input(

1947 + self, audio_input: NanoNemotronVLAudioFeatureInputs

1948 + ) -> tuple[torch.Tensor, ...]:

1949 + assert self.sound_encoder is not None

1950 + input_audio_features = audio_input.input_audio_features

1951 + feature_attention_mask = audio_input.feature_attention_mask

1952 + target_device = next(self.sound_encoder.parameters()).device

1953 +

1954 + # When cross-request batching combines audio clips with different

1955 + # time dimensions, _reduce_data returns a list instead of a stacked

1956 + # tensor. Pad to the max time dim and stack; the attention mask

1957 + # already marks valid positions so zero-padding is safe.

1958 + if isinstance(input_audio_features, list):

1959 + feature_sizes = [f.shape[-2] for f in input_audio_features]

1960 + max_t = max(feature_sizes)

1961 + padded_feats = [

1962 + torch.nn.functional.pad(feat, (0, 0, 0, max_t - feat_size))

1963 + for feat, feat_size in zip(

1964 + input_audio_features, feature_sizes, strict=True

1965 + )

1966 + ]

1967 + padded_masks = [

1968 + torch.nn.functional.pad(mask, (0, max_t - mask.shape[-1]))

1969 + for mask in feature_attention_mask

1970 + ]

1971 + input_audio_features = torch.stack(padded_feats)

1972 + feature_attention_mask = torch.stack(padded_masks)

1973 +

1974 + input_audio_features = input_audio_features.to(

1975 + dtype=self.llm_dtype, device=target_device

1976 + )

1977 + feature_attention_mask = feature_attention_mask.to(device=target_device)

1978 + sound_embeds = self.sound_encoder(input_audio_features, feature_attention_mask)

1979 +

1980 + valid_input_lens = feature_attention_mask.sum(dim=1)

1981 + valid_output_lens = self.sound_encoder.encoder._get_subsampling_output_length(

1982 + valid_input_lens

1983 + )

1984 + truncated_embeds = []

1985 + for i in range(sound_embeds.shape[0]):

1986 + valid_len = valid_output_lens[i].item()

1987 + truncated_embeds.append(sound_embeds[i, :valid_len])

1988 +

1989 + return tuple(truncated_embeds)

1990 +

1991 def _create_final_video_embeddings(

1992 self,

1993 video_embeddings: torch.Tensor,

@@ -1887,6 +2095,18 @@ class NemotronH_Nano_VL_V2(

2095 modalities["images"] = self._parse_and_validate_image_input(**kwargs)

2096 if input_key in ("pixel_values_flat_video",) and "videos" not in modalities:

2097 modalities["videos"] = self._parse_and_validate_video_input(**kwargs)

2098 + if (

2099 + input_key

2100 + in (

2101 + "input_audio_features",

2102 + "feature_attention_mask",

2103 + "audio_feature_lengths",

2104 + )

2105 + and "audios" not in modalities

2106 + ):

2107 + modalities["audios"] = NanoNemotronVLAudioFeatureInputs(

2108 + **kwargs, validate=False

2109 + )

2110

2111 return modalities

2112

@@ -1917,6 +2137,10 @@ class NemotronH_Nano_VL_V2(

2137 video_input = modalities["videos"]

2138 video_embeddings = self._process_video_input(video_input)

2139 multimodal_embeddings += tuple(video_embeddings)

2140 + if modality == "audios":

2141 + audio_input = modalities["audios"]

2142 + audio_embeddings = self._process_audio_input(audio_input)

2143 + multimodal_embeddings += tuple(audio_embeddings)

2144

2145 return multimodal_embeddings

2146

@@ -1947,8 +2171,8 @@ class NemotronH_Nano_VL_V2(

2171 """

2172 return MultiModelKeys.from_string_field(

2173 language_model="language_model",

1950 - connector="mlp1",

1951 - tower_model="vision_model",

2174 + connector=["mlp1", "sound_encoder.projection"],

2175 + tower_model=["vision_model", "sound_encoder.encoder"],

2176 )

2177

2178 def compute_logits(

@@ -1969,9 +2193,13 @@ class NemotronH_Nano_VL_V2(

2193 def is_vision_weights(name: str) -> bool:

2194 return name.startswith("vision_model.radio_model.")

2195

2196 + def is_sound_weights(name: str) -> bool:

2197 + return name.startswith("sound")

2198 +

2199 # Separate weights by component

2200 llm_weights = []

2201 vision_weights = []

2202 + sound_weights = []

2203

2204 for name, w in weights:

2205 if is_llm(name):

@@ -1987,9 +2215,15 @@ class NemotronH_Nano_VL_V2(

2215 # Convert: vision_model.radio_model.* → radio_model.*

2216 hf_key = name[len("vision_model.") :] # Remove "vision_model." prefix

2217 vision_weights.append((hf_key, w))

2218 + elif is_sound_weights(name):

2219 + assert self.sound_encoder is not None

2220 + sound_weights.append((name, w))

2221

2222 self.language_model.load_weights(llm_weights)

2223 self.vision_model.load_weights(vision_weights)

2224 + if self.sound_encoder is not None:

2225 + assert len(sound_weights) > 0

2226 + self.sound_encoder.load_weights(sound_weights)

2227

2228 def print_architecture(self, detailed: bool = True, save_to_file: str = None):

2229 """

vllm/model_executor/models/parakeet.py new

+145

@@ -0,0 +1,145 @@

1 +# SPDX-License-Identifier: Apache-2.0

2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 +"""

4 +Modules below used for the audio encoder component in: models/nano_nemotron_vl.py

5 +"""

6 +

7 +from collections.abc import Iterable

8 +from dataclasses import asdict

9 +

10 +import numpy as np

11 +import torch

12 +import torch.nn as nn

13 +from transformers import ParakeetEncoder as HFParakeetEncoder

14 +from transformers import ParakeetFeatureExtractor, PretrainedConfig

15 +

16 +from vllm.model_executor.layers.activation import ReLUSquaredActivation

17 +from vllm.model_executor.model_loader.weight_utils import default_weight_loader

18 +from vllm.transformers_utils.configs.parakeet import ExtractorConfig, ParakeetConfig

19 +

20 +

21 +class ParakeetProjection(nn.Module):

22 + def __init__(self, config: ParakeetConfig) -> None:

23 + super().__init__()

24 + sound_hidden_size = config.hidden_size

25 + proj_hidden_size = config.projection_hidden_size

26 + llm_hidden_size = config.llm_hidden_size

27 + bias = config.projection_bias

28 +

29 + self.norm = nn.LayerNorm(sound_hidden_size, eps=config.projection_eps)

30 + self.linear1 = nn.Linear(sound_hidden_size, proj_hidden_size, bias=bias)

31 + self.activation = ReLUSquaredActivation()

32 + self.linear2 = nn.Linear(proj_hidden_size, llm_hidden_size, bias=bias)

33 +

34 + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:

35 + hidden_states = self.norm(hidden_states)

36 + hidden_states = self.linear1(hidden_states)

37 + hidden_states = self.activation(hidden_states)

38 + hidden_states = self.linear2(hidden_states)

39 + return hidden_states

40 +

41 +

42 +class ProjectedParakeet(nn.Module):

43 + def __init__(

44 + self,

45 + config: PretrainedConfig,

46 + *,

47 + dtype: torch.dtype,

48 + llm_hidden_size: int,

49 + max_model_len: int,

50 + ) -> None:

51 + super().__init__()

52 + self.config = ParakeetConfig.from_hf_config(

53 + config, llm_hidden_size=llm_hidden_size, max_model_len=max_model_len

54 + )

55 + self.encoder = HFParakeetEncoder(self.config)

56 + self.encoder = self.encoder.to(dtype)

57 + self.projection = ParakeetProjection(self.config)

58 + self.projection = self.projection.to(dtype)

59 +

60 + def forward(

61 + self, input_features: torch.Tensor, attention_mask: torch.Tensor | None = None

62 + ) -> torch.Tensor:

63 + outputs = self.encoder(

64 + input_features=input_features, attention_mask=attention_mask

65 + )

66 + outputs = outputs.last_hidden_state

67 + outputs = outputs.to(dtype=torch.bfloat16)

68 + outputs = self.projection(outputs)

69 + return outputs

70 +

71 + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:

72 + loaded_params: set[str] = set()

73 + params_dict = dict(self.named_parameters())

74 + buffers_dict = dict(self.named_buffers())

75 +

76 + if isinstance(weights, dict):

77 + weights_list = list(weights.items())

78 + else:

79 + weights_list = list(weights)

80 +

81 + for name, weight in weights_list:

82 + if name.startswith("sound_encoder.encoder.feature_extractor."):

83 + # Feature extractor buffers are handled outside the encoder.

84 + continue

85 + if name.startswith("sound_encoder."):

86 + target_name = name[len("sound_encoder.") :]

87 + elif name.startswith("sound_projection."):

88 + target_name = f"projection.{name[len('sound_projection.') :]}"

89 + else:

90 + continue

91 +

92 + target = params_dict.get(target_name)

93 + if target is None:

94 + target = buffers_dict.get(target_name)

95 + if target is None:

96 + raise ValueError(f"Unknown weight: {name}")

97 + weight_loader = getattr(target, "weight_loader", default_weight_loader)

98 + with torch.no_grad():

99 + weight_loader(target, weight)

100 + loaded_params.add(target_name)

101 +

102 + return loaded_params

103 +

104 +

105 +class ParakeetExtractor(ParakeetFeatureExtractor):

106 + def __init__(self, config: PretrainedConfig) -> None:

107 + self.config = ExtractorConfig.from_hf_config(config)

108 + super().__init__(**asdict(self.config))

109 + self._clip_target_samples = int(

110 + round(self.config.clip_duration_s * self.sampling_rate)

111 + )

112 + self._tail_min_samples = int(

113 + round(self.config.clip_min_duration_s * self.sampling_rate)

114 + )

115 +

116 + def _normalize_audio_length(self, audio_len: int) -> int:

117 + # Match mcore's compute_params() logic for clip/minduration handling.

118 + target_len = max(audio_len, self._tail_min_samples)

119 + tail_remainder = target_len % self._clip_target_samples

120 + if 0 < tail_remainder < self._tail_min_samples:

121 + padding = self._tail_min_samples - tail_remainder

122 + target_len += padding

123 + assert isinstance(target_len, int)

124 + return target_len

125 +

126 + def audio_token_count(self, audio_len: int) -> int:

127 + audio_len = self._normalize_audio_length(audio_len)

128 + num_frames = audio_len // self.hop_length

129 + n_tokens = HFParakeetEncoder._get_subsampling_output_length(

130 + self, torch.tensor([num_frames], dtype=torch.float)

131 + )

132 + return max(1, n_tokens.item())

133 +

134 + def __call__(self, raw_speech: list[np.ndarray], *args, **kwargs):

135 + padded = []

136 + for p in raw_speech:

137 + assert p.ndim == 1

138 + audio_len = int(p.shape[0])

139 + target_len = self._normalize_audio_length(audio_len)

140 + p = np.pad(p, (0, target_len - audio_len))

141 + padded.append(p)

142 + return super().__call__(padded, *args, **kwargs)

143 +

144 + def audio_length(self, audio_tokens: int) -> int:

145 + return int(audio_tokens * self.config.subsampling_factor * self.hop_length)

vllm/transformers_utils/configs/parakeet.py new

+49

@@ -0,0 +1,49 @@

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 + llm_hidden_size: int

10 + projection_hidden_size: int

11 + projection_bias: bool

12 + projection_eps: float = 1e-5

13 + sampling_rate: int

14 +

15 + @staticmethod

16 + def from_hf_config(

17 + config: PretrainedConfig, *, llm_hidden_size: int, max_model_len: int

18 + ) -> "ParakeetConfig":

19 + assert isinstance(config, PretrainedConfig)

20 + return ParakeetConfig(

21 + **config.to_dict(),

22 + scale_input=False,

23 + attention_bias=False,

24 + llm_hidden_size=llm_hidden_size,

25 + max_position_embeddings=max_model_len

26 + + 1, # + 1 because it seems like max_model_len+1 can be passed

27 + )

28 +

29 +

30 +@dataclass(kw_only=True, frozen=True)

31 +class ExtractorConfig:

32 + feature_size: int

33 + sampling_rate: int

34 + subsampling_factor: int

35 + subsampling_conv_kernel_size: int

36 + subsampling_conv_stride: int

37 + clip_duration_s: int = 30

38 + clip_min_duration_s: float = 0.1

39 +

40 + @staticmethod

41 + def from_hf_config(config: PretrainedConfig) -> "ExtractorConfig":

42 + assert isinstance(config, PretrainedConfig)

43 + return ExtractorConfig(

44 + feature_size=config.num_mel_bins,

45 + sampling_rate=config.sampling_rate,

46 + subsampling_factor=config.subsampling_factor,

47 + subsampling_conv_kernel_size=config.subsampling_conv_kernel_size,

48 + subsampling_conv_stride=config.subsampling_conv_stride,

49 + )