Update call sites missing opt (#3715)

* update mha init in polyencoder * add polyencoder mha test * hit both multihead init sites in test
facebookresearch · Jun 15, 2021 · fe36d88 · fe36d88
1 parent 02d4ae7
commit fe36d88
Show file tree

Hide file tree

Showing 2 changed files with 24 additions and 2 deletions.
diff --git a/parlai/agents/transformer/polyencoder.py b/parlai/agents/transformer/polyencoder.py
@@ -346,7 +346,7 @@ def __init__(self, opt, dict_, null_idx):
             # The attention for the codes.
             if self.codes_attention_type == 'multihead':
                 self.code_attention = MultiHeadAttention(
-                    self.codes_attention_num_heads, embed_dim, opt['dropout']
+                    opt, self.codes_attention_num_heads, embed_dim, opt['dropout']
                 )
             elif self.codes_attention_type == 'sqrt':
                 self.code_attention = PolyBasicAttention(
@@ -360,7 +360,7 @@ def __init__(self, opt, dict_, null_idx):
         # The final attention (the one that takes the candidate as key)
         if self.attention_type == 'multihead':
             self.attention = MultiHeadAttention(
-                self.attention_num_heads, opt['embedding_size'], opt['dropout']
+                opt, self.attention_num_heads, opt['embedding_size'], opt['dropout']
             )
         else:
             self.attention = PolyBasicAttention(

diff --git a/tests/test_transformers.py b/tests/test_transformers.py
@@ -782,6 +782,28 @@ class TestPolyencoder(TestTransformerBase):
     def test_resize_embeddings(self):
         self._test_resize_embeddings('transformer/polyencoder')
 
+    def test_multi_head_attention(self):
+        with testing_utils.tempdir() as tmpdir:
+            model_file = os.path.join(tmpdir, 'model_file')
+            _, _ = testing_utils.train_model(
+                Opt(
+                    model='transformer/polyencoder',
+                    task='integration_tests:short_fixed',
+                    n_layers=1,
+                    n_encoder_layers=2,
+                    n_decoder_layers=4,
+                    num_epochs=1,
+                    dict_tokenizer='bytelevelbpe',
+                    bpe_vocab=DEFAULT_BYTELEVEL_BPE_VOCAB,
+                    bpe_merge=DEFAULT_BYTELEVEL_BPE_MERGE,
+                    bpe_add_prefix_space=False,
+                    model_file=model_file,
+                    save_after_valid=True,
+                    poly_attention_type='multihead',
+                    codes_attention_type='multihead',
+                )
+            )
+
 
 @testing_utils.skipUnlessVision
 class TestImagePolyencoder(unittest.TestCase):