model: Add NVIDIA Canary-1B-v2 to Transformers - #46825
Conversation
|
Good day @ebezzam! Just bumping this up for whenever you find the time; thanks! |
ebezzam
left a comment
ebezzam
left a comment
There was a problem hiding this comment.
hi @harshaljanjani thanks for the contribution! Here are some initial comments to start iterating. I didn't have time for an in-depth review, but hopefully these help to get things in the expected direction!
Main idea is to see how existing models define / implement the various parts, and trying to use modular as much as possible!
|
|
||
| def __init__( | ||
| self, | ||
| embed_dim: int, |
There was a problem hiding this comment.
I wonder if super_kwargs could be used here? I learned about this for here
There was a problem hiding this comment.
Thanks for the precedent link; done!
|
|
||
|
|
||
| @auto_docstring | ||
| class CanaryPreTrainedModel(PreTrainedModel): |
There was a problem hiding this comment.
(nit but good practice) could we do modular from an existing model? to save some lines
There was a problem hiding this comment.
Done as well.
| speech-to-text translation. | ||
| """ | ||
| ) | ||
| class CanaryForConditionalGeneration(CanaryPreTrainedModel, GenerationMixin): |
There was a problem hiding this comment.
can we do modular from WhisperForConditionalGeneration?
There was a problem hiding this comment.
IMO since WhisperForConditionalGeneration inherits from the custom WhisperGenerationMixin with tons of custom methods specific to Whisper (long-form chunking, timestamp extraction for Whisper-specific token IDs, etc., also bias = False), we'd be adding breakage here if we inherited from the class. Moonshine sets precedent as well by being standalone with the same WhisperModel derived shape. Happy to know if I'm missing something here though!
There was a problem hiding this comment.
ah yes I forgot the inheritance from WhisperGenerationMixin! Could we do modular from MoonshineForConditionalGeneration then?
There was a problem hiding this comment.
Done; only the necessary overrides remain now!
| init.copy_(module.positional_embeddings, module._build_table()) | ||
|
|
||
|
|
||
| class CanaryDecoder(CanaryPreTrainedModel): |
There was a problem hiding this comment.
TO CHECK: there isn't a similar module elsewhere in the lib? for example Whisper
There was a problem hiding this comment.
Deduped it with WhisperDecoder, thanks for the nudge!
| class CanaryModel(CanaryPreTrainedModel): | ||
| def __init__(self, config: CanaryConfig): | ||
| super().__init__(config) | ||
| self.encoder = AutoModel.from_config(config.encoder_config) | ||
| self.decoder = CanaryDecoder(config) | ||
| self.post_init() |
There was a problem hiding this comment.
could we do modular from WhisperModel?
| audio = make_list_of_audio(audio) | ||
| inputs = self.feature_extractor(audio, **output_kwargs["audio_kwargs"]) | ||
|
|
||
| prompt_tokens = self._build_prompt_tokens(source_lang, target_lang, pnc, timestamps) |
There was a problem hiding this comment.
can we use a chat template for specifying these prompts? See VibeVoice ASR or Qwen 3 ASR. And so we likely need to define an additional method called apply_transcription_request that will call apply_chat_template which then calls this __call__.
The purpose of this __call__ should be simply:
- apply feature extrator to audio and insert audio tokens
- tokenize text
- (optionally) prepare output labels
you will need to add this chat template to your checkpoint in the conversion script
There was a problem hiding this comment.
Done following the aforementioned lib patterns, happy to adjust if anything comes up!
|
|
||
| Canary reuses the [Fast Conformer](https://huggingface.co/papers/2305.05084) encoder from [Parakeet](./parakeet.md) (loaded through [`ParakeetEncoder`] / [`ParakeetEncoderConfig`]) and pairs it with a Transformer decoder that uses fixed sinusoidal positional embeddings, cross-attention to the encoder outputs and tied input/output embeddings. The task is selected through a decoder prompt prefix built by [`CanaryProcessor`] of the form `<|startofcontext|> <|startoftranscript|> <source_lang> <target_lang> <pnc|nopnc> <timestamp|notimestamp> ...`, where `source_lang == target_lang` selects transcription and otherwise selects translation. | ||
|
|
||
| The original implementation can be found in [NVIDIA NeMo](https://github.com/NVIDIA/NeMo). A Transformers-compatible checkpoint is available at [harshaljanjani/canary-1b-v2-hf](https://huggingface.co/harshaljanjani/canary-1b-v2-hf). |
There was a problem hiding this comment.
NOTE for later: we'll eventually want to transfer a checkpoint to NVIDIA's org
|
|
||
| The original implementation can be found in [NVIDIA NeMo](https://github.com/NVIDIA/NeMo). A Transformers-compatible checkpoint is available at [harshaljanjani/canary-1b-v2-hf](https://huggingface.co/harshaljanjani/canary-1b-v2-hf). | ||
|
|
||
| ## Usage |
There was a problem hiding this comment.
other example usage, we'd like to cover:
- batch
- training (simply showing forward/backward)
- torch compile
- model features such as timestamp, translation, diariarization
Check out existing models like AudioFlamingo3, VIbevoice ASR, Qwen3 ASR
There was a problem hiding this comment.
Done, fleshed out all the tasks with examples :)
|
Good day @ebezzam! Just checking in to see if there've been any updates :) |
ebezzam
left a comment
ebezzam
left a comment
There was a problem hiding this comment.
@harshaljanjani thank you for addressing my previous comments so quickly and thoroughly! Here are some more to iterate on, and some pointer to prepare for a possible merge with the original model card.
|
|
||
| Canary reuses the [Fast Conformer](https://huggingface.co/papers/2305.05084) encoder from [Parakeet](./parakeet.md) (loaded through [`ParakeetEncoder`] / [`ParakeetEncoderConfig`]) and pairs it with a Transformer decoder that uses fixed sinusoidal positional embeddings, cross-attention to the encoder outputs and tied input/output embeddings. The task is selected through a decoder prompt prefix built by [`CanaryProcessor`] of the form `<|startofcontext|> <|startoftranscript|> <|emo:undefined|> <source_lang> <target_lang> <pnc|nopnc> <|noitn|> <|notimestamp|> <|nodiarize|>`, where `source_lang == target_lang` selects transcription and otherwise selects translation. | ||
|
|
||
| The original implementation can be found in [NVIDIA NeMo](https://github.com/NVIDIA/NeMo). A model checkpoint is available at [harshaljanjani/canary-1b-v2-hf](https://huggingface.co/harshaljanjani/canary-1b-v2-hf). |
There was a problem hiding this comment.
TODO: let's check with NVIDIA about adding the Transformers checkpoint to their model card, as NeMo and Transformers files can exist in the same repo!
Can you open a PR on the HF Hub to add the Transformers file and usage? Like this and this as one PR. And if you can mention this PR + indicate that your HF PR is a draft so that they don't merge it just yet. Thanks!
There was a problem hiding this comment.
| from datasets import load_dataset, Audio | ||
| from transformers import AutoProcessor, AutoModelForSpeechSeq2Seq | ||
|
|
||
| processor = AutoProcessor.from_pretrained("harshaljanjani/canary-1b-v2-hf") |
There was a problem hiding this comment.
TODO (when HF model merged): set checkpoint here and elsewhere to nvidia/canary-1b-v2
| def batch_decode(self, *args, **kwargs): | ||
| return self.tokenizer.batch_decode(*args, **kwargs) |
There was a problem hiding this comment.
no need to have this anymore, decode also handles batch inputs!
| def decode(self, *args, **kwargs): | ||
| return self.tokenizer.decode(*args, **kwargs) |
There was a problem hiding this comment.
no need to specify as ProcessorMixin defines it!
There was a problem hiding this comment.
Removed, thanks!
| def __init__(self, config: CanaryConfig): | ||
| super().__init__(config) | ||
| self.max_source_positions = None | ||
| self.embed_positions = CanarySinusoidalPositionalEmbedding(self.max_target_positions, config.d_model) |
There was a problem hiding this comment.
how about calling this module CanaryPositionalEmbedding so that modular can take care of renaming the prefix and we don't need this line?
There was a problem hiding this comment.
Done, thanks!
| pass | ||
|
|
||
|
|
||
| class CanarySinusoidalPositionalEmbedding(nn.Module): |
There was a problem hiding this comment.
Done! Now subclasses with max_timescale = 10000 ** ((channels - 2) / channels) to make the table bit-exact with NeMo's interleaved table by reordering the hidden dim. Followed the glm4v example to update the converter accordingly. One note: I kept the 1/sqrt(d) scaling inside the module because LayerNorm's eps breaks scale invariance, so scale_embedding=True is not bit-identical. Re-exported and re-uploaded the checkpoint; checked that all baseline generations are byte-identical :)
|
Thank you for your contribution 🤗! CI Security Gate — automatic approval blockedThis PR was not automatically approved for CI because the security gate failed. Possible reasons:
See the workflow run for the exact violations. A maintainer can review and manually approve CI if a finding is a false positive. |
|
Good day @ebezzam, just a gentle ping regarding the review :) |
|
@ebezzam Maybe? |
ebezzam
left a comment
ebezzam
left a comment
There was a problem hiding this comment.
@harshaljanjani thanks for your patience! as switching between several things 😄
Hope you don't mind I've pushed a few changes:
- manually write the configuration file since we don't need modular for it and to avoid the duplicate code (thanks for looking into it with Whisper!)
- some processing refactoring to use common utilities across multiple models (we'll very likely get some comments on that from a core maintainer, but I think they'd be interested to see such utility sharing! cc @vasqu )
- test refactoring
Could you try the decoder config refactoring I suggested?
|
|
||
| @auto_docstring(checkpoint="harshaljanjani/canary-1b-v2-hf") | ||
| @strict | ||
| class CanaryConfig(PreTrainedConfig): |
There was a problem hiding this comment.
Alright thanks for trying! So if we don't use modular for config, let's take it out of here and just write the config file manually. I've pushed the change for that!
| d_model: int = 1024 | ||
| decoder_layers: int = 8 | ||
| decoder_attention_heads: int = 8 | ||
| decoder_ffn_dim: int = 4096 | ||
| decoder_layerdrop: float | int = 0.0 | ||
| activation_function: str = "relu" | ||
| max_target_positions: int = 1024 | ||
| dropout: float | int = 0.1 | ||
| attention_dropout: float | int = 0.1 | ||
| activation_dropout: float | int = 0.1 | ||
| scale_embedding: bool = False |
There was a problem hiding this comment.
it's not something you will have seen in other encoder-decoder models like Whisper/Moonshine (but they also didn't have a separate encoder_config), but let's try creating a separate config for the decoder, namely decoder_config and add it to sub_configs. We can call the class CanaryDecoderConfig and no need to register it with auto model so you can have
sub_configs = {
"encoder_config": AutoConfig,
"decoder_config": CanaryDecoderConfig,
}
It will improve readability and maybe future encoder-decoder models could benefit from it!
NOTE: as it's something new, we may get comments from other reviewers on how to further change (or go back on things). But I think there could be interest/support for it, as we've been tending towards splitting things in subconfigs in newer models! thanks for your patience on this and help in prototyping this 🙂
There was a problem hiding this comment.
Done, please do check if it's what you expected, thanks! Followed dia for the shape. It has model_type = "canary_decoder" so make fix-repo adds it to the tables exactly like dia_decoder
| position_ids = position_ids.unsqueeze(0) | ||
|
|
||
| positions = self.embed_positions(position_ids).to(inputs_embeds.dtype) | ||
| # unlike Whisper, NeMo normalizes the summed token and positional embeddings with a LayerNorm |
There was a problem hiding this comment.
so it's clearer
| # unlike Whisper, NeMo normalizes the summed token and positional embeddings with a LayerNorm | |
| # NOTE: unlike Whisper, NeMo normalizes the summed token and positional embeddings with a LayerNorm |
| past_key_values=past_key_values, | ||
| position_ids=position_ids, | ||
| ) | ||
| # unlike Whisper, the encoder outputs have variable length, so padded frames are masked in cross-attention |
There was a problem hiding this comment.
so it's clearer
| # unlike Whisper, the encoder outputs have variable length, so padded frames are masked in cross-attention | |
| # NOTE: unlike Whisper, the encoder outputs have variable length, so padded frames are masked in cross-attention |
# Conflicts: # src/transformers/models/qwen3_asr/processing_qwen3_asr.py
|
@ebezzam Thank you for your time, pushed the changes and made sure there are no regressions! The four |
ebezzam
left a comment
ebezzam
left a comment
There was a problem hiding this comment.
@harshaljanjani thanks for creating the decoder config! Some more small comments on the training example and double-checking that we don't need to shift the labels like in recent fixes.
| for output_flag in ("output_attentions", "output_hidden_states"): | ||
| if kwargs.get(output_flag) is None: | ||
| kwargs[output_flag] = getattr(self.config, output_flag, False) |
There was a problem hiding this comment.
is there a specific reason you needed this? it may already be handled within the encoder/decoder
| for output_flag in ("output_attentions", "output_hidden_states"): | ||
| if kwargs.get(output_flag) is None: | ||
| kwargs[output_flag] = getattr(self.config, output_flag, False) |
There was a problem hiding this comment.
is there a specific reason you needed this? it may already be handled within the encoder
There was a problem hiding this comment.
Only when the user enables it through the config, config.output_hidden_states = True returns None without it since _attn_implementation only is set recursively on the sub-configs, if you'd rather have me propagate them centrally I'm happy to do so
There was a problem hiding this comment.
I don't see Whisper or any other model with such logic, so I think we can altogether remove?
| loss = None | ||
| if labels is not None: | ||
| loss = self.loss_function(logits, labels, self.config.decoder_config.vocab_size) |
There was a problem hiding this comment.
Can you double check if any shifting is needed? There were some changes recently to Moonshine, Cohere, and encoder-decoder models in general
There was a problem hiding this comment.
Thank you for the precedent, fixed the same way as this and test_encoder_decoder_loss_no_double_shift is un-skipped now
# Conflicts: # src/transformers/models/vibevoice_asr/processing_vibevoice_asr.py
|
Thank you for your time @ebezzam, resolved :) |
ebezzam
left a comment
ebezzam
left a comment
There was a problem hiding this comment.
@harshaljanjani thanks for the changes! and actually looking at the loss, made me realize Cohere ASR would be better for the modular, as we can cut down on all of the forward methods I think!
I've left comments on how it could be done. It works for CanaryModel and CanaryForConditionalGeneration (tried locally and double-checked with the integration tests). And I've put some pointers for CanaryDecoder: this one will need updating the state dict and re-converting the checkpoint.
Sorry for bringing this up late in the review cycle, but we've very close to concise implementation 👌 Thanks again!
| specific head on top. | ||
| """ | ||
| ) | ||
| class CanaryModel(WhisperModel): |
There was a problem hiding this comment.
sorry just noticing this now! but let's instead do modular from Cohere, I think it might even be a straight pass through (and we drop the get_audio_features approach) like this
class CanaryModel(CohereAsrModel):
def __init__(self, config: CanaryConfig):
super().__init__(config)
self.decoder = CanaryDecoder(config.decoder_config)| class CanaryForConditionalGeneration(MoonshineForConditionalGeneration): | ||
| def __init__(self, config: CanaryConfig): | ||
| super().__init__(config) | ||
| self.proj_out = nn.Linear(config.decoder_config.d_model, config.decoder_config.vocab_size, bias=True) |
There was a problem hiding this comment.
this could also be a simpler modular from Cohere with only
class CanaryForConditionalGeneration(CohereAsrForConditionalGeneration):
def __init__(self, config: CanaryConfig):
super().__init__(config)
self.proj_out = nn.Linear(config.decoder_config.d_model, config.decoder_config.vocab_size, bias=True)| init.copy_(module.positional_embedding, position_embeddings) | ||
|
|
||
|
|
||
| class CanaryDecoder(WhisperDecoder): |
There was a problem hiding this comment.
we may be able to do a simpler modular from CohereAsrDecoder by:
- in the state dict: renaming
layernorm_embeddingtoembedding_layernorm - in the state dict: renaming
layer_normtonormin the state dict - in the state dict: renaming
embed_scaletopos_emb. Although maybe we still useCanaryPositionalEmbeddingyou've defined here to make it more explicit that it's a sinusoidal positional embedding -> overwritingself.pos_embin init - (not ideal) overwrite init and add self.proj as an identity layer (it would have been better if that layers weren't in Cohere's decoder as it doesn't depend on the decoder config params but oh well 🤷 )
- removing
config.scale_embeddingfrom the config (anyway it's false which makesself.embed_scale=1)
with such changes, we might be able to avoid defining the forward method!
There was a problem hiding this comment.
Thank you for taking the time to detail this out, it helped a ton! All points applied, the checkpoint re-converted and uploaded to the Hub, and the NVIDIA Hub PR has been updated.
There was a problem hiding this comment.
thanks for the changes! the modular is much cleaner and leaner now 👏
ebezzam
left a comment
ebezzam
left a comment
There was a problem hiding this comment.
@harshaljanjani thanks for the changes, the modular is looking much, much nicer now 🙂
@vasqu could we get your thoughts on some of the decisions we made for this modular?
|
|
||
| model_type = "canary" | ||
| keys_to_ignore_at_inference = ["past_key_values"] | ||
| sub_configs = {"encoder_config": AutoConfig, "decoder_config": CanaryDecoderConfig} |
There was a problem hiding this comment.
what do you think about making a separate decoder sub config? instead of having everything in the main config like this
|
|
||
| def __init__(self, config: CanaryDecoderConfig): | ||
| super().__init__(config) | ||
| self.pos_emb = CanaryPositionalEmbedding(config.max_position_embeddings, config.hidden_size) |
There was a problem hiding this comment.
how about explicitly create a sinusoidal embedding rather than a generic embedding layer like in the Cohere here
| def __init__(self, config: CanaryDecoderConfig): | ||
| super().__init__(config) | ||
| self.pos_emb = CanaryPositionalEmbedding(config.max_position_embeddings, config.hidden_size) | ||
| self.proj = nn.Identity() |
There was a problem hiding this comment.
is this ok? need to do this unideal layer for modular, but ideally that layer should have been in the encoder or between the encoder and decoder I feel?
| class CanaryModel(CohereAsrModel): | ||
| def __init__(self, config: CanaryConfig): | ||
| super().__init__(config) | ||
| self.decoder = CanaryDecoder(config.decoder_config) |
There was a problem hiding this comment.
since we make a separate subconfig for the decoder, otherwise would be straight modular
| class CanaryForConditionalGeneration(CohereAsrForConditionalGeneration): | ||
| def __init__(self, config: CanaryConfig): | ||
| super().__init__(config) | ||
| self.proj_out = nn.Linear(config.decoder_config.hidden_size, config.decoder_config.vocab_size, bias=True) |
There was a problem hiding this comment.
since we make a separate subconfig for the decoder, otherwise would be straight modular
|
run-slow: canary |
|
This comment contains models: ["models/canary"] |
|
The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update. |
|
[For maintainers] Suggested jobs to run (before merge) run-slow: audioflamingo3, auto, canary, glmasr, qwen3_asr, vibevoice_asr |
CI recapDashboard: View test results in Grafana |
5 participants
What does this PR do?
→ This PR adds Canary 1B v2 to Transformers!
References:
→ Model Checkpoints
→ GitHub Repository
→ Transformers Converted Checkpoints
→ Research Paper
→ Testing Script
Rebased on top of @eustlb's #40756
cc: @ebezzam @Rocketknight1
Multilingual Transcription and Test Results
transformersversion:5.13.0.dev0Linux-6.8.0-1060-gcp-x86_64-with-glibc2.353.11.15huggingface_hubversion:1.20.1safetensorsversion:0.8.0accelerateversion:1.14.0not installed2.11.0+cu130 (CUDA)NVIDIA L413.0Before submitting