Adapted to new fairseq releases

Author: jarfo

models/joint.py

@@ -15,13 +15,13 @@
 from fairseq.modules import PositionalEmbedding
 
 from fairseq.models import (
-    FairseqIncrementalDecoder, FairseqEncoder, FairseqModel, register_model, register_model_architecture
+    FairseqIncrementalDecoder, FairseqEncoder, FairseqEncoderDecoderModel, register_model, register_model_architecture
 )
 
 from .protected_multihead_attention import ProtectedMultiheadAttention
 
 @register_model('joint_attention')
-class JointAttentionModel(FairseqModel):
+class JointAttentionModel(FairseqEncoderDecoderModel):
     """
     Local Joint Source-Target model from
     `"Joint Source-Target Self Attention with Locality Constraints" (Fonollosa, et al, 2019)
@@ -225,7 +225,7 @@ def max_positions(self):
         """Maximum input length supported by the encoder."""
         if self.embed_positions is None:
             return self.max_source_positions
-        return min(self.max_source_positions, self.embed_positions.max_positions())
+        return min(self.max_source_positions, self.embed_positions.max_positions)
 
 
 class JointAttentionDecoder(FairseqIncrementalDecoder):
@@ -413,7 +413,7 @@ def max_positions(self):
         """Maximum output length supported by the decoder."""
         if self.embed_positions is None:
             return self.max_target_positions
-        return min(self.max_target_positions, self.embed_positions.max_positions())
+        return min(self.max_target_positions, self.embed_positions.max_positions)
 
     def buffered_future_mask(self, tensor):
         """Cached future mask."""

models/protected_multihead_attention.py

@@ -6,15 +6,17 @@
 # can be found in the PATENTS file in the same directory.
 
 import torch
+import torch.nn.functional as F
 from torch import nn
 from torch.nn import Parameter
-import torch.nn.functional as F
+from fairseq.incremental_decoding_utils import with_incremental_state
 
 from fairseq import utils
 
 # Adapted from faiserq/modules/multihead_attention to deal with local attention
 # Local attetion masking in combination with padding masking can lead to 
 # all -Inf attention rows. This version detects and corrects this situation
+@with_incremental_state
 class ProtectedMultiheadAttention(nn.Module):
     """Multi-headed attention.
 
@@ -247,15 +249,13 @@ def reorder_incremental_state(self, incremental_state, new_order):
             self._set_input_buffer(incremental_state, input_buffer)
 
     def _get_input_buffer(self, incremental_state):
-        return utils.get_incremental_state(
-            self,
+        return self.get_incremental_state(
             incremental_state,
             'attn_state',
         ) or {}
 
     def _set_input_buffer(self, incremental_state, buffer):
-        utils.set_incremental_state(
-            self,
+        return self.set_incremental_state(
             incremental_state,
             'attn_state',
             buffer,