diff --git a/test/asset/t5.base.encoder.output.pt b/test/asset/t5.base.encoder.output.pt new file mode 100644 index 0000000000..9d56557c42 Binary files /dev/null and b/test/asset/t5.base.encoder.output.pt differ diff --git a/test/asset/t5.base.output.pt b/test/asset/t5.base.output.pt new file mode 100644 index 0000000000..384b59074d Binary files /dev/null and b/test/asset/t5.base.output.pt differ diff --git a/test/prototype/integration_tests/__init__.py b/test/prototype/integration_tests/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/test/prototype/integration_tests/test_models.py b/test/prototype/integration_tests/test_models.py new file mode 100644 index 0000000000..01d728b179 --- /dev/null +++ b/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) diff --git a/test/prototype/models/__init__.py b/test/prototype/models/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/test/prototype/models/test_models.py b/test/prototype/models/test_models.py new file mode 100644 index 0000000000..ccd165230c --- /dev/null +++ b/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))