wehub-resource-sync
57 lines
1.8 KiB
Python
57 lines
1.8 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
|
|
|
from types import SimpleNamespace
|
|
|
|
import torch
|
|
|
|
from vllm.model_executor.models.molmo2 import build_flat_image_bool_length
|
|
|
|
|
|
def test_build_flat_image_bool_length_matches_molmoweb_processor_tokens():
|
|
hf_config = SimpleNamespace(
|
|
image_patch_id=151938,
|
|
low_res_image_start_token_id=151940,
|
|
image_start_token_id=151936,
|
|
image_col_id=151939,
|
|
image_end_token_id=151937,
|
|
)
|
|
image_grids = torch.tensor([[14, 14, 14, 23]], dtype=torch.long)
|
|
|
|
image_tokens, num_image_tokens = build_flat_image_bool_length(
|
|
image_grids,
|
|
hf_config,
|
|
image_use_col_tokens=True,
|
|
use_single_crop_col_tokens=None,
|
|
use_single_crop_start_token=False,
|
|
)
|
|
|
|
assert num_image_tokens.tolist() == [550]
|
|
assert len(image_tokens) == 550
|
|
assert image_tokens[0].item() == hf_config.image_start_token_id
|
|
assert (image_tokens == hf_config.image_col_id).sum().item() == 28
|
|
|
|
|
|
def test_build_flat_image_bool_length_respects_disabled_col_tokens():
|
|
hf_config = SimpleNamespace(
|
|
image_patch_id=151938,
|
|
low_res_image_start_token_id=151940,
|
|
image_start_token_id=151936,
|
|
image_col_id=151939,
|
|
image_end_token_id=151937,
|
|
)
|
|
image_grids = torch.tensor([[2, 3, 5, 7]], dtype=torch.long)
|
|
|
|
image_tokens, num_image_tokens = build_flat_image_bool_length(
|
|
image_grids,
|
|
hf_config,
|
|
image_use_col_tokens=False,
|
|
use_single_crop_col_tokens=False,
|
|
use_single_crop_start_token=True,
|
|
)
|
|
|
|
assert num_image_tokens.tolist() == [45]
|
|
assert len(image_tokens) == 45
|
|
assert image_tokens[0].item() == hf_config.low_res_image_start_token_id
|
|
assert (image_tokens == hf_config.image_col_id).sum().item() == 0
|