File size: 1,256 Bytes
f233443 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 | # Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.
#
# This source code is licensed under the BSD-style license found in the
# LICENSE file in the root directory of this source tree.
import pytest
import torch
from tests.test_utils import assert_expected, set_rng_seed
from torchmultimodal.modules.layers.text_embedding import BERTTextEmbeddings
@pytest.fixture(autouse=True)
def random():
set_rng_seed(4)
class TestTextEmbeddings:
@pytest.fixture
def text_embedding(self):
return BERTTextEmbeddings(hidden_size=3, vocab_size=3)
@pytest.fixture
def input_ids(self):
return torch.tensor([[1, 2], [0, 2]])
def test_forward(self, input_ids, text_embedding):
embs = text_embedding(input_ids)
actual = embs.shape
expected = torch.Size([2, 2, 3])
assert_expected(actual, expected)
def test_invalid_input(self, text_embedding):
with pytest.raises(ValueError):
_ = text_embedding()
def test_create_position_ids_from_input_ids(self, input_ids, text_embedding):
actual = text_embedding.create_position_ids_from_input_ids(input_ids)
expected = torch.tensor([[1, 2], [0, 1]])
assert_expected(actual, expected)
|