transformers icon indicating copy to clipboard operation
transformers copied to clipboard

Add Jukebox model (replaces #16875)

Open ArthurZucker opened this issue 2 years ago • 4 comments

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 function class JukeboxConverter(Converter):

ArthurZucker avatar Jun 22 '22 17:06 ArthurZucker

Replaces (#16875)

ArthurZucker avatar Jun 22 '22 17:06 ArthurZucker

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 avatar Jul 08 '22 07:07 ArthurZucker

@ArthurZucker let me know when this is ready for the final review

patrickvonplaten avatar Aug 31 '22 12:08 patrickvonplaten

@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

patrickvonplaten avatar Oct 20 '22 16:10 patrickvonplaten

Apart from the kwargs I think it is done! @patrickvonplaten feel free to review

ArthurZucker avatar Oct 21 '22 17:10 ArthurZucker

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.

ArthurZucker avatar Nov 03 '22 10:11 ArthurZucker

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.

ArthurZucker avatar Nov 09 '22 16:11 ArthurZucker

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.