transformers
transformers copied to clipboard
Add Jukebox model (replaces #16875)
This is a draft pull request.
What does this PR do?
This PR will progressively add the Jukebox model to the hub. It is linked to #16870.
Currently planned steps (WIP)
- [x] Create template files with
transformeres-cli add-new-model-like
- [x]
src/transformers/tokenization_jukebox.py
- [x]
src/transformers/test_tokenization_jukebox.py
- [x]
src/transformers/configuration_jukebox.py
- [x]
src/transformers/modeling_jukebox.py
- [x]
src/transformers/configuration_jukebox.py
- [x]
docs/source/model_doc/jukebox.rst
- [ ]
src/transformers/tokenization_jukebox_fast.py
(will most probably use WordLevel tokenizer). Also requires to implement a converter functionclass JukeboxConverter(Converter):
Replaces (#16875)
The docs for this PR live here. All of your documentation changes will be reflected on that endpoint.
Okay, the 1b lyrics
and 5b lyrics
match the original code. Just need refactoring to have better variable names and wrap the sampling kwargs for easy use.
@ArthurZucker let me know when this is ready for the final review
@ArthurZucker do you want to let me know once you want to have a final review? Let's try to not let it hang around for too long
Apart from the kwargs
I think it is done! @patrickvonplaten feel free to review
Slow tests are now passing, the only issue left to attend is the memory. The slow tests need a lot of RAM, and running inference with the model should also automatically send the unused Priors
and VQVAE
to the cpu
.
The documentation and models will be ready soon.
As it was previously requested, you can now instantiate JukeboxVQVAE
and JukeboxPrior
individually. This is convenient if people only want to use the VQVAE or generate form juste the top level prior.
The docs for this PR live here. All of your documentation changes will be reflected on that endpoint.
The docs for this PR live here. All of your documentation changes will be reflected on that endpoint.
The docs for this PR live here. All of your documentation changes will be reflected on that endpoint.