| import unittest |
| from types import SimpleNamespace |
|
|
| import torch |
| from torch import nn |
|
|
| from model.llava.prima_arch import LlavaMetaForCausalLM |
| from tools.utils import IGNORE_INDEX, IMAGE_TOKEN_INDEX |
|
|
|
|
| class TinyMultimodalModel(nn.Module, LlavaMetaForCausalLM): |
| def __init__(self, hidden_size=4, image_tokens=3): |
| super().__init__() |
| self.config = SimpleNamespace(with_region=False, mm_use_im_start_end=False, |
| tune_mm_mlp_adapter=False) |
| self.model = nn.Module() |
| self.model.embed_tokens = nn.Embedding(32, hidden_size) |
| self.model.mm_projector = nn.Linear(hidden_size, hidden_size) |
| self.model.vision_tower = nn.Identity() |
| self.model.get_vision_tower = lambda: self.model.vision_tower |
| self.model.use_qformer = True |
| self.model.Qformer = SimpleNamespace(config=SimpleNamespace(hidden_size=hidden_size)) |
| self.model.device = torch.device('cpu') |
| self.model.dtype = torch.float32 |
| self.image_features = nn.Parameter(torch.arange(image_tokens * hidden_size).reshape( |
| image_tokens, hidden_size).float() / 10) |
| self.encoder_calls = 0 |
|
|
| @property |
| def device(self): |
| return self.model.embed_tokens.weight.device |
|
|
| def get_model(self): |
| return self.model |
|
|
| def encode_images_qformer(self, images, *args): |
| self.encoder_calls += 1 |
| values = images.flatten(1).mean(1)[:, None, None] |
| return self.image_features.unsqueeze(0) + values, images |
|
|
|
|
| def prepare(model, ids, *, mask=None, labels=None, images=None, past=None): |
| offsets = torch.arange(len(ids) + 1) |
| return model.prepare_inputs_labels_for_multimodal( |
| ids, mask, past, labels, images, offsets, offsets, None, None) |
|
|
|
|
| class MultimodalTests(unittest.TestCase): |
| def test_image_features_replace_only_the_image_token(self): |
| model = TinyMultimodalModel() |
| ids = torch.tensor([[1, IMAGE_TOKEN_INDEX, 2, 3]]) |
| labels = ids.clone() |
| result = prepare(model, ids, mask=torch.ones_like(ids, dtype=torch.bool), |
| labels=labels, images=torch.zeros(1, 3, 2, 2)) |
| self.assertEqual(result[3].shape, (1, 6, 4)) |
| self.assertTrue(torch.equal(result[3][0, 1:4], model.image_features)) |
| self.assertTrue(torch.equal(result[3][0, [0, 4, 5]], model.model.embed_tokens(ids[0, [0, 2, 3]]))) |
| self.assertEqual(result[4].tolist(), [[1, IGNORE_INDEX, IGNORE_INDEX, IGNORE_INDEX, 2, 3]]) |
| self.assertEqual(result[5], [[slice(1, 4)]]) |
| self.assertTrue(result[1].all()) |
| result[3].sum().backward() |
| self.assertTrue(torch.equal(model.image_features.grad, torch.ones_like(model.image_features))) |
|
|
| def test_ragged_batch_preserves_left_padding_and_masks_new_right_padding(self): |
| model = TinyMultimodalModel() |
| ids = torch.tensor([[0, 1, IMAGE_TOKEN_INDEX, 2], [0, 0, 3, 4]]) |
| mask = ids.ne(0) |
| labels = ids.masked_fill(~mask, IGNORE_INDEX) |
| for include_labels in (True, False): |
| with self.subTest(labels=include_labels): |
| result = prepare(model, ids, mask=mask, |
| labels=labels if include_labels else None, |
| images=torch.zeros(1, 3, 2, 2)) |
| self.assertEqual(result[1].tolist(), [[False, True, True, True, True, True], |
| [False, False, True, True, False, False]]) |
| self.assertTrue(torch.equal(result[3][1, -2:], torch.zeros(2, 4))) |
| if include_labels: |
| self.assertEqual(result[4][1, -2:].tolist(), [IGNORE_INDEX, IGNORE_INDEX]) |
| else: |
| self.assertIsNone(result[4]) |
|
|
| def test_cached_generation_extends_attention_without_encoding_images(self): |
| model = TinyMultimodalModel() |
| ids = torch.tensor([[5]]) |
| past = ((torch.zeros(1, 2, 7, 4), torch.zeros(1, 2, 7, 4)),) |
| result = prepare(model, ids, mask=torch.ones(1, 7, dtype=torch.bool), |
| images=torch.zeros(1, 3, 2, 2), past=past) |
| self.assertIs(result[0], ids) |
| self.assertIs(result[2], past) |
| self.assertIsNone(result[3]) |
| self.assertEqual(result[1].shape, (1, 8)) |
| self.assertTrue(result[1].all()) |
| self.assertEqual(model.encoder_calls, 0) |
|
|
| def test_ragged_embeddings_preserve_dtype_promotion(self): |
| model = TinyMultimodalModel() |
| model.model.embed_tokens.to(dtype=torch.bfloat16) |
| ids = torch.tensor([[1, 2, 3, 4], [1, IMAGE_TOKEN_INDEX, 2, 3]]) |
| result = prepare(model, ids, mask=torch.ones_like(ids, dtype=torch.bool), |
| images=torch.zeros(1, 3, 2, 2)) |
| self.assertEqual(result[3].dtype, torch.float32) |
| self.assertTrue(torch.equal(result[3][1, 1:4], model.image_features)) |
|
|
| def test_grouped_images_preserve_feature_order(self): |
| model = TinyMultimodalModel() |
| ids = torch.tensor([[1, IMAGE_TOKEN_INDEX, 2], [3, IMAGE_TOKEN_INDEX, 4]]) |
| groups = [torch.stack((torch.zeros(3, 2, 2), torch.ones(3, 2, 2))), |
| torch.full((1, 3, 2, 2), 2.0)] |
| result = prepare(model, ids, mask=torch.ones_like(ids, dtype=torch.bool), images=groups) |
| self.assertEqual(result[3].shape, (2, 8, 4)) |
| self.assertTrue(torch.equal(result[3][0, 1:4], model.image_features)) |
| self.assertTrue(torch.equal(result[3][0, 4:7], model.image_features + 1)) |
| self.assertTrue(torch.equal(result[3][1, 1:4], model.image_features + 2)) |
|
|
|
|
| if __name__ == '__main__': |
| unittest.main() |
|
|
