    def test_mmt_txt(self):
        batch_size = 8

        class MockBert(nn.Module):
            """
            Mock BERT model in transformers
            """

            def __init__(self, hidden_size):
                super().__init__()
                self.hidden_size = hidden_size

                class Struct:
                    def __init__(self, **entries):
                        self.__dict__.update(entries)

                config = Struct(**{"hidden_size": hidden_size})
                self.config = config

            def forward(self, **kwargs):
                x = kwargs["input_ids"]
                return torch.rand(x.shape[0], x.shape[1], self.hidden_size), torch.rand(x.shape[0], self.hidden_size)

        class MockTokenizer:
            """
            Mock tokenizer in transformers
            """

            def __init__(self, vocab_size):
                self.vocab_size = vocab_size

            def __call__(self, text, max_length=30, truncation=True, add_special_tokens=None, padding=True,
                         return_tensors="pt"):
                print("tokenize text:", text)

                return {"input_ids": torch.randint(low=0, high=self.vocab_size, size=(batch_size, max_length,))}

        mmttxt_model = MMTTXT(
            txt_bert=MockBert(hidden_size=768),
            tokenizer=MockTokenizer(vocab_size=100),
            max_length=30,
            modalities=["CLIP", "tf_vggish"],
            add_special_tokens=True,
            add_dot=True,
            same_dim=512,
            dout_prob=0.2,
        )
        output = mmttxt_model([
            "hello world",
            "how are you"
        ])
        self.assertTrue(output.shape == (batch_size, 1024))
Mock BERT model in transformers
tests/unittests/models/mdmmt/test_mmt.py