Torchscriptable T5 generation #2146
Conversation
joecummings
force-pushed
the
torchscripting-gen-utils
branch
from
682b21a
to
2f1d38b
Compare
|
@Nayef211 Thoughts on the |
torchtext/prototype/generate.py
Outdated
| model_inputs = ( | ||
| self.model.prepare_inputs_for_generation(input_ids, model_kwargs=model_kwargs) | ||
| if torch.jit.is_scripting() | ||
| else self._call_to_prepare_inputs_for_generation_with_kwargs(input_ids, model_kwargs) | ||
| ) |
There was a problem hiding this comment.
Can we leave a comment here explaining what's going on? These helper method names are quite long and it's not immediately clear what the differences are in the model inputs based on if we're in script mode or eager mode.
| @@ -78,10 +80,12 @@ class T5Bundle: | |||
| def get_model( | |||
| self, | |||
| *, | |||
| with_generation_utils: bool = False, | |||
There was a problem hiding this comment.
Please add an explanation of this param to the docstring.
I think the |
There was a problem hiding this comment.
Minor comments.
+1 on @Nayef211 comment on naming non-script and script versions of preparation-for-decoding methods so that one can see that they do the same (up to scripting) at a glance.
| model_kwargs = { | ||
| "encoder_outputs": encoder_outputs, | ||
| "encoder_padding_mask": encoder_padding_mask, | ||
| "past_key_values": past, |
There was a problem hiding this comment.
Very minor: rename past into past_key_values to follow other key-value naming ?
joecummings
force-pushed
the
torchscripting-gen-utils
branch
from
2f1d38b
to
09d56fd
Compare
|
@atalman @osalpekar Why is |
joecummings
force-pushed
the
torchscripting-gen-utils
branch
from
09d56fd
to
e020b62
Compare
joecummings
force-pushed
the
torchscripting-gen-utils
branch
from
e020b62
to
5ab9025
Compare
Makes
GenerationUtilsTorchScript-compatible w/ T5 model.This PR makes the following changes:
GenerationUtilsannn.Moduleand TorchScript-compatiblewith_generation_utilstoT5Bundle.get_model()that returns a T5 model wrapped in a TorchScript-compatibleGenerationUtilsclassTesting:
GenerationUtilsget_modelwithwith_generation_utils=Trueand confirming that it is Torchscriptable and provides the correct results.