search

pytorch / text

Models, data loaders and abstractions for language processing, powered by PyTorch
https://pytorch.org/text
BSD 3-Clause "New" or "Revised" License
3.49k stars 815 forks source link

Added notebook torchscriptable_t5_with_torchtext.ipynb #2122

Closed rbahumi closed 8 months ago

rbahumi commented 1 year ago

This notebook is an example of a (working) "Hacky" solution for wrapping the full 'generate' functionality inside a "forward" function. The perpose of this is to start a discussion and be a suggention on how to make the this functionality TorchScriptable.

To do so I:

  1. Inherited from T5Model:
    • extracted the decoding code from t5.forward() function to a standalone 'decode' function that returns a specific type.
    • added the GenerationUtils's 'generate' functionality as a class method (similar to HuggingFace).
  2. Added TorchScriptableT5, a module that implements the full generative logic in the forward method.
  3. Helper classes that build a jit (TorchScript) model from a predefined T5 Bundle