evalstone/evalscope/tests/api/test_io_utils.py

import io
import json
from pathlib import Path

import numpy as np
import pyarrow as pa
import pytest
from datasets import Audio, Dataset, DatasetInfo, Features, Image, Sequence, Video, concatenate_datasets
from datasets.table import InMemoryTable
from PIL import Image as PILImage

from evalscope.utils.io_utils import jsonl_to_list, undecode_media


def _generate_jpeg_bytes(dim: int, quality: int = 80) -> bytes:
    """Return a square random-noise JPEG as raw bytes."""
    img = PILImage.fromarray(np.random.randint(0, 256, (dim, dim, 3), dtype=np.uint8))
    buf = io.BytesIO()
    img.save(buf, format='JPEG', quality=quality)
    return buf.getvalue()


def gen_dataset(
    media_factory: dict,
    extra_dict: dict | None = None,
    num_rows: int = 1,
) -> Dataset:
    features = {}
    row = {}

    # make media cols
    for col_name, (media_bytes, inner_feat, count) in media_factory.items():
        if count == 1:
            features[col_name] = inner_feat
            row[col_name] = {'bytes': media_bytes, 'path': None}
        else:
            features[col_name] = Sequence(inner_feat)
            row[col_name] = [{'bytes': media_bytes, 'path': None}] * count

    # inject non-media cols
    if extra_dict:
        features |= Dataset.from_list([extra_dict]).features
        row.update(extra_dict)

    dataset_features = Features(features)
    table = pa.Table.from_pylist([row], schema=dataset_features.arrow_schema)
    single_row = Dataset(InMemoryTable(table), info=DatasetInfo(features=dataset_features))
    return concatenate_datasets([single_row] * num_rows)


class TestArrowOffsetOverflow:
    """The default ``batch_size=1000`` in ``datasets`` can overflow Arrow's
    32-bit offset when a dataset contains many large images per row.

    These tests reproduce the overflow and confirm that a smaller batch
    size avoids it.
    """

    IMAGES_PER_ROW = 100

    @pytest.mark.parametrize(
        'dim, batch_size',
        [(512, None), (1024, 10)],
    )
    def test_default_batch_overflows_but_small_batch_succeeds(
        self,
        dim: int,
        batch_size: int | None,
    ):
        ds = gen_dataset(
            media_factory={'images': (_generate_jpeg_bytes(dim), Image(decode=True), self.IMAGES_PER_ROW)},
            num_rows=1000,
        )

        # expected overflow with default batch size (1k)
        with pytest.raises(ValueError, match='offset'):
            undecode_media(ds, media_type=['image'], batch_size=1000)

        # A smaller batch size avoids the overflow. One full batch is sufficient to verify the boundary.
        safe_batch_size = batch_size or 100
        result = undecode_media(
            ds.select(range(safe_batch_size)), media_type=['image'], batch_size=batch_size
        )
        assert result.features['images'].feature.decode is False


class TestUndecodeMediaIntegration:
    """End-to-end tests for ``undecode_media``.

    Every test verifies that:
    • target media columns get ``decode=False``
    • non-media columns keep their original feature type and data
    """

    @pytest.mark.parametrize(
        'example',
        [
            {'text': 'hello', 'label': 0},
            {'messages': [{'role': 'user', 'content': 'world'}], 'answer': [1]},
            {'tokens': [1, 2, 3, 4], 'answer': {'text': 'hello'}},
        ],
    )
    def test_leave_non_media_columns_unchanged(self, example: dict):
        """Given a dataset with only text / numeric columns,
        when undecode_media is called,
        then the original object is returned and features are preserved."""
        ds = Dataset.from_list([example])
        undecoded_ds = undecode_media(ds, media_type=['image', 'audio', 'video'])

        assert undecoded_ds.features == ds.features
        assert undecoded_ds[0] == example

    @pytest.mark.parametrize(
        'media_factory, non_media_data',
        [
            pytest.param(
                {'image': (_generate_jpeg_bytes(40), Image(decode=True), 1)},
                {'text': 'cat'},
                id='plain_Image',
            ),
            pytest.param(
                {'images': (_generate_jpeg_bytes(20), Image(decode=True), 4)},
                {'text': 'album'},
                id='Sequence_Image',
            ),
            pytest.param(
                {'audio': (b'dummy', Audio(decode=True), 1)},
                {'text': 'recording'},
                id='plain_Audio',
            ),
            pytest.param(
                {'audios': (b'dummy', Audio(decode=True), 2)},
                {'text': 'playlist'},
                id='Sequence_Audio',
            ),
            pytest.param(
                {'video': (b'dummy', Video(decode=True), 1)},
                {'text': 'clip'},
                id='plain_Video',
            ),
            pytest.param(
                {'videos': (b'dummy', Video(decode=True), 2)},
                {'text': 'clips'},
                id='Sequence_Video',
            ),
            pytest.param(
                {
                    'image': (_generate_jpeg_bytes(40), Image(decode=True), 1),
                    'audio': (b'dummy', Audio(decode=True), 1),
                },
                {'text': 'multi'},
                id='multiple_media_types',
            ),
        ],
    )
    def test_it_disables_decode_on_media_columns_and_preserves_others(
        self,
        media_factory: dict,
        non_media_data: dict,
    ):
        """Given a dataset with a media column (decode=True) and a text column,
        when undecode_media is called with the relevant media_type(s),
        then the media column's decode flag is set to False,
        and the text column is untouched."""
        ds = gen_dataset(media_factory, extra_dict=non_media_data)
        result = undecode_media(ds, media_type=['image', 'audio', 'video'])

        # -- Media columns: decode must be False --
        for col_name in media_factory:
            media_feat = result.features[col_name]
            if isinstance(media_feat, Sequence):
                assert isinstance(media_feat.feature, (Image, Audio, Video))
                assert media_feat.feature.decode is False
            else:
                assert isinstance(media_feat, (Image, Audio, Video))
                assert media_feat.decode is False

        for other_col in non_media_data:
            assert other_col in result.features, f'Column {other_col} missing'
            assert result.features[other_col] == ds.features[other_col]

        original_row = ds.to_list()[0]
        result_row = result.to_list()[0]
        for other_col in non_media_data:
            assert result_row[other_col] == original_row[other_col]


class TestJsonlToList:
    @staticmethod
    def _write(path: Path, content: str) -> str:
        path.write_text(content, encoding='utf-8')
        return str(path)

    def test_tolerant_read_skips_torn_tail(self, tmp_path: Path) -> None:
        file_path = self._write(tmp_path / 'torn.jsonl', '{"a": 1}\n{"a": 2}\n{"a": 3')

        records = jsonl_to_list(file_path, skip_invalid=True)

        assert records == [{'a': 1}, {'a': 2}]

    def test_tolerant_read_skips_null_empty_and_non_dict_json(self, tmp_path: Path) -> None:
        file_path = self._write(
            tmp_path / 'invalid-types.jsonl',
            '{"a": 1}\n\nnull\n[]\n"text"\n42\n{"a": 2}\n',
        )

        records = jsonl_to_list(file_path, skip_invalid=True)

        assert records == [{'a': 1}, {'a': 2}]

    def test_default_read_remains_strict_for_malformed_json(self, tmp_path: Path) -> None:
        file_path = self._write(tmp_path / 'strict.jsonl', '{"a": 1}\n{"a": 2')

        with pytest.raises(json.JSONDecodeError):
            jsonl_to_list(file_path)