Testing T5Model

test/prototype/integration_tests/test_models.py

@@ -0,0 +1,35 @@
+import torch
+from test.common.assets import get_asset_path
+from test.common.torchtext_test_case import TorchtextTestCase
+from torchtext.prototype.models import (
+    T5_BASE_ENCODER,
+    T5_BASE,
+)
+
+
+class TestT5(TorchtextTestCase):
+    def _t5_model(self, t5_model, expected_asset_name, model_input):
+        """Verify that pre-trained T5 models in torchtext produce
+        the same output as the HuggingFace reference implementation.
+        """
+        expected_asset_path = get_asset_path(expected_asset_name)
+        model = t5_model.get_model()
+        model = model.eval()
+
+        if model.encoder_only:
+            actual = model(model_input)["encoder_output"]
+        else:
+            actual = model(model_input)["decoder_output"]
+
+        expected = torch.load(expected_asset_path)
+        torch.testing.assert_close(actual, expected)
+
+    def test_t5_base_encoder_model(self):
+        expected_asset_name = "t5.base.encoder.output.pt"
+        model_input = torch.tensor([[1, 2, 3, 4, 5, 6], [7, 8, 9, 0, 0, 0]])
+        self._t5_model(t5_model=T5_BASE_ENCODER, expected_asset_name=expected_asset_name, model_input=model_input)
+
+    def test_t5_base_model(self):
+        expected_asset_name = "t5.base.output.pt"
+        model_input = torch.tensor([[1, 2, 3, 4, 5, 6], [7, 8, 9, 0, 0, 0]])
+        self._t5_model(t5_model=T5_BASE, expected_asset_name=expected_asset_name, model_input=model_input)

test/prototype/models/test_models.py

@@ -0,0 +1,117 @@
+from unittest.mock import patch
+
+from test.common.torchtext_test_case import TorchtextTestCase
+
+
+class TestModels(TorchtextTestCase):
+    def test_t5_bundler_build_model(self):
+        from torchtext.prototype.models import T5Conf, T5Model, T5Bundle
+
+        # case: user provide encoder checkpoint state dict
+        dummy_encoder_conf = T5Conf(
+            encoder_only=True,
+            vocab_size=10,
+            embedding_dim=16,
+            ffn_dimension=64,
+            num_attention_heads=2,
+            num_encoder_layers=2,
+        )
+        dummy_t5_encoder = T5Model(dummy_encoder_conf)
+        t5_encoder_model = T5Bundle.build_model(config=dummy_encoder_conf, checkpoint=dummy_t5_encoder.state_dict())
+        self.assertEqual(t5_encoder_model.state_dict(), dummy_t5_encoder.state_dict())
+
+        # case: user provide encoder-decoder checkpoint state dict
+        dummy_t5_conf = T5Conf(
+            encoder_only=False,
+            vocab_size=10,
+            embedding_dim=16,
+            ffn_dimension=64,
+            num_attention_heads=2,
+            num_encoder_layers=2,
+        )
+        dummy_t5 = T5Model(dummy_t5_conf)
+        t5_model = T5Bundle.build_model(config=dummy_t5_conf, checkpoint=dummy_t5.state_dict())
+        self.assertEqual(t5_model.state_dict(), dummy_t5.state_dict())
+
+    @patch("logging.Logger.warning")
+    def test_t5_bundler_get_model(self, mock):
+        from torchtext.prototype.models import T5Conf, T5Bundle
+
+        # encoder-only
+        dummy_encoder_conf = T5Conf(
+            encoder_only=True,
+            vocab_size=10,
+            embedding_dim=16,
+            ffn_dimension=64,
+            num_attention_heads=2,
+            num_encoder_layers=2,
+        )
+        encoder_bundle = T5Bundle(dummy_encoder_conf)
+        encoder_bundle.get_model(load_weights=False, freeze_model=True)
+        mock.assert_called_with(
+            "The model is not loaded with pre-trained weights. Setting freeze_model to True will hinder model from learning appropriate weights."
+        )
+
+        # encoder-decoder
+        dummy_t5_conf = T5Conf(
+            encoder_only=False,
+            vocab_size=10,
+            embedding_dim=16,
+            ffn_dimension=64,
+            num_attention_heads=2,
+            num_encoder_layers=2,
+        )
+        t5_bundle = T5Bundle(dummy_t5_conf)
+        t5_bundle.get_model(load_weights=False, freeze_model=True)
+        mock.assert_called_with(
+            "The model is not loaded with pre-trained weights. Setting freeze_model to True will hinder model from learning appropriate weights."
+        )
+
+    def test_t5_bundler_raise_checkpoint(self):
+        from torchtext.prototype.models import T5Conf, T5Bundle
+
+        # encoder-only
+        with self.assertRaises(TypeError):
+            dummy_encoder_conf = T5Conf(
+                encoder_only=True,
+                vocab_size=10,
+                embedding_dim=16,
+                ffn_dimension=64,
+                num_attention_heads=2,
+                num_encoder_layers=2,
+            )
+            T5Bundle.build_model(
+                config=dummy_encoder_conf,
+                freeze_model=True,
+                checkpoint=1,
+            )
+
+        # encoder-decoder
+        with self.assertRaises(TypeError):
+            dummy_t5_conf = T5Conf(
+                encoder_only=False,
+                vocab_size=10,
+                embedding_dim=16,
+                ffn_dimension=64,
+                num_attention_heads=2,
+                num_encoder_layers=2,
+            )
+            T5Bundle.build_model(
+                config=dummy_t5_conf,
+                freeze_model=True,
+                checkpoint=1,
+            )
+
+    def test_t5_bundler_conf_property(self):
+        from torchtext.prototype.models import T5Conf, T5Bundle
+
+        dummy_t5_conf = T5Conf(
+            encoder_only=False,
+            vocab_size=10,
+            embedding_dim=16,
+            ffn_dimension=64,
+            num_attention_heads=2,
+            num_encoder_layers=2,
+        )
+        t5_bundle = T5Bundle(dummy_t5_conf)
+        self.assertTrue(isinstance(t5_bundle.config, T5Conf))