|
|
import os |
|
|
from unittest import TestCase |
|
|
|
|
|
import torch |
|
|
|
|
|
from apps.plm.dataset_conf import dataset_config as DATASET_CONFIGS |
|
|
from apps.plm.tokenizer import PLMTokenizer |
|
|
from core.data.dataloader import DataloadArgs, get_dataloader |
|
|
|
|
|
|
|
|
|
|
|
class DataloaderTest(TestCase): |
|
|
def setUp(self): |
|
|
self.seq_len = 8196 |
|
|
self.patch_size = 14 |
|
|
self.pooling_ratio = 2 |
|
|
self.max_num_tiles = 9 |
|
|
self.image_res = 448 |
|
|
self.mllm_tokenizer = PLMTokenizer( |
|
|
os.environ["TOKENIZER_PATH"], |
|
|
seq_len=self.seq_len, |
|
|
patch_size=self.patch_size, |
|
|
pooling_ratio=self.pooling_ratio, |
|
|
) |
|
|
|
|
|
def test_jsonl_image_text_dataloader(self): |
|
|
dataloader_args = DataloadArgs( |
|
|
datamix="dummy_image:1", |
|
|
num_workers=1, |
|
|
vision_input_type="thumb+tile", |
|
|
image_res=self.image_res, |
|
|
max_num_tiles=self.max_num_tiles, |
|
|
batch_size=1, |
|
|
) |
|
|
dataloader = get_dataloader( |
|
|
dataloader_args, |
|
|
dp_rank=0, |
|
|
dp_world_size=1, |
|
|
dataset_configs=DATASET_CONFIGS, |
|
|
tokenizer=self.mllm_tokenizer, |
|
|
) |
|
|
batch_iterator = iter(dataloader) |
|
|
expected_num_image_tokens = ( |
|
|
self.image_res // self.patch_size // self.pooling_ratio |
|
|
) ** 2 * (self.max_num_tiles + 1) |
|
|
print(f"expected_num_image_tokens: {expected_num_image_tokens}") |
|
|
for i in range(3): |
|
|
mllm_batch = next(batch_iterator) |
|
|
|
|
|
image_token_mask = mllm_batch.x == self.mllm_tokenizer.image_token_id |
|
|
num_image_tokens = image_token_mask.sum(dim=1) |
|
|
self.assertTrue((mllm_batch.image_pos_index[~image_token_mask] == -1).all()) |
|
|
|
|
|
for i in range(mllm_batch.x.shape[0]): |
|
|
print(f"num_image_tokens in example {i}", num_image_tokens[i]) |
|
|
self.assertEqual( |
|
|
num_image_tokens[i], |
|
|
expected_num_image_tokens, |
|
|
) |
|
|
cur_x_is_image = image_token_mask[i] |
|
|
cur_image_pos_index = mllm_batch.image_pos_index[i] |
|
|
self.assertTrue( |
|
|
torch.equal( |
|
|
cur_image_pos_index[cur_x_is_image], |
|
|
torch.arange(num_image_tokens[i]), |
|
|
) |
|
|
) |
|
|
|
