wehub-resource-sync
80 lines
2.8 KiB
Python
80 lines
2.8 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
|
"""Inference-only IBM Granite Speech Plus model."""
|
|
|
|
import torch
|
|
from transformers import PretrainedConfig
|
|
|
|
from vllm.model_executor.layers.quantization import QuantizationConfig
|
|
from vllm.multimodal import MULTIMODAL_REGISTRY
|
|
|
|
from .granite_speech import (
|
|
GraniteSpeechCTCEncoder,
|
|
GraniteSpeechDummyInputsBuilder,
|
|
GraniteSpeechForConditionalGeneration,
|
|
GraniteSpeechMultiModalProcessingInfo,
|
|
GraniteSpeechMultiModalProcessor,
|
|
)
|
|
|
|
ISO639_1_SUPPORTED_LANGS = {
|
|
"en": "English",
|
|
"fr": "French",
|
|
"de": "German",
|
|
"pt": "Portuguese",
|
|
"es": "Spanish",
|
|
}
|
|
|
|
|
|
class GraniteSpeechPlusCTCEncoder(GraniteSpeechCTCEncoder):
|
|
def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
|
|
hidden_states = self.input_linear(hidden_states)
|
|
# cat_hidden_layers selects non-negative layer indices (0 = encoder
|
|
# input, N = output of layer N) whose hidden states are concatenated
|
|
# along the feature dim *in addition to* the final hidden states,
|
|
# which are always appended last.
|
|
cat_layers = set(self.config.cat_hidden_layers or [])
|
|
exported_hidden_states = []
|
|
|
|
if 0 in cat_layers:
|
|
exported_hidden_states.append(hidden_states)
|
|
|
|
for idx, layer in enumerate(self.layers, start=1):
|
|
hidden_states = layer(hidden_states, attention_dists=self.attention_dists)
|
|
|
|
# Skip the final layer here since its output is always appended
|
|
# below; capturing it twice would double-append.
|
|
if idx in cat_layers and idx != self.num_layers:
|
|
exported_hidden_states.append(hidden_states)
|
|
|
|
if idx == self.num_layers // 2:
|
|
hidden_states_mid = hidden_states.clone()
|
|
hidden_states_mid, _ = self.out(hidden_states_mid)
|
|
hidden_states_mid = self.softmax(hidden_states_mid)
|
|
hidden_states_mid, _ = self.out_mid(hidden_states_mid)
|
|
hidden_states += hidden_states_mid
|
|
|
|
if exported_hidden_states:
|
|
hidden_states = torch.cat([*exported_hidden_states, hidden_states], dim=-1)
|
|
return hidden_states
|
|
|
|
|
|
@MULTIMODAL_REGISTRY.register_processor(
|
|
GraniteSpeechMultiModalProcessor,
|
|
info=GraniteSpeechMultiModalProcessingInfo,
|
|
dummy_inputs=GraniteSpeechDummyInputsBuilder,
|
|
)
|
|
class GraniteSpeechPlusForConditionalGeneration(GraniteSpeechForConditionalGeneration):
|
|
supported_languages = ISO639_1_SUPPORTED_LANGS
|
|
|
|
def _build_encoder(
|
|
self,
|
|
config: PretrainedConfig,
|
|
quant_config: QuantizationConfig | None,
|
|
prefix: str,
|
|
) -> GraniteSpeechCTCEncoder:
|
|
return GraniteSpeechPlusCTCEncoder(
|
|
config=config,
|
|
quant_config=quant_config,
|
|
prefix=prefix,
|
|
)
|