Optimize T5 for sequence generation #2054
Conversation
joecummings
force-pushed
the
optimize-t5
branch
3 times, most recently
from
daa19d7
to
73465d8
Compare
joecummings
force-pushed
the
optimize-t5
branch
from
e807b7e
to
d0e866e
Compare
joecummings
requested review from
Nayef211,
abhinavarora,
rshraga and
mthrok
| @@ -79,7 +76,7 @@ def _t5_get_encoder(self, model, model_input, encoder_output): | |||
| encoder = model.get_encoder() | |||
| # Need to set the tgt_key_padding_mask to ensure the same results | |||
| encoder_padding_mask = model_input.eq(model.padding_idx) | |||
| output_from_get_encoder = encoder(tgt=model_input, tgt_key_padding_mask=encoder_padding_mask)["encoder_output"] | |||
| output_from_get_encoder = encoder(model_input, src_key_padding_mask=encoder_padding_mask)["encoder_output"] | |||
There was a problem hiding this comment.
This change, is it using different set of existing arguments or changing the name of the arguments?
If changing the name of the arguments, that's BC-breaking unless it's prototype.
There was a problem hiding this comment.
changing the name of arguments, but yes this is prototype until tomorrow :)
| @@ -56,13 +55,13 @@ def __post_init__(self): | |||
| self.activation = "gelu_new" | |||
|
|
|||
|
|
|||
| # NOTE: Comparable HuggingFace implentation can be found at https://github.com/huggingface/transformers/blob/8581a798c0a48fca07b29ce2ca2ef55adcae8c7e/src/transformers/models/t5/modeling_t5.py#L1269 | |||
There was a problem hiding this comment.
Is this context no-longer applicable?
There was a problem hiding this comment.
We include in the header that several functions are based on HF and I call it out in the docstring of those functions, as well. No need to say that there is a comparable HF implementation for ones that are just the normal Enc/Dec forward functions.
| @torch.jit.export | ||
| def _reorder_cache( | ||
| self, past: List[Tuple[Tensor, Tensor, Tensor, Tensor]], beam_idx: Tensor | ||
| ) -> List[Tuple[Tensor, Tensor, Tensor, Tensor]]: |
There was a problem hiding this comment.
would be nice if there is a comment/docstring of why and what, for the future developer.
| for layer_past_states in past: | ||
| # get the correct batch idx from layer past batch dim | ||
| # batch dim of `past` is at 2nd position | ||
| reordered_layer_past_states = () |
There was a problem hiding this comment.
List would be semantically better, but is it for TorchScript compaibility?
| ) -> Dict[ | ||
| str, | ||
| Union[ | ||
| Tensor, | ||
| Dict[str, Union[Optional[Tensor], List[Tensor], List[Optional[Tensor]]]], | ||
| Optional[List[Tuple[Tensor, Tensor, Tensor, Tensor]]], | ||
| bool, | ||
| ], | ||
| ]: |
There was a problem hiding this comment.
This annotation is complex and it seems to be repeated. Can we define variable to store the annotation?
This PR makes the following changes to T5 to improve generation capabilities.
prepare_inputs_for_generationfunction to be compliant w/GenerationWrapperAPIget_encoderandget_decoderhelper functions.past_key_valuesto implement incremental decoding. This involves also a custom reorder cache function that can be used for beam search.T5WrapperTesting:
Todo: