From 5a1bf20336e773aab78feadd77de156033ac0c3c Mon Sep 17 00:00:00 2001 From: supermario_leo Date: Tue, 9 Jun 2026 03:14:47 +0800 Subject: [PATCH] fix(data): resolve data_type aliases and 'auto' to dispatchable dataset ids MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit DataLoaderFactory.create_data_loader documented 'text', 'multimodal' and 'auto' as valid data_type values (with 'auto' as the parameter default), but auto-detection produced 'text'/'multimodal' while the dataset dispatch only matched the class-name keys ('TextDataset', 'MultiModalDataset', ...). As a result every create_data_loader(data_type='auto') call — and the documented short aliases — fell through to the final else branch and raised ValueError: Unsupported data type. Extract _resolve_data_type() to perform auto-detection and normalize the documented short aliases to canonical dataset ids via a DATA_TYPE_ALIASES map, validated against SUPPORTED_DATA_TYPES with a descriptive error. Config-supplied class names pass through unchanged. Add CPU-only regression tests covering auto detection, alias resolution, class-name pass-through and the unknown-type error. --- angelslim/data/dataloader.py | 63 +++++++++++++--- tests/test_dataloader.py | 140 +++++++++++++++++++++++++++++++++++ 2 files changed, 194 insertions(+), 9 deletions(-) create mode 100644 tests/test_dataloader.py diff --git a/angelslim/data/dataloader.py b/angelslim/data/dataloader.py index 2bd4cf93..a9e8c617 100644 --- a/angelslim/data/dataloader.py +++ b/angelslim/data/dataloader.py @@ -25,10 +25,57 @@ from .text2image_dataset import Text2ImageDataset from .text_dataset import TextDataset +# Canonical dataset identifiers accepted by ``create_data_loader``. +SUPPORTED_DATA_TYPES = ( + "TextDataset", + "MultiModalDataset", + "Text2ImageDataset", + "OmniDataset", + "AudioDataset", +) + +# Friendly aliases (including the values produced by ``"auto"`` detection and +# the short names documented in ``create_data_loader``) mapped to their +# canonical dataset identifier. +DATA_TYPE_ALIASES = { + "text": "TextDataset", + "multimodal": "MultiModalDataset", + "text2image": "Text2ImageDataset", + "omni": "OmniDataset", + "audio": "AudioDataset", +} + class DataLoaderFactory: """Factory for creating PyTorch DataLoaders from various data sources""" + @staticmethod + def _resolve_data_type(data_type: str, data_source: Union[str, Dict]) -> str: + """Resolve a user-supplied ``data_type`` to a canonical dataset id. + + Handles ``"auto"`` detection and normalizes the documented short + aliases (``"text"``, ``"multimodal"``, ...) to the dataset class names + used by the dispatch in :meth:`create_data_loader`. Raises a clear + ``ValueError`` for anything unrecognized. + """ + if data_type == "auto": + if isinstance(data_source, str) and ( + ".parquet" in data_source.lower() or ".json" in data_source.lower() + ): + data_type = "text" + else: + data_type = "multimodal" + + data_type = DATA_TYPE_ALIASES.get(data_type, data_type) + + if data_type not in SUPPORTED_DATA_TYPES: + raise ValueError( + f"Unsupported data type: {data_type}. Expected one of " + f"{list(SUPPORTED_DATA_TYPES)}, an alias in " + f"{sorted(DATA_TYPE_ALIASES)}, or 'auto'." + ) + return data_type + @staticmethod def create_data_loader( processor: ProcessorMixin, @@ -58,7 +105,11 @@ def create_data_loader( shuffle: Whether to shuffle data num_samples: Limit number of samples (-1 for all) data_source: File path or HF dataset dict - data_type: "text", "multimodal" or "auto" + data_type: a canonical dataset id (one of + ``TextDataset``, ``MultiModalDataset``, ``Text2ImageDataset``, + ``OmniDataset``, ``AudioDataset``), a short alias + (``text``, ``multimodal``, ``text2image``, ``omni``, ``audio``) + or ``auto`` to infer from ``data_source`` num_workers: Number of workers for DataLoader inference_settings: Settings for text-to-image inference - height: Image height @@ -70,14 +121,8 @@ def create_data_loader( Returns: PyTorch DataLoader ready for use """ - # Auto detect data type if not specified - if data_type == "auto": - if isinstance(data_source, str) and ( - ".parquet" in data_source.lower() or ".json" in data_source - ): - data_type = "text" - else: - data_type = "multimodal" + # Resolve "auto" detection and documented aliases to a canonical id + data_type = DataLoaderFactory._resolve_data_type(data_type, data_source) # Create appropriate dataset if data_type == "TextDataset": diff --git a/tests/test_dataloader.py b/tests/test_dataloader.py new file mode 100644 index 00000000..ee2f7e8a --- /dev/null +++ b/tests/test_dataloader.py @@ -0,0 +1,140 @@ +# Copyright 2025 Tencent Inc. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Unit tests for ``DataLoaderFactory`` data-type resolution. + +These tests are CPU-only and require neither a GPU, model weights, nor the +heavy ``torch``/``transformers`` stack: the dataset classes and third-party +imports pulled in by ``angelslim.data.dataloader`` are stubbed so that the +pure ``_resolve_data_type`` string logic can be exercised in isolation. +""" + +import importlib.util +import os +import sys +import types + +import pytest + +_REPO_ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) +_DATALOADER_PATH = os.path.join(_REPO_ROOT, "angelslim", "data", "dataloader.py") + + +def _install_stubs(): + """Register lightweight stand-ins so ``dataloader.py`` imports cleanly.""" + + def _module(name): + mod = types.ModuleType(name) + sys.modules[name] = mod + return mod + + # torch.utils.data.DataLoader + if "torch" not in sys.modules: + torch = _module("torch") + torch.utils = _module("torch.utils") + torch_data = _module("torch.utils.data") + torch_data.DataLoader = object + torch.utils.data = torch_data + + # transformers.ProcessorMixin + if "transformers" not in sys.modules: + transformers = _module("transformers") + transformers.ProcessorMixin = object + + # angelslim.data. siblings imported by dataloader.py + for pkg in ("angelslim", "angelslim.data"): + if pkg not in sys.modules: + mod = _module(pkg) + mod.__path__ = [] # mark as package + for leaf, cls in ( + ("audio_dataset", "AudioDataset"), + ("base_dataset", "BaseDataset"), + ("multimodal_dataset", "MultiModalDataset"), + ("omni_dataset", "OmniDataset"), + ("text2image_dataset", "Text2ImageDataset"), + ("text_dataset", "TextDataset"), + ): + name = f"angelslim.data.{leaf}" + if name not in sys.modules: + mod = _module(name) + setattr(mod, cls, type(cls, (), {})) + + +def _load_dataloader_module(): + _install_stubs() + spec = importlib.util.spec_from_file_location("angelslim.data.dataloader", _DATALOADER_PATH) + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + return module + + +@pytest.fixture(scope="module") +def factory(): + return _load_dataloader_module().DataLoaderFactory + + +@pytest.mark.parametrize( + "source,expected", + [ + ("calib.json", "TextDataset"), + ("CALIB.JSON", "TextDataset"), + ("data.parquet", "TextDataset"), + ("images/", "MultiModalDataset"), + ({"split": "train"}, "MultiModalDataset"), + ], +) +def test_auto_resolves_to_a_dispatchable_id(factory, source, expected): + """``"auto"`` must resolve to a canonical id the dispatch can match. + + Regression: auto-detection previously produced ``"text"``/``"multimodal"`` + which no dispatch branch matched, so every ``data_type="auto"`` call raised + ``ValueError: Unsupported data type``. + """ + assert factory._resolve_data_type("auto", source) == expected + + +@pytest.mark.parametrize( + "alias,expected", + [ + ("text", "TextDataset"), + ("multimodal", "MultiModalDataset"), + ("text2image", "Text2ImageDataset"), + ("omni", "OmniDataset"), + ("audio", "AudioDataset"), + ], +) +def test_documented_short_aliases_resolve(factory, alias, expected): + """The short aliases named in the docstring must be accepted.""" + assert factory._resolve_data_type(alias, "calib.json") == expected + + +@pytest.mark.parametrize( + "canonical", + [ + "TextDataset", + "MultiModalDataset", + "Text2ImageDataset", + "OmniDataset", + "AudioDataset", + ], +) +def test_canonical_ids_pass_through(factory, canonical): + """Existing config-supplied class names must remain valid (no regression).""" + assert factory._resolve_data_type(canonical, "calib.json") == canonical + + +def test_unknown_data_type_raises_value_error(factory): + """An unrecognized value raises a descriptive error, not a silent miss.""" + with pytest.raises(ValueError, match="Unsupported data type"): + factory._resolve_data_type("nonexistent", "calib.json")