[translation] Use only layer encoder decoder for testing.

PiperOrigin-RevId: 388614364

[translation] Use only layer encoder decoder for testing.
PiperOrigin-RevId: 388614364
638863bd · Frederick Liu · A. Unique TensorFlower · bb832334 · 638863bd · 638863bd
Commit 638863bd authored Aug 03, 2021 by Frederick Liu Committed by A. Unique TensorFlower Aug 03, 2021
Hide whitespace changes
Inline Side-by-side

Showing with 9 additions and 5 deletions

official/nlp/tasks/translation.py official/nlp/tasks/translation.py +1 -1

official/nlp/tasks/translation_test.py official/nlp/tasks/translation_test.py +8 -4

No files found.
--- a/official/nlp/tasks/translation.py
+++ b/official/nlp/tasks/translation.py
@@ -13,11 +13,11 @@
 # limitations under the License.

 """Defines the translation task."""
+import dataclasses
 import os
 from typing import Optional

 from absl import logging
-import dataclasses
 import sacrebleu
 import tensorflow as tf
 import tensorflow_text as tftxt

--- a/official/nlp/tasks/translation_test.py
+++ b/official/nlp/tasks/translation_test.py
@@ -85,7 +85,8 @@ class TranslationTaskTest(tf.test.TestCase):
  def test_task(self):
    config = translation.TranslationConfig(
        model=translation.ModelConfig(
-            encoder=translation.EncDecoder(), decoder=translation.EncDecoder()),
+            encoder=translation.EncDecoder(num_layers=1),
+            decoder=translation.EncDecoder(num_layers=1)),
        train_data=wmt_dataloader.WMTDataConfig(
            input_path=self._record_input_path,
            src_lang="en", tgt_lang="reverse_en",
@@ -102,7 +103,8 @@ class TranslationTaskTest(tf.test.TestCase):
  def test_no_sentencepiece_path(self):
    config = translation.TranslationConfig(
        model=translation.ModelConfig(
-            encoder=translation.EncDecoder(), decoder=translation.EncDecoder()),
+            encoder=translation.EncDecoder(num_layers=1),
+            decoder=translation.EncDecoder(num_layers=1)),
        train_data=wmt_dataloader.WMTDataConfig(
            input_path=self._record_input_path,
            src_lang="en", tgt_lang="reverse_en",
@@ -122,7 +124,8 @@ class TranslationTaskTest(tf.test.TestCase):
        sentencepeice_model_prefix)
    config = translation.TranslationConfig(
        model=translation.ModelConfig(
-            encoder=translation.EncDecoder(), decoder=translation.EncDecoder()),
+            encoder=translation.EncDecoder(num_layers=1),
+            decoder=translation.EncDecoder(num_layers=1)),
        train_data=wmt_dataloader.WMTDataConfig(
            input_path=self._record_input_path,
            src_lang="en", tgt_lang="reverse_en",
@@ -137,7 +140,8 @@ class TranslationTaskTest(tf.test.TestCase):
  def test_evaluation(self):
    config = translation.TranslationConfig(
        model=translation.ModelConfig(
-            encoder=translation.EncDecoder(), decoder=translation.EncDecoder(),
+            encoder=translation.EncDecoder(num_layers=1),
+            decoder=translation.EncDecoder(num_layers=1),
            padded_decode=False,
            decode_max_length=64),
        validation_data=wmt_dataloader.WMTDataConfig(