otter

1
import unittest
2
from unittest.mock import Mock
3
from pipeline.mimicit_utils.data import get_mmc4_dataset
4

5

6
class TestGetMMC4Dataset(unittest.TestCase):
7
    def test_get_mmc4_dataset(self):
8
        # Mock the required inputs
9
        args = Mock(
10
            mmc4_shards="/home/luodian/projects/Otter/archived/000000000.tar",
11
            train_num_samples_mmc4=1000,
12
            mmc4_textsim_threshold=0.32,
13
            batch_size_mmc4=10,
14
            seed=0,
15
            workers=2,
16
            world_size=1,
17
        )
18
        image_processor = Mock()
19
        tokenizer = Mock()
20

21
        # Call the function to test
22
        data_info = get_mmc4_dataset(args, image_processor, tokenizer)
23

24
        # Check if the dataloader's attributes are as expected
25
        self.assertEqual(data_info.dataloader.num_batches, 100)
26
        self.assertEqual(data_info.dataloader.num_samples, 1000)
27