-
Notifications
You must be signed in to change notification settings - Fork 26.3k
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
[WIP] Moshi integration #33624
base: main
Are you sure you want to change the base?
[WIP] Moshi integration #33624
Conversation
"for speech-to-speech.", | ||
MOSHI_START_DOCSTRING, | ||
) | ||
class MoshiForConditionalGeneration(MoshiPreTrainedModel): |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
cc @gante, this is the model that'll use a weird generation!
This is roughly how I envision generation. It works already, but there'll be some changes that will make the code a bit heavier
).audio_values | ||
|
||
|
||
return output_text_ids, output_values |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
I actually would like to allow dynamic outputs depending on the type of generation (beam, sample) etc., do you think I can do a nested ModelOutput
?
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
My suggestion would be to make it as close as possible to the return structure from the original generate
. Users transitioning from other models to moshi
would then have as little friction as possible 🤗
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
the generation part makes sense to me!
suggestion: because it is quite convoluted with nested generate calls, adding a block diagram explaining the workflow and linking it to the docstring in def generate()
will likely make the life easier for us (long-term maintenance) and our user (can quickly understand what's going on)
).audio_values | ||
|
||
|
||
return output_text_ids, output_values |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
My suggestion would be to make it as close as possible to the return structure from the original generate
. Users transitioning from other models to moshi
would then have as little friction as possible 🤗
What does this PR do?
Moshi is the latest Kyutai model. It is a streaming speech-to-speech model, that can also do an inner dialogue (i.e it outputs text as well).
In particular, it means that Moshi deals with 3 streams of information:
Similarly to
Musicgen
, audio is represented with audio codebooks, which can be interpreted like tokens. The main difference between text tokens and audio codebooks is that audio codebooks introduce an additional dimension of information.Text tokens are typically of dim
(batch_size, sequence_length)
but audio tokens are of dim(batch_size, num_codebooks, sequence_length)
.It's made of 3 components:
1. The main decoder (Helium in the paper)
Here, it corresponds to
MoshiForCausalLM
. It is strictly a classic text LLM, that uses an architecture similar toGemma
. In other words, it takes text tokens, embeds them, pass them through the decoder and a language head, to get text logits.2. The depth decoder
On its own, it's also a classic LLM, but this time, instead of generating over the time dimension, it generates over the codebook dimension.
It also means that its context length is
num_codebooks
-> it can't generate more thannum_codebooks
.Another interesting difference with a classic LLM is that each timestamp (here it correspond to each codebook) got its own set of Linear Layers and Embeddings.
3. Mimi
It's the audio encoder from Kyutai, that has recently been integrated to transformers, which is used to "tokenize" audio. It has the same use that
Encodec
has inMusicgen
.Architecture choice:
MoshiForCausalLM
corresponds to the main decoder, it can be used as a textual LLM.MoshiDepthDecoder
is the depth decoder mentioned aboveMoshiForConditionalGeneration
encapsulates the main decoder, the depth decoder and the audio encoder.Conceptually,
MoshiForConditionalGeneration
takes as input one stream of text and two streams of audio inputs - what the user has said so far, and what the model have generated so far - and generates two streams - a text stream and an audio stream.How does it work:
-> The input streams are embedded and combined into
inputs_embeds
.->
inputs_embeds
is passed through the main decoder. There's nothing special done here, it's the same operation as Gemma or so on.-> The main decoder outputs
text logits
but also itslast hidden state
which is calledtemporal context
in the picture above.-> the depth decoder switches the dimension on which we generate (codebooks instead of time). It uses the token generated from
text logits
and thetemporal context
to auto-regressively generate audio codebooks.