@MULTIMODAL_REGISTRY.register_processor(
CanaryMultiModalProcessor,
info=CanaryProcessingInfo,
dummy_inputs=CanaryDummyInputsBuilder,
)
class CanaryForConditionalGeneration(
nn.Module, SupportsTranscription, SupportsMultiModal
):
packed_modules_mapping = {
"qkv_proj": ["q_proj", "k_proj", "v_proj"],
"kv_proj": ["k_proj", "v_proj"],
}
# Only the decoder is built from vLLM layers; encoder weights keep their
# Transformers names.
hf_to_vllm_mapper = WeightsMapper(
orig_to_new_regex={
re.compile(r"(decoder\.layers\.\d+\.\w+_attn)\.o_proj\."): (
r"\1.out_projection."
),
re.compile(r"(decoder\.layers\.\d+\.mlp)\.fc1\."): r"\1.dense_in.",
re.compile(r"(decoder\.layers\.\d+\.mlp)\.fc2\."): r"\1.dense_out.",
re.compile(r"\.num_batches_tracked$"): None,
},
)
supported_languages = CANARY_SUPPORTED_LANGS
supports_transcription_only = True
skip_warmup_audio_preprocessing = True
# TODO: canary-1b-v2 handles long-form audio with dynamic chunking in NeMo;
# confirm the right per-chunk limit for vLLM's chunker.
max_audio_clip_s: ClassVar[int] = 40
@classmethod
def validate_language(cls, language: str | None) -> str | None:
if language is None:
logger.warning(
"Defaulting to language='en'. If you wish to transcribe audio in "
"a different language, pass the `language` field in the "
"TranscriptionRequest."
)
language = "en"
# The decoder prompt needs one of the trained language tokens, so
# unlike the base class, do not fall back for other languages.
if language not in cls.supported_languages:
raise ValueError(
f"Unsupported language: {language!r}. Canary supports "
f"{list(cls.supported_languages)}."
)
return language
@classmethod
def get_decoder_prompt_tokens(
cls, source_language: str, target_language: str, punctuation: bool = True
) -> list[str]:
"""Same token sequence as the checkpoint's ``chat_template.jinja``."""
return [
"<|startofcontext|>",
"<|startoftranscript|>",
"<|emo:undefined|>",
f"<|{source_language}|>",
f"<|{target_language}|>",
"<|pnc|>" if punctuation else "<|nopnc|>",
"<|noitn|>",
"<|notimestamp|>",
"<|nodiarize|>",
]
@classmethod
def get_generation_prompt(cls, stt_params: SpeechToTextParams) -> PromptType:
source_language = stt_params.language
if source_language is None:
raise ValueError("Language must be specified for Canary")
target_language = source_language
if stt_params.task_type == "translate":
target_language = stt_params.to_language or "en"
if target_language not in cls.supported_languages:
raise ValueError(
f"Unsupported target language {target_language!r}; expected "
f"one of {list(cls.supported_languages)}."
)
# Canary cannot detect the audio's language, and `validate_language`
# defaults a missing one to "en", so a translation without
# `language` would silently become English transcription.
if source_language == target_language:
raise ValueError(
"Canary translation needs `language` (the audio's language) "
f"to differ from `to_language` ({target_language!r}); it "
"cannot detect the source language."
)
# canary-1b-v2 is trained to translate to or from English only.
if "en" not in (source_language, target_language):
raise ValueError(
"Canary translates to or from English only; got "
f"{source_language!r} -> {target_language!r}."
)
tokenizer = cached_tokenizer_from_config(stt_params.model_config)
tokens = cls.get_decoder_prompt_tokens(source_language, target_language)
prompt_token_ids = tokenizer.convert_tokens_to_ids(tokens)
unk_token_id = getattr(tokenizer, "unk_token_id", None)
if unk_token_id is not None and unk_token_id in prompt_token_ids:
raise ValueError("Failed to resolve the Canary decoder control tokens.")
audio = stt_params.audio
return TokensPrompt(
prompt_token_ids=prompt_token_ids,
multi_modal_data={"audio": (audio, stt_params.stt_config.sample_rate)},
)
@classmethod
def get_placeholder_str(cls, modality: str, i: int) -> str | None:
if modality.startswith("audio"):
return None
raise ValueError("Only audio modality is supported")
@classmethod
def get_speech_to_text_config(
cls, model_config: ModelConfig, task_type: str
) -> SpeechToTextConfig:
return SpeechToTextConfig(
max_audio_clip_s=cls.max_audio_clip_s,
sample_rate=16_000,
)
@classmethod
def get_num_audio_tokens(
cls,
audio_duration_s: float,
stt_config: SpeechToTextConfig,
model_config: ModelConfig,
) -> int | None:
# Matches ParakeetFeatureExtractor defaults (n_fft=512, hop_length=160).
num_frames = int(audio_duration_s * stt_config.sample_rate) // 160
return _subsampled_length(num_frames, model_config.hf_config.encoder_config)
def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
super().__init__()
config: CanaryConfig = vllm_config.model_config.hf_config
quant_config = vllm_config.quant_config
self.config = config
self.dtype = vllm_config.model_config.dtype
with self._mark_composite_model(
vllm_config,
language_targets=CanaryDecoder,
tower_targets={"audio": HFParakeetEncoder},
):
self.model = CanaryModel(
vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model")
)
decoder_config = config.decoder_config
self.proj_out = ParallelLMHead(
decoder_config.vocab_size,
decoder_config.hidden_size,
quant_config=quant_config,
bias=True,
prefix=maybe_prefix(prefix, "proj_out"),
)
self.logits_processor = LogitsProcessor(decoder_config.vocab_size)
def forward(
self,
input_ids: torch.Tensor,
positions: torch.Tensor,
encoder_outputs: list[torch.Tensor] | None = None,
**kwargs,
) -> torch.Tensor:
return self.model(
input_ids=input_ids,
positions=positions,
encoder_outputs=encoder_outputs or [],
)
def get_language_model(self) -> torch.nn.Module:
return self.model.decoder
def embed_multimodal(self, **kwargs: object) -> MultiModalEmbeddings:
input_features = kwargs.pop("input_features", None)
feature_lengths = kwargs.pop("feature_lengths", None)
if input_features is None:
return []
if not isinstance(feature_lengths, torch.Tensor):
raise ValueError("Audio feature lengths must be a tensor.")
if isinstance(input_features, torch.Tensor):
features = list(input_features)
elif is_list_of(input_features, torch.Tensor):
features = input_features
else:
raise ValueError(
f"Incorrect type of audio features: {type(input_features)}"
)
padded = torch.nn.utils.rnn.pad_sequence(
[f.to(self.dtype) for f in features], batch_first=True, padding_value=0.0
)
return self.model.get_encoder_outputs(padded, feature_lengths.reshape(-1))
def compute_logits(self, hidden_states: torch.Tensor) -> torch.Tensor:
return self.logits_processor(self.proj_out, hidden_states, self.proj_out.bias)
def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
stacked_params_mapping = [
# (param_name, shard_name, shard_id)
(".self_attn.qkv_proj", ".self_attn.q_proj", "q"),
(".self_attn.qkv_proj", ".self_attn.k_proj", "k"),
(".self_attn.qkv_proj", ".self_attn.v_proj", "v"),
(".encoder_attn.kv_proj", ".encoder_attn.k_proj", "k"),
(".encoder_attn.kv_proj", ".encoder_attn.v_proj", "v"),
]
params_dict = dict(self.named_parameters())
params_dict.update(self.named_buffers())
loaded_params: set[str] = set()
for name, loaded_weight in self.hf_to_vllm_mapper.apply(weights):
# The encoder is the Transformers module and keeps unfused q/k/v.
is_decoder = name.startswith("model.decoder.")
for param_name, shard_name, shard_id in stacked_params_mapping:
if not is_decoder or shard_name not in name:
continue
name = name.replace(shard_name, param_name)
param = params_dict[name]
param.weight_loader(param, loaded_weight, shard_id)
break
else:
param = params_dict[name]
weight_loader = getattr(param, "weight_loader", default_weight_loader)
weight_loader(param, loaded_weight)
loaded_params.add(name)
return loaded_params