| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """Bailing prompt expansion tests without downloading a checkpoint.""" |
| 4 | |
| 5 | from types import SimpleNamespace |
| 6 | |
| 7 | import pytest |
| 8 | import torch |
| 9 | from PIL import Image |
| 10 | from tokenizers import Tokenizer, decoders, models, pre_tokenizers |
| 11 | from transformers import TokenizersBackend |
| 12 | |
| 13 | from vllm.model_executor.models.bailing_moe_v3_vl import ( |
| 14 | IMAGE_PLACEHOLDER, |
| 15 | IMAGE_TOKEN, |
| 16 | BailingMoeV3VLMultiModalProcessor, |
| 17 | ) |
| 18 | from vllm.multimodal.inputs import MultiModalFieldConfig, MultiModalKwargsItems |
| 19 | from vllm.multimodal.parse import ImageProcessorItems, MultiModalDataItems |
| 20 | |
| 21 | pytestmark = pytest.mark.cpu_test |
| 22 | |
| 23 | |
| 24 | @pytest.fixture |
| 25 | def processor(): |
| 26 | vocab = { |
| 27 | token: index |
| 28 | for index, token in enumerate(sorted(pre_tokenizers.ByteLevel.alphabet())) |
| 29 | } |
| 30 | # A newline added by the processor must merge with an existing newline. |
| 31 | vocab["ĊĊ"] = len(vocab) |
| 32 | backend = Tokenizer(models.BPE(vocab=vocab, merges=[("Ċ", "Ċ")])) |
| 33 | backend.pre_tokenizer = pre_tokenizers.ByteLevel( |
| 34 | add_prefix_space=False, use_regex=False |
| 35 | ) |
| 36 | backend.decoder = decoders.ByteLevel() |
| 37 | tokenizer = TokenizersBackend( |
| 38 | tokenizer_object=backend, |
| 39 | additional_special_tokens=[ |
| 40 | "<|vision_start|>", |
| 41 | IMAGE_TOKEN, |
| 42 | "<|vision_end|>", |
| 43 | ], |
| 44 | clean_up_tokenization_spaces=False, |
| 45 | ) |
| 46 | processor = object.__new__(BailingMoeV3VLMultiModalProcessor) |
| 47 | processor.info = SimpleNamespace( |
| 48 | get_tokenizer=lambda: tokenizer, |
| 49 | get_hf_config=lambda: SimpleNamespace( |
| 50 | vision_config=SimpleNamespace(spatial_merge_size=2) |
| 51 | ), |
| 52 | ) |
| 53 | return processor |
| 54 | |
| 55 | |
| 56 | @pytest.mark.parametrize( |
| 57 | ("prompt", "expected_text", "token_counts"), |
| 58 | [ |
| 59 | ( |
| 60 | f"before{IMAGE_PLACEHOLDER}after", |
| 61 | f"before<|vision_start|>{IMAGE_TOKEN * 2}<|vision_end|>\nafter", |
| 62 | [2], |
| 63 | ), |
| 64 | ( |
| 65 | f"before{IMAGE_TOKEN}after", |
| 66 | f"before<|vision_start|>{IMAGE_TOKEN * 2}<|vision_end|>\nafter", |
| 67 | [2], |
| 68 | ), |
| 69 | ( |
| 70 | f"before{IMAGE_PLACEHOLDER}\nafter", |
| 71 | f"before<|vision_start|>{IMAGE_TOKEN * 2}<|vision_end|>\n\nafter", |
| 72 | [2], |
| 73 | ), |
| 74 | ( |
| 75 | f"before{IMAGE_PLACEHOLDER}{IMAGE_PLACEHOLDER}after", |
| 76 | f"before<|vision_start|>{IMAGE_TOKEN * 2}<|vision_end|>\n" |
| 77 | f"<|vision_start|>{IMAGE_TOKEN * 3}<|vision_end|>\nafter", |
| 78 | [2, 3], |
| 79 | ), |
| 80 | ], |
| 81 | ids=["wrapped", "bare", "existing_newline", "adjacent_images"], |
| 82 | ) |
| 83 | def test_image_expansion_matches_wrappers_and_newline_tokenization( |
| 84 | processor, prompt, expected_text, token_counts |
| 85 | ): |
| 86 | tokenizer = processor.info.get_tokenizer() |
| 87 | image_id = tokenizer.get_vocab()[IMAGE_TOKEN] |
| 88 | grids = torch.tensor([[1, 2, 2 * count] for count in token_counts]) |
| 89 | mm_kwargs = MultiModalKwargsItems.from_hf_inputs( |
| 90 | {"image_grid_thw": grids}, |
| 91 | {"image_grid_thw": MultiModalFieldConfig.batched("image")}, |
| 92 | ) |
| 93 | mm_items = MultiModalDataItems( |
| 94 | {"image": ImageProcessorItems([Image.new("RGB", (1, 1))] * len(token_counts))} |
| 95 | ) |
| 96 | updates = processor._get_mm_prompt_updates(mm_items, {}, mm_kwargs) |
| 97 | prompt_ids = tokenizer.encode(prompt, add_special_tokens=False) |
| 98 | expected_ids = tokenizer.encode(expected_text, add_special_tokens=False) |
| 99 | |
| 100 | actual_ids, placeholders = processor._apply_prompt_updates(prompt_ids, updates) |
| 101 | |
| 102 | assert actual_ids == expected_ids |
| 103 | assert len(placeholders["image"]) == len(token_counts) |
| 104 | found = processor._find_mm_placeholders(actual_ids, updates) |
| 105 | search_offset = 0 |
| 106 | for index, count in enumerate(token_counts): |
| 107 | placeholder = placeholders["image"][index] |
| 108 | expected_offset = expected_ids.index(image_id, search_offset) |
| 109 | assert placeholder.start_idx == expected_offset |
| 110 | assert placeholder.tokens == [image_id] * count |
| 111 | assert placeholder.to_range().get_num_embeds() == count |
| 112 | assert placeholder.to_range() == found["image"][index].to_range() |
| 113 | search_offset = expected_offset + count |
| 114 | |
| 115 | |
| 116 | def test_text_only_prompt_preserves_original_token_ids(processor): |
| 117 | tokenizer = processor.info.get_tokenizer() |
| 118 | prompt_ids = tokenizer.encode("before\n", add_special_tokens=False) |
| 119 | prompt_ids += tokenizer.encode("\nafter", add_special_tokens=False) |
| 120 | # Deliberately use two newline tokens instead of their merged BPE token. |
| 121 | assert prompt_ids != tokenizer.encode("before\n\nafter", add_special_tokens=False) |
| 122 | |
| 123 | actual_ids, placeholders = processor._apply_prompt_updates(prompt_ids, {}) |
| 124 | |
| 125 | assert actual_ids == prompt_ids |
| 126 | assert placeholders == {} |