| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """Inference-only IBM Granite Speech Plus model.""" |
| 4 | |
| 5 | import torch |
| 6 | from transformers import PreTrainedConfig |
| 7 | |
| 8 | from vllm.model_executor.layers.quantization import QuantizationConfig |
| 9 | from vllm.multimodal import MULTIMODAL_REGISTRY |
| 10 | |
| 11 | from .granite_speech import ( |
| 12 | GraniteSpeechCTCEncoder, |
| 13 | GraniteSpeechDummyInputsBuilder, |
| 14 | GraniteSpeechForConditionalGeneration, |
| 15 | GraniteSpeechMultiModalProcessor, |
| 16 | GraniteSpeechProcessingInfo, |
| 17 | ) |
| 18 | |
| 19 | ISO639_1_SUPPORTED_LANGS = { |
| 20 | "en": "English", |
| 21 | "fr": "French", |
| 22 | "de": "German", |
| 23 | "pt": "Portuguese", |
| 24 | "es": "Spanish", |
| 25 | } |
| 26 | |
| 27 | |
| 28 | class GraniteSpeechPlusCTCEncoder(GraniteSpeechCTCEncoder): |
| 29 | def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: |
| 30 | hidden_states = self.input_linear(hidden_states) |
| 31 | # cat_hidden_layers selects non-negative layer indices (0 = encoder |
| 32 | # input, N = output of layer N) whose hidden states are concatenated |
| 33 | # along the feature dim *in addition to* the final hidden states, |
| 34 | # which are always appended last. |
| 35 | cat_layers = set(self.config.cat_hidden_layers or []) |
| 36 | exported_hidden_states = [] |
| 37 | |
| 38 | if 0 in cat_layers: |
| 39 | exported_hidden_states.append(hidden_states) |
| 40 | |
| 41 | for idx, layer in enumerate(self.layers, start=1): |
| 42 | hidden_states = layer(hidden_states, attention_dists=self.attention_dists) |
| 43 | |
| 44 | # Skip the final layer here since its output is always appended |
| 45 | # below; capturing it twice would double-append. |
| 46 | if idx in cat_layers and idx != self.num_layers: |
| 47 | exported_hidden_states.append(hidden_states) |
| 48 | |
| 49 | if idx == self.num_layers // 2: |
| 50 | hidden_states_mid = hidden_states.clone() |
| 51 | hidden_states_mid, _ = self.out(hidden_states_mid) |
| 52 | hidden_states_mid = self.softmax(hidden_states_mid) |
| 53 | hidden_states_mid, _ = self.out_mid(hidden_states_mid) |
| 54 | hidden_states += hidden_states_mid |
| 55 | |
| 56 | if exported_hidden_states: |
| 57 | hidden_states = torch.cat([*exported_hidden_states, hidden_states], dim=-1) |
| 58 | return hidden_states |
| 59 | |
| 60 | |
| 61 | @MULTIMODAL_REGISTRY.register_processor( |
| 62 | GraniteSpeechMultiModalProcessor, |
| 63 | info=GraniteSpeechProcessingInfo, |
| 64 | dummy_inputs=GraniteSpeechDummyInputsBuilder, |
| 65 | ) |
| 66 | class GraniteSpeechPlusForConditionalGeneration(GraniteSpeechForConditionalGeneration): |
| 67 | supported_languages = ISO639_1_SUPPORTED_LANGS |
| 68 | |
| 69 | def _build_encoder( |
| 70 | self, |
| 71 | config: PreTrainedConfig, |
| 72 | quant_config: QuantizationConfig | None, |
| 73 | prefix: str, |
| 74 | ) -> GraniteSpeechCTCEncoder: |
| 75 | return GraniteSpeechPlusCTCEncoder( |
| 76 | config=config, |
| 77 | quant_config=quant_config, |
| 78 | prefix=prefix, |
| 79 | ) |