| |
| |
| |
| |
|
|
| import logging |
| from typing import Dict, Optional |
|
|
| from fairseq.incremental_decoding_utils import with_incremental_state |
| from fairseq.models import FairseqDecoder |
| from torch import Tensor |
|
|
|
|
| logger = logging.getLogger(__name__) |
|
|
|
|
| @with_incremental_state |
| class FairseqIncrementalDecoder(FairseqDecoder): |
| """Base class for incremental decoders. |
| |
| Incremental decoding is a special mode at inference time where the Model |
| only receives a single timestep of input corresponding to the previous |
| output token (for teacher forcing) and must produce the next output |
| *incrementally*. Thus the model must cache any long-term state that is |
| needed about the sequence, e.g., hidden states, convolutional states, etc. |
| |
| Compared to the standard :class:`FairseqDecoder` interface, the incremental |
| decoder interface allows :func:`forward` functions to take an extra keyword |
| argument (*incremental_state*) that can be used to cache state across |
| time-steps. |
| |
| The :class:`FairseqIncrementalDecoder` interface also defines the |
| :func:`reorder_incremental_state` method, which is used during beam search |
| to select and reorder the incremental state based on the selection of beams. |
| |
| To learn more about how incremental decoding works, refer to `this blog |
| <http://www.telesens.co/2019/04/21/understanding-incremental-decoding-in-fairseq/>`_. |
| """ |
|
|
| def __init__(self, dictionary): |
| super().__init__(dictionary) |
|
|
| def forward( |
| self, prev_output_tokens, encoder_out=None, incremental_state=None, **kwargs |
| ): |
| """ |
| Args: |
| prev_output_tokens (LongTensor): shifted output tokens of shape |
| `(batch, tgt_len)`, for teacher forcing |
| encoder_out (dict, optional): output from the encoder, used for |
| encoder-side attention |
| incremental_state (dict, optional): dictionary used for storing |
| state during :ref:`Incremental decoding` |
| |
| Returns: |
| tuple: |
| - the decoder's output of shape `(batch, tgt_len, vocab)` |
| - a dictionary with any model-specific outputs |
| """ |
| raise NotImplementedError |
|
|
| def extract_features( |
| self, prev_output_tokens, encoder_out=None, incremental_state=None, **kwargs |
| ): |
| """ |
| Returns: |
| tuple: |
| - the decoder's features of shape `(batch, tgt_len, embed_dim)` |
| - a dictionary with any model-specific outputs |
| """ |
| raise NotImplementedError |
|
|
| def reorder_incremental_state( |
| self, |
| incremental_state: Dict[str, Dict[str, Optional[Tensor]]], |
| new_order: Tensor, |
| ): |
| """Reorder incremental state. |
| |
| This will be called when the order of the input has changed from the |
| previous time step. A typical use case is beam search, where the input |
| order changes between time steps based on the selection of beams. |
| """ |
| pass |
|
|
| def reorder_incremental_state_scripting( |
| self, |
| incremental_state: Dict[str, Dict[str, Optional[Tensor]]], |
| new_order: Tensor, |
| ): |
| """Main entry point for reordering the incremental state. |
| |
| Due to limitations in TorchScript, we call this function in |
| :class:`fairseq.sequence_generator.SequenceGenerator` instead of |
| calling :func:`reorder_incremental_state` directly. |
| """ |
| for module in self.modules(): |
| if hasattr(module, "reorder_incremental_state"): |
| result = module.reorder_incremental_state(incremental_state, new_order) |
| if result is not None: |
| incremental_state = result |
|
|
| def set_beam_size(self, beam_size): |
| """Sets the beam size in the decoder and all children.""" |
| if getattr(self, "_beam_size", -1) != beam_size: |
| seen = set() |
|
|
| def apply_set_beam_size(module): |
| if ( |
| module != self |
| and hasattr(module, "set_beam_size") |
| and module not in seen |
| ): |
| seen.add(module) |
| module.set_beam_size(beam_size) |
|
|
| self.apply(apply_set_beam_size) |
| self._beam_size = beam_size |
|
|
