mirror of
https://github.com/modelscope/modelscope.git
synced 2026-08-29 10:08:40 +02:00
134 lines
5.5 KiB
Python
134 lines
5.5 KiB
Python
# Copyright (c) Alibaba, Inc. and its affiliates.
|
|
|
|
import unittest
|
|
from unittest import mock
|
|
|
|
from huggingface_hub.hf_file_system import HfFileSystem
|
|
|
|
from modelscope import MsDataset
|
|
from modelscope.utils.logger import get_logger
|
|
from modelscope.utils.test_utils import test_level
|
|
|
|
logger = get_logger()
|
|
|
|
|
|
class TestStreamLoad(unittest.TestCase):
|
|
|
|
@staticmethod
|
|
def _reset_hf_filesystem_patch(hf_datasets_util):
|
|
if (HfFileSystem._open is hf_datasets_util._hf_fs_open
|
|
and hf_datasets_util._hf_fs_open_original is not None):
|
|
HfFileSystem._open = hf_datasets_util._hf_fs_open_original
|
|
hf_datasets_util._hf_fs_open_original = None
|
|
if (HfFileSystem.__init__ is hf_datasets_util._hf_fs_init_with_cookie
|
|
and hf_datasets_util._hf_fs_init_original is not None):
|
|
HfFileSystem.__init__ = hf_datasets_util._hf_fs_init_original
|
|
hf_datasets_util._hf_fs_init_original = None
|
|
|
|
def test_hf_filesystem_patch_idempotent_for_repeated_streaming_loads(self):
|
|
from modelscope.msdatasets.utils import hf_datasets_util
|
|
|
|
hf_fs_open_before = HfFileSystem._open
|
|
hf_fs_init_before = HfFileSystem.__init__
|
|
open_original_before = hf_datasets_util._hf_fs_open_original
|
|
init_original_before = hf_datasets_util._hf_fs_init_original
|
|
try:
|
|
self._reset_hf_filesystem_patch(hf_datasets_util)
|
|
with mock.patch.object(
|
|
hf_datasets_util.DatasetsWrapperHF,
|
|
'load_dataset',
|
|
return_value=object()):
|
|
with hf_datasets_util.load_dataset_with_ctx(streaming=True):
|
|
pass
|
|
with hf_datasets_util.load_dataset_with_ctx(streaming=True):
|
|
pass
|
|
|
|
self.assertIs(HfFileSystem._open, hf_datasets_util._hf_fs_open)
|
|
self.assertIsNot(hf_datasets_util._hf_fs_open_original,
|
|
hf_datasets_util._hf_fs_open)
|
|
self.assertIs(HfFileSystem.__init__,
|
|
hf_datasets_util._hf_fs_init_with_cookie)
|
|
self.assertIsNot(hf_datasets_util._hf_fs_init_original,
|
|
hf_datasets_util._hf_fs_init_with_cookie)
|
|
finally:
|
|
HfFileSystem._open = hf_fs_open_before
|
|
HfFileSystem.__init__ = hf_fs_init_before
|
|
hf_datasets_util._hf_fs_open_original = open_original_before
|
|
hf_datasets_util._hf_fs_init_original = init_original_before
|
|
|
|
def test_hf_filesystem_patch_restored_when_streaming_load_fails(self):
|
|
from modelscope.msdatasets.utils import hf_datasets_util
|
|
|
|
hf_fs_open_before = HfFileSystem._open
|
|
hf_fs_init_before = HfFileSystem.__init__
|
|
open_original_before = hf_datasets_util._hf_fs_open_original
|
|
init_original_before = hf_datasets_util._hf_fs_init_original
|
|
try:
|
|
self._reset_hf_filesystem_patch(hf_datasets_util)
|
|
hf_fs_open_clean = HfFileSystem._open
|
|
hf_fs_init_clean = HfFileSystem.__init__
|
|
with mock.patch.object(
|
|
hf_datasets_util.DatasetsWrapperHF,
|
|
'load_dataset',
|
|
side_effect=RuntimeError('load failed')):
|
|
with self.assertRaises(RuntimeError):
|
|
with hf_datasets_util.load_dataset_with_ctx(
|
|
streaming=True):
|
|
pass
|
|
|
|
self.assertIs(HfFileSystem._open, hf_fs_open_clean)
|
|
self.assertIs(HfFileSystem.__init__, hf_fs_init_clean)
|
|
self.assertIsNone(hf_datasets_util._hf_fs_open_original)
|
|
self.assertIsNone(hf_datasets_util._hf_fs_init_original)
|
|
finally:
|
|
HfFileSystem._open = hf_fs_open_before
|
|
HfFileSystem.__init__ = hf_fs_init_before
|
|
hf_datasets_util._hf_fs_open_original = open_original_before
|
|
hf_datasets_util._hf_fs_init_original = init_original_before
|
|
|
|
def setUp(self):
|
|
...
|
|
|
|
def tearDown(self):
|
|
...
|
|
|
|
@unittest.skipUnless(test_level() >= 1, 'skip test in current test level')
|
|
def test_stream_read_zstd(self):
|
|
repo_id: str = 'swift/chinese-c4'
|
|
ds = MsDataset.load(repo_id, split='train', use_streaming=True)
|
|
sample = next(iter(ds))
|
|
logger.info(sample)
|
|
|
|
assert sample['url'], f'Failed to load sample from {repo_id}'
|
|
|
|
@unittest.skipUnless(test_level() >= 1, 'skip test in current test level')
|
|
def test_stream_imagefolder(self):
|
|
repo_id: str = 'wangxingjun778/test_new_dataset'
|
|
ds = MsDataset.load(repo_id, split='train', use_streaming=True)
|
|
sample = next(iter(ds))
|
|
logger.info(sample)
|
|
|
|
assert sample['image'], f'Failed to load sample from {repo_id}'
|
|
|
|
@unittest.skipUnless(test_level() >= 2, 'skip test in current test level')
|
|
def test_stream_parquet(self):
|
|
repo_id: str = 'swift/A-OKVQA'
|
|
ds = MsDataset.load(repo_id, split='train', use_streaming=True)
|
|
sample = next(iter(ds))
|
|
logger.info(sample)
|
|
|
|
assert sample['question'], f'Failed to load sample from {repo_id}'
|
|
|
|
@unittest.skipUnless(test_level() >= 1, 'skip test in current test level')
|
|
def test_stream_swift_jsonl(self):
|
|
repo_id: str = 'iic/MSAgent-MultiRole'
|
|
ds = MsDataset.load(repo_id, split='train', use_streaming=True)
|
|
sample = next(iter(ds))
|
|
logger.info(sample)
|
|
|
|
assert sample['id'], f'Failed to load sample from {repo_id}'
|
|
|
|
|
|
if __name__ == '__main__':
|
|
unittest.main()
|