wehub-resource-sync
793 lines
26 KiB
Python
793 lines
26 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||
"""
|
||
FireRedLID – Language Identification model adapted for vLLM.
|
||
|
||
Architecture: ConformerEncoder + TransformerDecoder (6-layer cross-attn)
|
||
Vocabulary: 120 LID tokens (dict.txt)
|
||
Output: Up to 2 tokens (e.g. "en", "zh mandarin")
|
||
|
||
This implementation follows the Whisper-style encoder-decoder pattern:
|
||
• Encoder processes audio features (Fbank + CMVN via FeatureExtractor)
|
||
• Decoder performs single-step autoregressive forward
|
||
• vLLM's generation loop handles beam search / sampling
|
||
"""
|
||
|
||
from __future__ import annotations
|
||
|
||
from collections.abc import Iterable, Mapping, Sequence
|
||
from typing import Annotated, Literal
|
||
|
||
import numpy as np
|
||
import torch
|
||
from torch import nn
|
||
from transformers import BatchFeature
|
||
|
||
from vllm.config import ModelConfig, VllmConfig
|
||
from vllm.config.multimodal import BaseDummyOptions
|
||
from vllm.config.speech_to_text import SpeechToTextConfig
|
||
from vllm.distributed import get_tensor_model_parallel_world_size
|
||
from vllm.inputs import MultiModalDataDict, PromptType
|
||
from vllm.logger import init_logger
|
||
from vllm.model_executor.layers.attention import Attention, CrossAttention
|
||
from vllm.model_executor.layers.linear import (
|
||
ColumnParallelLinear,
|
||
ReplicatedLinear,
|
||
RowParallelLinear,
|
||
)
|
||
from vllm.model_executor.layers.logits_processor import LogitsProcessor
|
||
from vllm.model_executor.layers.vocab_parallel_embedding import ParallelLMHead
|
||
from vllm.multimodal import MULTIMODAL_REGISTRY
|
||
from vllm.multimodal.inputs import (
|
||
MultiModalFieldConfig,
|
||
MultiModalKwargsItems,
|
||
)
|
||
from vllm.multimodal.parse import MultiModalDataItems, MultiModalDataParser
|
||
from vllm.multimodal.processing import (
|
||
BaseDummyInputsBuilder,
|
||
BaseProcessingInfo,
|
||
EncDecMultiModalProcessor,
|
||
PromptReplacement,
|
||
PromptUpdate,
|
||
)
|
||
from vllm.transformers_utils.processor import cached_processor_from_config
|
||
from vllm.utils.tensor_schema import TensorSchema, TensorShape
|
||
|
||
from .conformer_encoder import ConformerEncoder
|
||
from .interfaces import (
|
||
MultiModalEmbeddings,
|
||
SupportsMultiModal,
|
||
SupportsTranscription,
|
||
)
|
||
from .utils import (
|
||
AutoWeightsLoader,
|
||
WeightsMapper,
|
||
maybe_prefix,
|
||
)
|
||
from .whisper_utils import ISO639_1_SUPPORTED_LANGS
|
||
|
||
logger = init_logger(__name__)
|
||
|
||
|
||
class FireRedLIDAudioInputs(TensorSchema):
|
||
"""
|
||
Dimensions:
|
||
- b: Batch size
|
||
- t: Time frames (variable across utterances)
|
||
- nmb: Number of mel bins (80)
|
||
"""
|
||
|
||
input_features: Annotated[
|
||
list[torch.Tensor] | None,
|
||
TensorShape("b", "t", "nmb", dynamic_dims={"t"}),
|
||
]
|
||
speech_lengths: Annotated[
|
||
list[torch.Tensor] | None,
|
||
TensorShape("b"),
|
||
]
|
||
fake_token_lengths: Annotated[
|
||
list[torch.Tensor] | None,
|
||
TensorShape("b"),
|
||
]
|
||
|
||
|
||
FireRedLIDEncoder = ConformerEncoder
|
||
|
||
|
||
class FireRedLIDPositionalEmbedding(nn.Module):
|
||
"""Absolute sinusoidal positional embedding indexed by `positions`."""
|
||
|
||
def __init__(self, d_model: int, max_len: int = 5000):
|
||
super().__init__()
|
||
assert d_model % 2 == 0
|
||
pe = torch.zeros(max_len, d_model, requires_grad=False)
|
||
position = torch.arange(0, max_len).unsqueeze(1).float()
|
||
div_term = torch.exp(
|
||
torch.arange(0, d_model, 2).float()
|
||
* -(torch.log(torch.tensor(10000.0)).item() / d_model)
|
||
)
|
||
pe[:, 0::2] = torch.sin(position * div_term)
|
||
pe[:, 1::2] = torch.cos(position * div_term)
|
||
self.register_buffer("pe", pe, persistent=False)
|
||
|
||
def forward(self, position_ids: torch.Tensor) -> torch.Tensor:
|
||
return self.pe[position_ids]
|
||
|
||
|
||
class FireRedLIDAttention(nn.Module):
|
||
"""Base attention with shared QKV/FC projections for the LID decoder."""
|
||
|
||
def __init__(
|
||
self,
|
||
d_model: int,
|
||
n_head: int,
|
||
*,
|
||
vllm_config: VllmConfig,
|
||
prefix: str = "",
|
||
):
|
||
super().__init__()
|
||
tp_size = get_tensor_model_parallel_world_size()
|
||
assert n_head % tp_size == 0
|
||
self.total_num_heads = n_head
|
||
self.num_heads = n_head // tp_size
|
||
self.num_kv_heads = max(1, n_head // tp_size)
|
||
self.head_dim = d_model // n_head
|
||
self.scaling = self.head_dim**-0.5
|
||
|
||
cache_config = vllm_config.cache_config
|
||
quant_config = vllm_config.quant_config
|
||
|
||
self.w_qs = ColumnParallelLinear(
|
||
d_model,
|
||
d_model,
|
||
bias=True,
|
||
quant_config=quant_config,
|
||
prefix=f"{prefix}.w_qs",
|
||
)
|
||
self.w_ks = ColumnParallelLinear(
|
||
d_model,
|
||
d_model,
|
||
bias=False,
|
||
quant_config=quant_config,
|
||
prefix=f"{prefix}.w_ks",
|
||
)
|
||
self.w_vs = ColumnParallelLinear(
|
||
d_model,
|
||
d_model,
|
||
bias=True,
|
||
quant_config=quant_config,
|
||
prefix=f"{prefix}.w_vs",
|
||
)
|
||
self.fc = RowParallelLinear(
|
||
d_model,
|
||
d_model,
|
||
bias=True,
|
||
quant_config=quant_config,
|
||
prefix=f"{prefix}.fc",
|
||
)
|
||
self._init_attn(cache_config, quant_config, prefix)
|
||
|
||
def _init_attn(self, cache_config, quant_config, prefix: str) -> None:
|
||
raise NotImplementedError
|
||
|
||
|
||
class FireRedLIDSelfAttention(FireRedLIDAttention):
|
||
def _init_attn(self, cache_config, quant_config, prefix: str) -> None:
|
||
self.attn = Attention(
|
||
self.num_heads,
|
||
self.head_dim,
|
||
self.scaling,
|
||
num_kv_heads=self.num_kv_heads,
|
||
cache_config=cache_config,
|
||
quant_config=quant_config,
|
||
prefix=f"{prefix}.attn",
|
||
)
|
||
|
||
def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
|
||
q, _ = self.w_qs(hidden_states)
|
||
k, _ = self.w_ks(hidden_states)
|
||
v, _ = self.w_vs(hidden_states)
|
||
attn_output = self.attn(q, k, v)
|
||
output, _ = self.fc(attn_output)
|
||
return output
|
||
|
||
|
||
class FireRedLIDCrossAttention(FireRedLIDAttention):
|
||
def _init_attn(self, cache_config, quant_config, prefix: str) -> None:
|
||
self.attn = CrossAttention(
|
||
self.num_heads,
|
||
self.head_dim,
|
||
self.scaling,
|
||
num_kv_heads=self.num_kv_heads,
|
||
cache_config=cache_config,
|
||
quant_config=quant_config,
|
||
prefix=f"{prefix}.attn",
|
||
)
|
||
|
||
def forward(
|
||
self,
|
||
hidden_states: torch.Tensor,
|
||
encoder_hidden_states: torch.Tensor | None,
|
||
) -> torch.Tensor:
|
||
q, _ = self.w_qs(hidden_states)
|
||
if encoder_hidden_states is not None:
|
||
k, _ = self.w_ks(encoder_hidden_states)
|
||
v, _ = self.w_vs(encoder_hidden_states)
|
||
else:
|
||
k = v = None
|
||
|
||
attn_output = self.attn(q, k, v)
|
||
output, _ = self.fc(attn_output)
|
||
return output
|
||
|
||
|
||
class FireRedLIDFFN(nn.Module):
|
||
def __init__(self, d_model: int, d_ff: int):
|
||
super().__init__()
|
||
self.w_1 = ReplicatedLinear(d_model, d_ff, bias=True)
|
||
self.act = nn.GELU()
|
||
self.w_2 = ReplicatedLinear(d_ff, d_model, bias=True)
|
||
|
||
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
||
x, _ = self.w_1(x)
|
||
x = self.act(x)
|
||
x, _ = self.w_2(x)
|
||
return x
|
||
|
||
|
||
class FireRedLIDDecoderLayer(nn.Module):
|
||
"""vLLM-native decoder layer while preserving FireRedLID parameter names."""
|
||
|
||
def __init__(
|
||
self,
|
||
d_model: int,
|
||
n_head: int,
|
||
*,
|
||
vllm_config: VllmConfig,
|
||
prefix: str = "",
|
||
):
|
||
super().__init__()
|
||
self.self_attn_norm = nn.LayerNorm(d_model)
|
||
self.self_attn = FireRedLIDSelfAttention(
|
||
d_model,
|
||
n_head,
|
||
vllm_config=vllm_config,
|
||
prefix=f"{prefix}.self_attn",
|
||
)
|
||
|
||
self.cross_attn_norm = nn.LayerNorm(d_model)
|
||
self.cross_attn = FireRedLIDCrossAttention(
|
||
d_model,
|
||
n_head,
|
||
vllm_config=vllm_config,
|
||
prefix=f"{prefix}.cross_attn",
|
||
)
|
||
|
||
self.mlp_norm = nn.LayerNorm(d_model)
|
||
self.mlp = FireRedLIDFFN(d_model, d_model * 4)
|
||
|
||
def forward(
|
||
self,
|
||
hidden_states: torch.Tensor,
|
||
encoder_hidden_states: torch.Tensor | None,
|
||
) -> torch.Tensor:
|
||
residual = hidden_states
|
||
hidden_states = self.self_attn_norm(hidden_states)
|
||
hidden_states = self.self_attn(hidden_states)
|
||
hidden_states = residual + hidden_states
|
||
|
||
residual = hidden_states
|
||
hidden_states = self.cross_attn_norm(hidden_states)
|
||
hidden_states = self.cross_attn(hidden_states, encoder_hidden_states)
|
||
hidden_states = residual + hidden_states
|
||
|
||
residual = hidden_states
|
||
hidden_states = self.mlp_norm(hidden_states)
|
||
hidden_states = residual + self.mlp(hidden_states)
|
||
|
||
return hidden_states
|
||
|
||
|
||
class FireRedLIDDecoder(nn.Module):
|
||
def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
|
||
super().__init__()
|
||
config = vllm_config.model_config.hf_config
|
||
self.pad_id = getattr(config, "pad_token_id", 2)
|
||
self.n_layers = getattr(config, "n_layers_lid_dec", 6)
|
||
self.d_model = getattr(config, "d_model", 1280)
|
||
self.scale = self.d_model**0.5
|
||
|
||
self.tgt_word_emb = nn.Embedding(
|
||
getattr(config, "vocab_size", 120),
|
||
self.d_model,
|
||
padding_idx=self.pad_id,
|
||
)
|
||
self.positional_encoding = FireRedLIDPositionalEmbedding(
|
||
self.d_model,
|
||
max_len=getattr(config, "pe_maxlen", 5000),
|
||
)
|
||
|
||
self.layer_stack = nn.ModuleList(
|
||
[
|
||
FireRedLIDDecoderLayer(
|
||
self.d_model,
|
||
getattr(config, "n_head", 20),
|
||
vllm_config=vllm_config,
|
||
prefix=f"{prefix}.layer_stack.{idx}",
|
||
)
|
||
for idx in range(self.n_layers)
|
||
]
|
||
)
|
||
self.layer_norm_out = nn.LayerNorm(self.d_model)
|
||
|
||
def forward(
|
||
self,
|
||
input_ids: torch.Tensor,
|
||
positions: torch.Tensor,
|
||
encoder_hidden_states: torch.Tensor | None,
|
||
) -> torch.Tensor:
|
||
hidden_states = self.tgt_word_emb(input_ids) * self.scale
|
||
hidden_states = hidden_states + self.positional_encoding(positions)
|
||
|
||
for layer in self.layer_stack:
|
||
hidden_states = layer(hidden_states, encoder_hidden_states)
|
||
|
||
hidden_states = self.layer_norm_out(hidden_states)
|
||
return hidden_states
|
||
|
||
def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
|
||
return self.tgt_word_emb(input_ids)
|
||
|
||
|
||
class FireRedLIDModel(nn.Module):
|
||
def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
|
||
super().__init__()
|
||
config = vllm_config.model_config.hf_config
|
||
|
||
self.encoder = FireRedLIDEncoder(
|
||
idim=getattr(config, "idim", 80),
|
||
n_layers_enc=getattr(config, "n_layers_enc", 16),
|
||
n_head=getattr(config, "n_head", 20),
|
||
d_model=getattr(config, "d_model", 1280),
|
||
kernel_size=getattr(config, "kernel_size", 33),
|
||
pe_maxlen=getattr(config, "pe_maxlen", 5000),
|
||
)
|
||
|
||
self.decoder = FireRedLIDDecoder(
|
||
vllm_config=vllm_config,
|
||
prefix=maybe_prefix(prefix, "decoder"),
|
||
)
|
||
|
||
def forward(
|
||
self,
|
||
input_ids: torch.Tensor,
|
||
positions: torch.Tensor,
|
||
encoder_outputs: list[torch.Tensor] | None = None,
|
||
) -> torch.Tensor:
|
||
enc_states = (
|
||
torch.cat(encoder_outputs, dim=0)
|
||
if encoder_outputs and len(encoder_outputs) > 0
|
||
else None
|
||
)
|
||
decoder_outputs = self.decoder(
|
||
input_ids=input_ids,
|
||
positions=positions,
|
||
encoder_hidden_states=enc_states,
|
||
)
|
||
return decoder_outputs
|
||
|
||
def get_encoder_outputs(
|
||
self,
|
||
speech: torch.Tensor | list[torch.Tensor],
|
||
speech_lengths: torch.Tensor | list[torch.Tensor],
|
||
) -> tuple[torch.Tensor, torch.Tensor]:
|
||
"""Run the encoder and return padded outputs plus true sequence lengths."""
|
||
enc_output, enc_lengths, _ = self.encoder(speech, speech_lengths)
|
||
return enc_output, enc_lengths
|
||
|
||
|
||
class FireRedLIDProcessingInfo(BaseProcessingInfo):
|
||
def get_hf_config(self):
|
||
return self.ctx.get_hf_config()
|
||
|
||
def get_supported_mm_limits(self) -> Mapping[str, int | None]:
|
||
return {"audio": 1}
|
||
|
||
def get_feature_extractor(self, **kwargs):
|
||
hf_processor = self.get_hf_processor(**kwargs)
|
||
feature_extractor = hf_processor.feature_extractor
|
||
return feature_extractor
|
||
|
||
def get_data_parser(self) -> MultiModalDataParser:
|
||
feature_extractor = self.get_feature_extractor()
|
||
return MultiModalDataParser(
|
||
target_sr=feature_extractor.sampling_rate,
|
||
target_channels=1,
|
||
)
|
||
|
||
@property
|
||
def skip_prompt_length_check(self) -> bool:
|
||
return True
|
||
|
||
def get_num_audio_tokens(self) -> int:
|
||
# For encoder profiling – return a reasonable dummy length.
|
||
# This doesn't affect actual inference since encoder processes
|
||
# variable-length features.
|
||
return 1
|
||
|
||
|
||
class FireRedLIDDummyInputsBuilder(BaseDummyInputsBuilder[FireRedLIDProcessingInfo]):
|
||
def get_dummy_text(self, mm_counts: Mapping[str, int]) -> str:
|
||
return "<sos>"
|
||
|
||
def get_dummy_mm_data(
|
||
self,
|
||
seq_len: int,
|
||
mm_counts: Mapping[str, int],
|
||
mm_options: Mapping[str, BaseDummyOptions],
|
||
) -> MultiModalDataDict:
|
||
feature_extractor = self.info.get_feature_extractor()
|
||
sampling_rate = feature_extractor.sampling_rate
|
||
audio_len = feature_extractor.chunk_length * sampling_rate
|
||
num_audios = mm_counts.get("audio", 0)
|
||
audio_overrides = mm_options.get("audio")
|
||
return {
|
||
"audio": self._get_dummy_audios(
|
||
length=audio_len,
|
||
num_audios=num_audios,
|
||
overrides=audio_overrides,
|
||
)
|
||
}
|
||
|
||
|
||
class FireRedLIDMultiModalProcessor(
|
||
EncDecMultiModalProcessor[FireRedLIDProcessingInfo]
|
||
):
|
||
def create_encoder_prompt(
|
||
self,
|
||
prompt: str | list[int],
|
||
mm_items: MultiModalDataItems,
|
||
) -> str | list[int]:
|
||
# Dummy encoder prompt for profiling (encoder only processes audio).
|
||
return [0]
|
||
|
||
def _call_hf_processor(
|
||
self,
|
||
prompt: str,
|
||
mm_data: Mapping[str, object],
|
||
mm_kwargs: Mapping[str, object],
|
||
tok_kwargs: Mapping[str, object],
|
||
) -> BatchFeature:
|
||
if mm_data:
|
||
feature_extractor = self.info.get_feature_extractor(**mm_kwargs)
|
||
mm_data = dict(audio=mm_data.pop("audios"))
|
||
mm_kwargs = dict(
|
||
**mm_kwargs,
|
||
sampling_rate=feature_extractor.sampling_rate,
|
||
)
|
||
processed_outputs = super()._call_hf_processor(
|
||
prompt=prompt,
|
||
mm_data=mm_data,
|
||
mm_kwargs=mm_kwargs,
|
||
tok_kwargs=tok_kwargs,
|
||
)
|
||
if "labels" in processed_outputs:
|
||
processed_outputs["input_ids"] = processed_outputs.pop("labels")
|
||
return processed_outputs
|
||
|
||
def _get_mm_fields_config(
|
||
self,
|
||
hf_inputs: BatchFeature,
|
||
hf_processor_mm_kwargs: Mapping[str, object],
|
||
) -> Mapping[str, MultiModalFieldConfig]:
|
||
return dict(
|
||
input_features=MultiModalFieldConfig.batched("audio"),
|
||
speech_lengths=MultiModalFieldConfig.batched("audio"),
|
||
fake_token_lengths=MultiModalFieldConfig.batched("audio"),
|
||
)
|
||
|
||
def _get_prompt_updates(
|
||
self,
|
||
mm_items: MultiModalDataItems,
|
||
hf_processor_mm_kwargs: Mapping[str, object],
|
||
out_mm_kwargs: MultiModalKwargsItems,
|
||
) -> Sequence[PromptUpdate]:
|
||
out_mm_data = out_mm_kwargs.get_data()
|
||
fake_token_lengths = out_mm_data.get("fake_token_lengths")
|
||
|
||
if fake_token_lengths is None:
|
||
# Fallback to max encoder output length if not available
|
||
audio_output_lengths = []
|
||
else:
|
||
assert isinstance(fake_token_lengths, torch.Tensor)
|
||
audio_output_lengths = fake_token_lengths.tolist()
|
||
|
||
def get_replacement(item_idx: int):
|
||
if audio_output_lengths:
|
||
num_tokens = int(audio_output_lengths[item_idx])
|
||
else:
|
||
num_tokens = self.info.get_num_audio_tokens()
|
||
return [0] * num_tokens
|
||
|
||
return [
|
||
PromptReplacement(
|
||
modality="audio",
|
||
target=[0],
|
||
replacement=get_replacement,
|
||
)
|
||
]
|
||
|
||
|
||
# FireRedLID supports a wider set of languages than Whisper's shared list.
|
||
# Only ISO 639-1 codes are listed; FireRedLID's dialect tokens (mandarin,
|
||
# xinan, wu, …) are output tokens but not valid language *request* codes.
|
||
_FIREREDLID_SUPPORTED_LANGUAGES: Mapping[str, str] = {
|
||
**ISO639_1_SUPPORTED_LANGS,
|
||
"am": "Amharic",
|
||
"as": "Assamese",
|
||
"ba": "Bashkir",
|
||
"bn": "Bengali",
|
||
"bo": "Tibetan",
|
||
"br": "Breton",
|
||
"eu": "Basque",
|
||
"fo": "Faroese",
|
||
"gu": "Gujarati",
|
||
"ha": "Hausa",
|
||
"haw": "Hawaiian",
|
||
"ht": "Haitian Creole",
|
||
"jw": "Javanese",
|
||
"ka": "Georgian",
|
||
"km": "Khmer",
|
||
"la": "Latin",
|
||
"lb": "Luxembourgish",
|
||
"ln": "Lingala",
|
||
"lo": "Lao",
|
||
"mg": "Malagasy",
|
||
"ml": "Malayalam",
|
||
"mn": "Mongolian",
|
||
"mt": "Maltese",
|
||
"my": "Myanmar",
|
||
"nn": "Nynorsk",
|
||
"oc": "Occitan",
|
||
"pa": "Panjabi",
|
||
"ps": "Pashto",
|
||
"sa": "Sanskrit",
|
||
"sd": "Sindhi",
|
||
"si": "Sinhala",
|
||
"sn": "Shona",
|
||
"so": "Somali",
|
||
"sq": "Albanian",
|
||
"su": "Sundanese",
|
||
"te": "Telugu",
|
||
"tg": "Tajik",
|
||
"tk": "Turkmen",
|
||
"tt": "Tatar",
|
||
"uz": "Uzbek",
|
||
"yi": "Yiddish",
|
||
"yo": "Yoruba",
|
||
"yue": "Cantonese",
|
||
}
|
||
|
||
|
||
@MULTIMODAL_REGISTRY.register_processor(
|
||
FireRedLIDMultiModalProcessor,
|
||
info=FireRedLIDProcessingInfo,
|
||
dummy_inputs=FireRedLIDDummyInputsBuilder,
|
||
)
|
||
class FireRedLIDForConditionalGeneration(
|
||
nn.Module, SupportsTranscription, SupportsMultiModal
|
||
):
|
||
# -- SupportsTranscription protocol attributes --
|
||
supports_transcription_only = True
|
||
supported_languages = _FIREREDLID_SUPPORTED_LANGUAGES
|
||
|
||
hf_to_vllm_mapper = WeightsMapper(
|
||
orig_to_new_substr={
|
||
"encoder.": "model.encoder.",
|
||
"lid_decoder.": "model.decoder.",
|
||
# Encoder FFN: nn.Sequential indices → named children
|
||
"net.0": "pre_layer_norm",
|
||
"net.1": "linear_expand",
|
||
"net.4": "linear_project",
|
||
}
|
||
)
|
||
|
||
def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
|
||
super().__init__()
|
||
config = vllm_config.model_config.hf_config
|
||
self.config = config
|
||
self.dtype = vllm_config.model_config.dtype
|
||
|
||
with self._mark_composite_model(
|
||
vllm_config,
|
||
language_targets=FireRedLIDDecoder,
|
||
tower_targets={"audio": FireRedLIDEncoder},
|
||
):
|
||
self.model = FireRedLIDModel(
|
||
vllm_config=vllm_config,
|
||
prefix=maybe_prefix(prefix, "model"),
|
||
)
|
||
|
||
self.proj_out = ParallelLMHead(
|
||
getattr(config, "vocab_size", 120),
|
||
getattr(config, "d_model", 1280),
|
||
quant_config=vllm_config.quant_config,
|
||
prefix=maybe_prefix(prefix, "proj_out"),
|
||
)
|
||
self.proj_out = self.proj_out.tie_weights(self.model.decoder.tgt_word_emb)
|
||
|
||
logit_scale = getattr(config, "logit_scale", 1.0)
|
||
self.logits_processor = LogitsProcessor(
|
||
getattr(config, "vocab_size", 120),
|
||
scale=logit_scale,
|
||
)
|
||
|
||
def forward(
|
||
self,
|
||
input_ids: torch.Tensor,
|
||
positions: torch.Tensor,
|
||
encoder_outputs: list[torch.Tensor] | None = None,
|
||
**kwargs,
|
||
) -> torch.Tensor:
|
||
if encoder_outputs is None:
|
||
encoder_outputs = []
|
||
decoder_outputs = self.model(
|
||
input_ids=input_ids,
|
||
positions=positions,
|
||
encoder_outputs=encoder_outputs,
|
||
)
|
||
return decoder_outputs
|
||
|
||
def embed_multimodal(self, **kwargs: object) -> MultiModalEmbeddings:
|
||
"""Run encoder on audio features and return per-item embeddings."""
|
||
audio_input = self._parse_and_validate_audio_input(**kwargs)
|
||
|
||
speech = audio_input["input_features"]
|
||
speech_lengths = audio_input["speech_lengths"]
|
||
if speech is None or speech_lengths is None:
|
||
return []
|
||
|
||
# When audio items have different time lengths, vLLM's
|
||
# MultiModalBatchedField._reduce_data returns a plain
|
||
# list[Tensor] instead of a stacked Tensor. The encoder
|
||
# expects a padded [B, Tmax, feat_dim] Tensor, so we
|
||
# normalise both speech and speech_lengths here.
|
||
if isinstance(speech, (list, tuple)):
|
||
# Each element: [Ti, feat_dim] (or [1, Ti, feat_dim])
|
||
tensors = [
|
||
s.squeeze(0) if s.dim() == 3 and s.size(0) == 1 else s for s in speech
|
||
]
|
||
device = tensors[0].device
|
||
dtype = tensors[0].dtype
|
||
feat_dim = tensors[0].shape[-1]
|
||
lengths = torch.tensor(
|
||
[t.size(0) for t in tensors],
|
||
device=device,
|
||
dtype=torch.int32,
|
||
)
|
||
t_max = int(lengths.max().item())
|
||
# Pre-allocate zero-padded batch tensor
|
||
speech = torch.zeros(
|
||
(len(tensors), t_max, feat_dim),
|
||
device=device,
|
||
dtype=dtype,
|
||
)
|
||
for i, t in enumerate(tensors):
|
||
speech[i, : t.size(0)] = t
|
||
speech_lengths = lengths
|
||
else:
|
||
# Already a batched Tensor [B, T, feat_dim]
|
||
if speech.dim() == 2:
|
||
speech = speech.unsqueeze(0)
|
||
|
||
speech_lengths = torch.as_tensor(
|
||
speech_lengths, dtype=torch.int32, device=speech.device
|
||
)
|
||
|
||
enc_output, enc_lengths = self.model.get_encoder_outputs(
|
||
speech=speech,
|
||
speech_lengths=speech_lengths,
|
||
)
|
||
|
||
# vLLM expects one 2D tensor per multimodal item. Slice each batch entry
|
||
# by the true encoder length so cross-attention never sees padded frames.
|
||
return tuple(
|
||
enc_output[i, : max(0, int(enc_lengths[i].item()))]
|
||
for i in range(enc_output.size(0))
|
||
)
|
||
|
||
def embed_input_ids(
|
||
self,
|
||
input_ids: torch.Tensor,
|
||
multimodal_embeddings: MultiModalEmbeddings | None = None,
|
||
*,
|
||
is_multimodal: torch.Tensor | None = None,
|
||
) -> torch.Tensor:
|
||
return self.model.decoder.embed_input_ids(input_ids)
|
||
|
||
def _parse_and_validate_audio_input(
|
||
self, **kwargs: object
|
||
) -> FireRedLIDAudioInputs:
|
||
input_features = kwargs.pop("input_features", None)
|
||
speech_lengths = kwargs.pop("speech_lengths", None)
|
||
fake_token_lengths = kwargs.pop("fake_token_lengths", None)
|
||
return FireRedLIDAudioInputs(
|
||
input_features=input_features,
|
||
speech_lengths=speech_lengths,
|
||
fake_token_lengths=fake_token_lengths,
|
||
)
|
||
|
||
def compute_logits(self, hidden_states: torch.Tensor) -> torch.Tensor:
|
||
logits = self.logits_processor(self.proj_out, hidden_states)
|
||
return logits
|
||
|
||
@classmethod
|
||
def validate_language(cls, language: str | None) -> str | None:
|
||
# FireRedLID is a language *identification* model – the caller does
|
||
# not need to specify a language up-front. Accept None silently.
|
||
if language is None:
|
||
return None
|
||
return super().validate_language(language)
|
||
|
||
@classmethod
|
||
def get_generation_prompt(
|
||
cls,
|
||
audio: np.ndarray,
|
||
stt_config: SpeechToTextConfig,
|
||
model_config: ModelConfig,
|
||
language: str | None,
|
||
task_type: Literal["transcribe", "translate"],
|
||
request_prompt: str,
|
||
to_language: str | None,
|
||
) -> PromptType:
|
||
"""Build the prompt for the FireRedLID encoder-decoder model.
|
||
|
||
The decoder receives a single <sos> token; the encoder processes
|
||
the raw audio waveform via the multimodal pipeline.
|
||
"""
|
||
prompt: PromptType = {
|
||
"encoder_prompt": {
|
||
"prompt": "",
|
||
"multi_modal_data": {
|
||
"audio": (audio, int(stt_config.sample_rate)),
|
||
},
|
||
},
|
||
"decoder_prompt": {
|
||
"prompt": "<sos>",
|
||
},
|
||
}
|
||
return prompt
|
||
|
||
@classmethod
|
||
def get_speech_to_text_config(
|
||
cls,
|
||
model_config: ModelConfig,
|
||
task_type: Literal["transcribe", "translate"],
|
||
) -> SpeechToTextConfig:
|
||
processor = cached_processor_from_config(model_config)
|
||
return SpeechToTextConfig(
|
||
max_audio_clip_s=processor.feature_extractor.chunk_length,
|
||
sample_rate=processor.feature_extractor.sampling_rate,
|
||
# LID output is at most 2 tokens – no chunking needed.
|
||
min_energy_split_window_size=None,
|
||
)
|
||
|
||
@classmethod
|
||
def post_process_output(cls, text: str) -> str:
|
||
# Strip any leading/trailing whitespace from the raw LID output.
|
||
return text.strip()
|
||
|
||
def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
|
||
loader = AutoWeightsLoader(
|
||
self,
|
||
skip_prefixes=[
|
||
# Position encoding buffers are rebuilt at init
|
||
"model.encoder.positional_encoding.pe",
|
||
"model.decoder.positional_encoding.pe",
|
||
# Tied output projection (shared with embedding)
|
||
"model.decoder.tgt_word_prj.weight",
|
||
"proj_out.",
|
||
],
|
||
)
|
||
return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)
|