diff --git a/docs/add_new_schema.md b/docs/add_new_schema.md index faac2b5..32a3d15 100644 --- a/docs/add_new_schema.md +++ b/docs/add_new_schema.md @@ -28,9 +28,12 @@ Create a file named `schema.yaml`. Start with the basic required fields: ```yaml dataset_id: "your-dataset-id" # The unique ID of the dataset on MDC -task: "ASR" # ASR, TTS, or OTH +task: "ASR" # e.g. ASR, TTS, or OTH ``` +The `task` is used to validate that the loaded DataFrame contains +the task's required columns (e.g. ASR/TTS: `audio_path` + `transcription`). + Then add the fields for your chosen strategy. #### Example: Index-based ASR diff --git a/docs/api.md b/docs/api.md index f3de53f..89f529e 100644 --- a/docs/api.md +++ b/docs/api.md @@ -23,7 +23,15 @@ ::: datacollective.schema_loaders.cache_schema -::: datacollective.schema_loaders.tasks.asr +::: datacollective.schema_loaders.contracts -::: datacollective.schema_loaders.tasks.tts +::: datacollective.schema_loaders.strategies.index + +::: datacollective.schema_loaders.strategies.multi_split + +::: datacollective.schema_loaders.strategies.multi_sections + +::: datacollective.schema_loaders.strategies.paired_glob + +::: datacollective.schema_loaders.strategies.glob diff --git a/docs/extend_schema_loading_logic.md b/docs/extend_schema_loading_logic.md index 91f1830..77d378e 100644 --- a/docs/extend_schema_loading_logic.md +++ b/docs/extend_schema_loading_logic.md @@ -1,14 +1,16 @@ # Extending Schema Loading Logic -This document is for **developers** who want to add support for new tasks or implement new loading strategies within the MDC Python SDK. +This document is for **developers** who want to add new loading strategies or task contracts within the MDC Python SDK. -## 1. How to add a new task type +Dispatch is **strategy-based**: the schema's `root_strategy` field (default: `"index"`) selects the loader class, and any strategy can be combined with any task. The optional `task` field only adds a validation step — the loaded DataFrame must contain the task's required logical columns. -Supporting a new task (e.g., **MT** — Machine Translation) involves creating a new loader class and registering it. +## 1. How to add a new strategy + +Supporting a new strategy (e.g. **manifest** — one JSON manifest per dataset) involves creating a new loader class and registering it. ### Step 1: Create the loader class -Create a new file under `src/datacollective/schema_loaders/tasks/`, for example `mt.py`: +Create a new file under `src/datacollective/schema_loaders/strategies/`, for example `manifest.py`: ```python from __future__ import annotations @@ -17,20 +19,20 @@ import pandas as pd from datacollective.schema import DatasetSchema from datacollective.schema_loaders.base import BaseSchemaLoader -class MTLoader(BaseSchemaLoader): - """Load a machine-translation dataset.""" +class ManifestLoader(BaseSchemaLoader): + """Load a dataset described by a single JSON manifest.""" def __init__(self, schema: DatasetSchema, extract_dir: Path) -> None: super().__init__(schema, extract_dir) - # Validate required schema fields + # Validate required schema fields up front if not schema.index_file: - raise ValueError("MT schema must specify 'index_file'") + raise ValueError("manifest schema must specify 'index_file'") def load(self) -> pd.DataFrame: # BaseSchemaLoader provides shared helpers: # 1. Locate and read the index file raw_df = self._load_index_file() - + # 2. Apply column mappings and dtypes return self._apply_column_mappings(raw_df) ``` @@ -47,39 +49,49 @@ When implementing `load()`, you can leverage these methods from the base class: ### Step 3: Register the loader -Register your new class in `src/datacollective/schema_loaders/registry.py`: +1. **Add to the Enum**: add your strategy to the `Strategy` enum in `src/datacollective/schema_loaders/base.py`. +2. **Register the class** in `src/datacollective/schema_loaders/registry.py`: ```python -from datacollective.schema_loaders.tasks.mt import MTLoader +from datacollective.schema_loaders.strategies.manifest import ManifestLoader -_TASK_REGISTRY: dict[str, Type[BaseSchemaLoader]] = { - "ASR": ASRLoader, - "TTS": TTSLoader, - "OTH": OTHLoader, - "MT": MTLoader, # Add your new task here +_STRATEGY_REGISTRY: dict[Strategy, Type[BaseSchemaLoader]] = { + Strategy.INDEX: IndexLoader, + ... + Strategy.MANIFEST: ManifestLoader, # Add your new strategy here } ``` -## 2. How to extend or update strategies - -Strategies define the high-level approach to locating data (e.g., using an index file vs. globbing). +3. **Update Schema**: if the strategy requires new YAML fields, add them to `DatasetSchema` in `src/datacollective/schema.py`. ### Loading Strategies (`Strategy` enum) -Strategies are defined in the `Strategy` enum in `src/datacollective/schema_loaders/base.py`: - | Enum Member | YAML Value | Description | |---|---|---| -| `Strategy.MULTI_SPLIT` | `"multi_split"` | Loads multiple split files matching a pattern (ASR). | -| `Strategy.MULTI_SECTIONS` | `"multi_sections"` | Loads one index file per section directory, adding a `section` column (TTS). | -| `Strategy.PAIRED_GLOB` | `"paired_glob"` | Pairs audio files with sidecar files: `.txt` transcriptions (TTS) or JSON metadata/utterance files (ASR, with `format: "json"` and optional `record_path`). | -| `Strategy.GLOB` | `"glob"` | Walks directory-structured datasets, deriving metadata from the path hierarchy (OTH). | +| `Strategy.INDEX` | `"index"` (default) | Loads a single delimited index file, optionally applying column mappings. | +| `Strategy.MULTI_SPLIT` | `"multi_split"` | Loads multiple split files matching a pattern, adding a `split` column. | +| `Strategy.MULTI_SECTIONS` | `"multi_sections"` | Loads one index file per section directory, adding a `section` column. | +| `Strategy.PAIRED_GLOB` | `"paired_glob"` | Pairs audio files with sidecar files: `.txt` transcriptions, or JSON metadata/utterance files (with `format: "json"` and optional `record_path`). | +| `Strategy.GLOB` | `"glob"` | Walks directory-structured datasets, deriving metadata from the path hierarchy. | + +## 2. How to add a task contract + +The optional `task` field validates the loaded DataFrame against the task's required logical columns. Contracts live in `src/datacollective/schema_loaders/contracts.py`: + +```python +TASK_CONTRACTS: dict[str, frozenset[str]] = { + "ASR": frozenset({"audio_path", "transcription"}), + "TTS": frozenset({"audio_path", "transcription"}), + "LLM": frozenset({"text"}), + # Add your new task contract here, e.g. + # "MT": frozenset({"source_text", "target_text"}), +} +``` -### Adding a new strategy +A contract violation raises `TaskValidationError`. Schemas whose task has no contract (e.g. `OTH`), or with no task at all, load without validation. -1. **Add to the Enum**: Add your new strategy to the `Strategy` class in `base.py`. -2. **Implement Logic**: Add a branch in the relevant loader's `load()` method to handle the new strategy. -3. **Update Schema**: If the strategy requires new YAML fields, add them to `DatasetSchema` in `src/datacollective/schema.py`. +> **Note for registry schemas:** keep the `task` field in existing `schema.yaml` +> files — older SDK versions still require it. ## 3. Architecture Overview @@ -91,10 +103,12 @@ When a user calls `load_dataset("id")`: 2. **`_extract_archive()`**: Extracts it to a local directory. (Skipped if already extracted) 3. **`_resolve_schema()`**: Locates or downloads `schema.yaml`. 4. **`_parse_schema()`**: Validates YAML into a `DatasetSchema` object. -5. **`_load_dataset_from_schema()`**: +5. **`_load_dataset_from_schema()`**: - If the schema specifies `extract_files`, extracts inner archives (skipped when already extracted). - - Finds the correct loader in the **Registry**. + - Resolves the strategy loader from the **Registry** (`root_strategy`, default `"index"`). + - If the task has a contract and the schema declares column mappings, fails fast when the declared logical columns cannot satisfy the contract. - Calls `loader.load()`. + - Validates the loaded DataFrame against the task contract (when one exists). - Returns the final **pandas DataFrame**. ### Module Map @@ -102,7 +116,8 @@ When a user calls `load_dataset("id")`: | Module | Responsibility | |---|---| | `datacollective.schema` | Pydantic models and YAML parsing. | -| `datacollective.schema_loaders.base` | Abstract base class and strategy definitions. | -| `datacollective.schema_loaders.registry` | Task-to-loader mapping. | +| `datacollective.schema_loaders.base` | Abstract base class, shared helpers, and strategy definitions. | +| `datacollective.schema_loaders.registry` | Strategy-to-loader mapping and load orchestration. | +| `datacollective.schema_loaders.contracts` | Task contracts and their validation. | | `datacollective.schema_loaders.cache_schema` | Local schema caching and checksum validation. | -| `datacollective.schema_loaders.tasks.*` | Implementation of task-specific logic (ASR, TTS, OTH). | +| `datacollective.schema_loaders.strategies.*` | One loader per strategy (index, multi_split, multi_sections, paired_glob, glob). | diff --git a/docs/loaders/asr.md b/docs/loaders/asr.md index 64cefbc..918609e 100644 --- a/docs/loaders/asr.md +++ b/docs/loaders/asr.md @@ -1,5 +1,10 @@ # ASR Loader +> **Note:** since v0.6 loading strategies are **task-agnostic** — the schema's +> `root_strategy` selects the loader, and any strategy can be used with any +> task. This page describes the strategies commonly used for ASR datasets. +> The `task` field now only validates that the loaded DataFrame contains the +> task's required columns (`audio_path` + `transcription`). Loader for **ASR** (Automatic Speech Recognition) datasets. There are three parsing strategies for ASR datasets, controlled by the `root_strategy` field in the schema. diff --git a/docs/loaders/oth.md b/docs/loaders/oth.md index 9a27c72..50bd942 100644 --- a/docs/loaders/oth.md +++ b/docs/loaders/oth.md @@ -1,5 +1,10 @@ # OTH Loader +> **Note:** since v0.6 loading strategies are **task-agnostic** — the schema's +> `root_strategy` selects the loader, and any strategy can be used with any +> task. This page describes the strategies commonly used for OTH datasets. +> The `task` field now only validates that the loaded DataFrame contains the +> task's required columns (no contract — OTH loads without validation). Loader for tasks classified as **OTH** (other). There are two parsing strategies, controlled by the `root_strategy` field in the schema. diff --git a/docs/loaders/tts.md b/docs/loaders/tts.md index ea7614e..f7a4107 100644 --- a/docs/loaders/tts.md +++ b/docs/loaders/tts.md @@ -1,5 +1,10 @@ # TTS Loader +> **Note:** since v0.6 loading strategies are **task-agnostic** — the schema's +> `root_strategy` selects the loader, and any strategy can be used with any +> task. This page describes the strategies commonly used for TTS datasets. +> The `task` field now only validates that the loaded DataFrame contains the +> task's required columns (`audio_path` + `transcription`). Loader for **TTS** (Text-to-Speech) datasets. There are three parsing strategies for TTS datasets, controlled by the `root_strategy` field in the schema. diff --git a/docs/schema_documentation.md b/docs/schema_documentation.md index 756fb89..a80d7e3 100644 --- a/docs/schema_documentation.md +++ b/docs/schema_documentation.md @@ -19,13 +19,13 @@ Under the hood, `load_dataset()` performs the following steps automatically: 1. **Resolve the schema**: check local cache or the schema registry for `schema.yaml`. If the dataset is not registered this step raises a warning, so we never download an unsupported archive. 2. **Download** the archive (with resume support). The schema we fetched in step 1 tells the loader how the files are structured. 3. **Extract** the `.tar.gz` / `.zip` to a local directory. -4. **Parse** the YAML into a validated `DatasetSchema` (Pydantic model) and dispatch to the task-specific loader (ASR, TTS, OTH, …), which returns the final **DataFrame**. +4. **Parse** the YAML into a validated `DatasetSchema` (Pydantic model) and dispatch to the loader for the schema's `root_strategy` (index, glob, …), which returns the final **DataFrame**. When the schema declares a `task` with a known contract (ASR, TTS, LLM), the loaded DataFrame is validated to contain the task's required logical columns. The schema file describes: -- **What task** the dataset is for (ASR, TTS, …). - **How to find** the data files (index file path, glob pattern, etc.). - **How to map** raw columns / files into a clean DataFrame. +- Optionally, **what task** the dataset is for (ASR, TTS, …), which the loaded DataFrame is validated against. ### Minimal example @@ -56,7 +56,8 @@ Every schema **must** have: | Field | Type | Required | Description | |---|---|---|---| | `dataset_id` | `str` | ✓ | Unique dataset identifier on MDC. | -| `task` | `str` | ✓ | Task type: determines which loader is used (`"ASR"`, `"TTS"`, `"OTH"`). | +| `task` | `str` | ✗ | *(optional)* Task type as defined on the MDC Platform (`"ASR"`, `"TTS"`, …). When set to a task with a known contract, the loaded DataFrame is validated to contain the task's required logical columns (e.g. ASR/TTS: `audio_path` + `transcription`; LLM: `text`). Tasks without a contract (e.g. `"OTH"`) load without validation. | + ### Loading strategies diff --git a/src/datacollective/datasets.py b/src/datacollective/datasets.py index 6284cc4..2b80955 100644 --- a/src/datacollective/datasets.py +++ b/src/datacollective/datasets.py @@ -161,7 +161,7 @@ def load_dataset( If there is a directory with the same name as the archive file without the suffix extension, we assume it has already been extracted, and it will not be re-extracted unless `overwrite_extracted=True`. - Uses the dataset schema to determine task-specific loading logic. + Uses the dataset schema to determine the loading strategy. Automatically resumes interrupted downloads if a .checksum file exists from a previous attempt. diff --git a/src/datacollective/errors.py b/src/datacollective/errors.py index 609a1f9..f109d6c 100644 --- a/src/datacollective/errors.py +++ b/src/datacollective/errors.py @@ -42,6 +42,10 @@ def __str__(self) -> str: return "Download failed. Unfortunately this dataset does not support resuming downloads — please try again." +class TaskValidationError(ValueError): + """Raised when a loaded dataset does not satisfy its task's column contract.""" + + class MissingDependencyError(ImportError): """Raised when an optional dependency required for a feature is not installed.""" diff --git a/src/datacollective/schema.py b/src/datacollective/schema.py index fa1d892..639cec3 100644 --- a/src/datacollective/schema.py +++ b/src/datacollective/schema.py @@ -67,13 +67,13 @@ class DatasetSchema(BaseModel): """ Task-agnostic representation of a dataset schema, as defined by a ``schema.yaml`` file. - Every schema **must** have ``dataset_id`` and ``task``. The remaining - fields depend on the task type and the ``root_strategy`` - (``"index"`` vs ``"glob"``). + Every schema **must** have ``dataset_id``. The remaining fields depend on + the ``root_strategy`` (``"index"`` by default); the loader registered for + that strategy decides which fields are required at load time. - New task types only need to populate the fields they care about; - the loader registered for that task will decide which fields are - required at load time. + ``task`` is optional. When set to a task with a known contract (e.g. ASR, + TTS), the loaded DataFrame is validated to contain the task's required + logical columns. """ model_config = ConfigDict(frozen=False) @@ -81,8 +81,13 @@ class DatasetSchema(BaseModel): dataset_id: str = Field( description="Unique identifier for the dataset in the registry" ) - task: str = Field( - description="A task as defined in the MDC Platform e.g. ASR, TTS etc" + task: str | None = Field( + default=None, + description=( + "Optional task as defined in the MDC Platform e.g. ASR, TTS etc. " + "When set to a task with a known contract, the loaded dataset is " + "validated against it." + ), ) # --- Index-based strategy (ASR / TTS) --- @@ -114,9 +119,13 @@ class DatasetSchema(BaseModel): default="utf-8", description='file encoding (e.g. "utf-8-sig" for BOM)' ) - # --- Glob-based strategy (LM, paired-file TTS) --- + # --- Loading strategy --- root_strategy: str | None = Field( - default=None, description='"glob" | "paired_glob" | "multi_split"' + default=None, + description=( + '"index" (default) | "glob" | "paired_glob" | "multi_split" | ' + '"multi_sections"' + ), ) file_pattern: str | None = Field(default=None, description='e.g. "**/*.txt"') audio_extension: str | None = Field( @@ -230,9 +239,9 @@ def _parse_schema(raw: str | dict[str, Any] | Path) -> DatasetSchema: data: dict[str, Any] = raw dataset_id = data.get("dataset_id") + if not dataset_id: + raise ValueError("schema.yaml must contain 'dataset_id'") task = data.get("task") - if not dataset_id or not task: - raise ValueError("schema.yaml must contain 'dataset_id' and 'task'") # Columns (index-based) columns: dict[str, ColumnMapping] = {} @@ -286,7 +295,7 @@ def _parse_schema(raw: str | dict[str, Any] | Path) -> DatasetSchema: return DatasetSchema( dataset_id=str(dataset_id), - task=str(task).upper(), + task=str(task).upper() if task else None, format=data.get("format"), index_file=data.get("index_file"), base_audio_path=data.get("base_audio_path"), diff --git a/src/datacollective/schema_loaders/__init__.py b/src/datacollective/schema_loaders/__init__.py index f610439..f2cfbf5 100644 --- a/src/datacollective/schema_loaders/__init__.py +++ b/src/datacollective/schema_loaders/__init__.py @@ -3,8 +3,9 @@ FORMAT_SEP, Strategy, ) +from datacollective.schema_loaders.contracts import TASK_CONTRACTS from datacollective.schema_loaders.registry import ( - _get_task_loader, + _get_strategy_loader, _load_dataset_from_schema, ) @@ -12,6 +13,7 @@ "BaseSchemaLoader", "FORMAT_SEP", "Strategy", - "_get_task_loader", + "TASK_CONTRACTS", + "_get_strategy_loader", "_load_dataset_from_schema", ] diff --git a/src/datacollective/schema_loaders/base.py b/src/datacollective/schema_loaders/base.py index 406bcd3..72215d0 100644 --- a/src/datacollective/schema_loaders/base.py +++ b/src/datacollective/schema_loaders/base.py @@ -33,6 +33,7 @@ class Strategy(StrEnum): """Loading strategies recognised by schema loaders.""" + INDEX = "index" MULTI_SPLIT = "multi_split" MULTI_SECTIONS = "multi_sections" PAIRED_GLOB = "paired_glob" @@ -41,7 +42,7 @@ class Strategy(StrEnum): class BaseSchemaLoader(abc.ABC): """ - Interface that every task-specific loader must implement. + Interface that every strategy loader must implement. Args: schema (DatasetSchema): The parsed schema for the dataset. @@ -69,8 +70,8 @@ def _load_index_file(self) -> pd.DataFrame: ``schema.format`` via `FORMAT_SEP`, then delegates the file lookup to `_resolve_index_file`. - Used by all index-based loaders (ASR, TTS, ...) so that each loader - only needs to call `_apply_column_mappings` on the result. + Used by index-based strategies so that each loader only needs to call + `_apply_column_mappings` on the result. Returns: A raw (unmapped) DataFrame exactly as read from the index file. @@ -106,43 +107,6 @@ def _resolve_index_file(self) -> Path: ) return self._resolved_index_file - def _load_multi_sections(self) -> pd.DataFrame: - """ - Parsing logic for archives with multiple directories, and each directory - has its own index file. The section name is inferred from the parent directory of the index file. - """ - sections = self._resolve_sections() - parts: list[pd.DataFrame] = [] - for section_path in sections: - section_df = self._read_delimited_file(section_path) - section_df["section"] = section_path.parents[0].name - parts.append(section_df) - - return pd.concat(parts, ignore_index=True) - - def _resolve_sections(self) -> list: - """ - Get a list of valid sections, i.e. subdirectories that include an index file. - """ - - assert self.schema.sections is not None - assert self.schema.index_file is not None - assert self.schema.section_root is not None - sections = self.schema.sections - section_paths = [] - for section in sections: - section_path = ( - self.extract_dir - / Path(self.schema.section_root) - / Path(section) - / self.schema.index_file - ) - if not section_path.exists(): - raise FileNotFoundError(f"Index file '{section_path}' not found ") - section_paths.append(section_path) - - return section_paths - def _apply_column_mappings(self, raw_df: pd.DataFrame) -> pd.DataFrame: """Select and rename columns according to the schema, applying dtype conversions. @@ -414,12 +378,12 @@ def _build_direct_file_candidates( if relative_candidate.is_absolute(): path_candidates = [relative_candidate] else: - path_candidates = list( + path_candidates = [ root / relative_candidate for root in self._get_audio_search_roots( row=row, template_value=template_value or raw_value ) - ) + ] dataset_root = self._get_dataset_root() path_candidates.append(dataset_root / relative_candidate) if dataset_root != self.extract_dir: diff --git a/src/datacollective/schema_loaders/contracts.py b/src/datacollective/schema_loaders/contracts.py new file mode 100644 index 0000000..3850e52 --- /dev/null +++ b/src/datacollective/schema_loaders/contracts.py @@ -0,0 +1,82 @@ +from __future__ import annotations + +import pandas as pd + +from datacollective.errors import TaskValidationError +from datacollective.logging_utils import get_logger +from datacollective.schema import DatasetSchema +from datacollective.schema_loaders.base import Strategy + +logger = get_logger(__name__) + +#: Logical columns a loaded DataFrame must contain for each known task. +TASK_CONTRACTS: dict[str, frozenset[str]] = { + "ASR": frozenset({"audio_path", "transcription"}), + "TTS": frozenset({"audio_path", "transcription"}), + "LLM": frozenset({"text"}), +} + + +def _validate_task_contract(df: pd.DataFrame, task: str | None) -> None: + """Check that the loaded DataFrame satisfies the task's column contract. + + Tasks without a contract (e.g. ``OTH``) and schemas without a task are + accepted as-is. + + Raises: + TaskValidationError: If contract columns are missing from *df*. + """ + if not task: + logger.debug("Schema has no task — skipping task contract validation.") + return + contract = TASK_CONTRACTS.get(task.upper()) + if contract is None: + logger.debug( + f"No contract defined for task '{task}' — skipping task contract validation." + ) + return + + missing = sorted(column for column in contract if column not in df.columns) + if missing: + raise TaskValidationError( + f"Loaded dataset does not satisfy the '{task.upper()}' task contract: " + f"missing column(s) {missing}. " + f"Available columns: {list(df.columns)}" + ) + + +def _validate_declared_contract(schema: DatasetSchema, strategy: Strategy) -> None: + """Fail fast when the declared column mappings cannot satisfy the task contract. + + Only applies when the task has a contract, the schema declares column + mappings, and the strategy actually applies those mappings — catching + misconfigured schemas before expensive per-row file resolution. + + Raises: + TaskValidationError: If contract columns are absent from the declared + logical column names. + """ + if not schema.task or not schema.columns: + return + if not _strategy_applies_mappings(schema, strategy): + return + contract = TASK_CONTRACTS.get(schema.task.upper()) + if contract is None: + return + + missing = sorted(column for column in contract if column not in schema.columns) + if missing: + raise TaskValidationError( + f"Schema for task '{schema.task.upper()}' declares column mappings that " + f"cannot satisfy the task contract: missing logical column(s) {missing}. " + f"Declared columns: {sorted(schema.columns)}" + ) + + +def _strategy_applies_mappings(schema: DatasetSchema, strategy: Strategy) -> bool: + if strategy == Strategy.GLOB: + return False + if strategy == Strategy.PAIRED_GLOB: + # Only the JSON-sidecar variant applies column mappings + return (schema.format or "").casefold() == "json" + return True diff --git a/src/datacollective/schema_loaders/registry.py b/src/datacollective/schema_loaders/registry.py index 8c8358b..4111b20 100644 --- a/src/datacollective/schema_loaders/registry.py +++ b/src/datacollective/schema_loaders/registry.py @@ -9,43 +9,75 @@ from datacollective.logging_utils import get_logger from datacollective.schema import DatasetSchema -from datacollective.schema_loaders.base import BaseSchemaLoader -from datacollective.schema_loaders.tasks.asr import ASRLoader -from datacollective.schema_loaders.tasks.oth import OTHLoader -from datacollective.schema_loaders.tasks.tts import TTSLoader +from datacollective.schema_loaders.base import BaseSchemaLoader, Strategy +from datacollective.schema_loaders.contracts import ( + _validate_declared_contract, + _validate_task_contract, +) +from datacollective.schema_loaders.strategies import ( + GlobLoader, + IndexLoader, + MultiSectionsLoader, + MultiSplitLoader, + PairedGlobLoader, +) logger = get_logger(__name__) -_TASK_REGISTRY: dict[str, Type[BaseSchemaLoader]] = { - "ASR": ASRLoader, - "OTH": OTHLoader, - "TTS": TTSLoader, +_STRATEGY_REGISTRY: dict[Strategy, Type[BaseSchemaLoader]] = { + Strategy.INDEX: IndexLoader, + Strategy.MULTI_SPLIT: MultiSplitLoader, + Strategy.MULTI_SECTIONS: MultiSectionsLoader, + Strategy.PAIRED_GLOB: PairedGlobLoader, + Strategy.GLOB: GlobLoader, } -def _get_task_loader(task: str) -> Type[BaseSchemaLoader]: +def _resolve_strategy(schema: DatasetSchema) -> Strategy: """ - Return the loader class for *task*. + Resolve the loading strategy for *schema*. + + Defaults to `Strategy.INDEX` when ``root_strategy`` is not set. Raises: - ValueError: If no loader is registered for the given task. + ValueError: If ``root_strategy`` names an unknown strategy. """ - key = task.upper() - if key not in _TASK_REGISTRY: - supported = ", ".join(sorted(_TASK_REGISTRY)) + raw = schema.root_strategy or Strategy.INDEX + try: + return Strategy(raw) + except ValueError: + supported = ", ".join(member.value for member in Strategy) raise ValueError( - f"No schema loader registered for task '{key}'. " - f"Supported tasks: {supported}" + f"Unknown root_strategy '{raw}'. Supported strategies: {supported}" + ) from None + + +def _get_strategy_loader(strategy: Strategy) -> Type[BaseSchemaLoader]: + """ + Return the loader class for *strategy*. + + Raises: + ValueError: If no loader is registered for the given strategy. + """ + if strategy not in _STRATEGY_REGISTRY: + supported = ", ".join(sorted(_STRATEGY_REGISTRY)) + raise ValueError( + f"No schema loader registered for strategy '{strategy}'. " + f"Supported strategies: {supported}" ) - return _TASK_REGISTRY[key] + return _STRATEGY_REGISTRY[strategy] def _load_dataset_from_schema(schema: DatasetSchema, extract_dir: Path) -> pd.DataFrame: """ - Instantiate the appropriate loader for *schema.task* and return the + Instantiate the loader for the schema's ``root_strategy`` and return the loaded `~pandas.DataFrame`. + When the schema declares a task with a known contract (see + `~datacollective.schema_loaders.contracts.TASK_CONTRACTS`), the loaded + DataFrame is validated against it. + Args: schema: Parsed dataset schema. extract_dir: Root directory where the dataset archive was extracted. @@ -56,10 +88,16 @@ def _load_dataset_from_schema(schema: DatasetSchema, extract_dir: Path) -> pd.Da if schema.extract_files: _extract_inner_archives(schema.extract_files, extract_dir) - loader_cls = _get_task_loader(schema.task) + strategy = _resolve_strategy(schema) + loader_cls = _get_strategy_loader(strategy) + _validate_declared_contract(schema, strategy) + loader = loader_cls(schema=schema, extract_dir=extract_dir) logger.info(f"Loading dataset '{schema.dataset_id}' with {loader_cls.__name__}") - return loader.load() + df = loader.load() + + _validate_task_contract(df, schema.task) + return df def _extract_inner_archives(extract_files: list[str], extract_dir: Path) -> None: diff --git a/src/datacollective/schema_loaders/strategies/__init__.py b/src/datacollective/schema_loaders/strategies/__init__.py new file mode 100644 index 0000000..9f75b0e --- /dev/null +++ b/src/datacollective/schema_loaders/strategies/__init__.py @@ -0,0 +1,13 @@ +from datacollective.schema_loaders.strategies.glob import GlobLoader +from datacollective.schema_loaders.strategies.index import IndexLoader +from datacollective.schema_loaders.strategies.multi_sections import MultiSectionsLoader +from datacollective.schema_loaders.strategies.multi_split import MultiSplitLoader +from datacollective.schema_loaders.strategies.paired_glob import PairedGlobLoader + +__all__ = [ + "GlobLoader", + "IndexLoader", + "MultiSectionsLoader", + "MultiSplitLoader", + "PairedGlobLoader", +] diff --git a/src/datacollective/schema_loaders/tasks/oth.py b/src/datacollective/schema_loaders/strategies/glob.py similarity index 65% rename from src/datacollective/schema_loaders/tasks/oth.py rename to src/datacollective/schema_loaders/strategies/glob.py index f58d46e..1e351dc 100644 --- a/src/datacollective/schema_loaders/tasks/oth.py +++ b/src/datacollective/schema_loaders/strategies/glob.py @@ -1,49 +1,29 @@ +from __future__ import annotations + from pathlib import Path import pandas as pd from datacollective.logging_utils import get_logger from datacollective.schema import DatasetSchema -from datacollective.schema_loaders.base import BaseSchemaLoader, Strategy +from datacollective.schema_loaders.base import BaseSchemaLoader logger = get_logger(__name__) -class OTHLoader(BaseSchemaLoader): - """Loader for tasks classified as OTH. - - Supports two strategies: +class GlobLoader(BaseSchemaLoader): + """Load a directory-structured dataset by globbing for files. - 1) Glob (``root_strategy: glob``): for directory-structured datasets where - metadata (e.g. speaker ID, language) is encoded in the path hierarchy rather - than in an index file or text-file pairing. - - 2) Index-file: reads a single delimited index file and applies the schema's - column mappings. Example datasets: Mozilla Common Voice Text Language - Identification dataset. + Metadata (e.g. speaker ID, language) is derived from the path hierarchy + rather than from an index file or sidecar pairing. """ def __init__(self, schema: DatasetSchema, extract_dir: Path) -> None: super().__init__(schema, extract_dir) - - if schema.root_strategy == Strategy.GLOB: - if not schema.file_pattern: - raise ValueError("OTH glob schema must specify 'file_pattern'") - elif not schema.index_file: - raise ValueError( - "OTH schema must specify either 'root_strategy: glob' with " - "'file_pattern', or an 'index_file'" - ) + if not schema.file_pattern: + raise ValueError("glob schema must specify 'file_pattern'") def load(self) -> pd.DataFrame: - if self.schema.root_strategy == Strategy.GLOB: - return self._load_glob() - raw_df = self._load_index_file() - if not self.schema.columns: - return raw_df - return self._apply_column_mappings(raw_df) - - def _load_glob(self) -> pd.DataFrame: """Glob for files and derive metadata from the directory hierarchy. When ``splits`` is set, each split name is treated as a subdirectory @@ -56,8 +36,6 @@ def _load_glob(self) -> pd.DataFrame: - ``language``: parent directory name - ``split`` (when splits are configured): source split directory """ - assert self.schema.file_pattern is not None - if self.schema.splits: return self._load_glob_splits() diff --git a/src/datacollective/schema_loaders/strategies/index.py b/src/datacollective/schema_loaders/strategies/index.py new file mode 100644 index 0000000..63bc815 --- /dev/null +++ b/src/datacollective/schema_loaders/strategies/index.py @@ -0,0 +1,32 @@ +from __future__ import annotations + +from pathlib import Path + +import pandas as pd + +from datacollective.logging_utils import get_logger +from datacollective.schema import DatasetSchema +from datacollective.schema_loaders.base import BaseSchemaLoader + +logger = get_logger(__name__) + + +class IndexLoader(BaseSchemaLoader): + """Load a dataset from a single delimited index file (the default strategy). + + An index file (e.g. CSV/TSV) holds one row per sample. When the schema + declares column mappings they are applied (renaming, dtype conversion, + file-path resolution); otherwise the raw DataFrame is returned as-is. + """ + + def __init__(self, schema: DatasetSchema, extract_dir: Path) -> None: + super().__init__(schema, extract_dir) + if not schema.index_file: + raise ValueError("index strategy schema must specify 'index_file'") + + def load(self) -> pd.DataFrame: + raw_df = self._load_index_file() + if not self.schema.columns: + # No column mapping -> return the raw dataframe as-is + return raw_df + return self._apply_column_mappings(raw_df) diff --git a/src/datacollective/schema_loaders/strategies/multi_sections.py b/src/datacollective/schema_loaders/strategies/multi_sections.py new file mode 100644 index 0000000..5ca0a3f --- /dev/null +++ b/src/datacollective/schema_loaders/strategies/multi_sections.py @@ -0,0 +1,67 @@ +from __future__ import annotations + +from pathlib import Path + +import pandas as pd + +from datacollective.logging_utils import get_logger +from datacollective.schema import DatasetSchema +from datacollective.schema_loaders.base import BaseSchemaLoader + +logger = get_logger(__name__) + + +class MultiSectionsLoader(BaseSchemaLoader): + """Load a dataset organised as one index file per section directory. + + Each section directory under ``section_root`` holds its own index file. + A ``section`` column (the directory name) is added to each part, column + mappings are applied when declared, and the parts are concatenated. + """ + + def __init__(self, schema: DatasetSchema, extract_dir: Path) -> None: + super().__init__(schema, extract_dir) + if not schema.sections: + raise ValueError( + "multi_sections schema must specify 'sections' (list of section names)" + ) + if not schema.section_root: + raise ValueError("multi_sections schema must specify 'section_root'") + if not schema.index_file: + raise ValueError("multi_sections schema must specify 'index_file'") + + def load(self) -> pd.DataFrame: + sections = self._resolve_sections() + parts: list[pd.DataFrame] = [] + for section_path in sections: + section_df = self._read_delimited_file(section_path) + section_name = section_path.parents[0].name + + if self.schema.columns: + section_df = self._apply_column_mappings(section_df) + section_df["section"] = section_name + parts.append(section_df) + + return pd.concat(parts, ignore_index=True) + + def _resolve_sections(self) -> list[Path]: + """ + Get a list of valid sections, i.e. subdirectories that include an index file. + """ + assert self.schema.sections is not None + assert self.schema.index_file is not None + assert self.schema.section_root is not None + + section_paths = [] + for section in self.schema.sections: + section_path = ( + self.extract_dir + / Path(self.schema.section_root) + / Path(section) + / self.schema.index_file + ) + if not section_path.exists(): + raise FileNotFoundError(f"Index file '{section_path}' not found ") + section_paths.append(section_path) + + return section_paths diff --git a/src/datacollective/schema_loaders/strategies/multi_split.py b/src/datacollective/schema_loaders/strategies/multi_split.py new file mode 100644 index 0000000..6e6191a --- /dev/null +++ b/src/datacollective/schema_loaders/strategies/multi_split.py @@ -0,0 +1,64 @@ +from __future__ import annotations + +from pathlib import Path + +import pandas as pd + +from datacollective.logging_utils import get_logger +from datacollective.schema import DatasetSchema +from datacollective.schema_loaders.base import BaseSchemaLoader + +logger = get_logger(__name__) + + +class MultiSplitLoader(BaseSchemaLoader): + """Load a dataset spread across one delimited file per split. + + All split files whose stems match the ``splits`` list are read, a + ``split`` column is added to each, column mappings are applied when + declared, and the parts are concatenated. + """ + + def __init__(self, schema: DatasetSchema, extract_dir: Path) -> None: + super().__init__(schema, extract_dir) + if not schema.splits: + raise ValueError( + "multi_split schema must specify 'splits' (list of split names)" + ) + + def load(self) -> pd.DataFrame: + assert self.schema.splits is not None + + pattern = self.schema.splits_file_pattern or "**/*.tsv" + allowed_splits = set(self.schema.splits) + + split_files: dict[str, Path] = {} + for path in self.extract_dir.rglob(pattern): + if path.stem in allowed_splits: + # Prefer the shallowest match per split name + if path.stem not in split_files or len(path.parts) < len( + split_files[path.stem].parts + ): + split_files[path.stem] = path + + if not split_files: + raise RuntimeError( + f"No split files matching pattern '{pattern}' with stems in " + f"{sorted(allowed_splits)} found under '{self.extract_dir}'" + ) + + frames: list[pd.DataFrame] = [] + + for split_name, file_path in sorted(split_files.items()): + logger.debug(f"Reading split '{split_name}' from {file_path}") + raw_df = self._read_delimited_file(file_path) + raw_df["split"] = split_name + + if self.schema.columns: + mapped = self._apply_column_mappings(raw_df) + mapped["split"] = split_name + frames.append(mapped) + else: + frames.append(raw_df) + + return pd.concat(frames, ignore_index=True) diff --git a/src/datacollective/schema_loaders/strategies/paired_glob.py b/src/datacollective/schema_loaders/strategies/paired_glob.py new file mode 100644 index 0000000..7e585e9 --- /dev/null +++ b/src/datacollective/schema_loaders/strategies/paired_glob.py @@ -0,0 +1,154 @@ +from __future__ import annotations + +import json +from pathlib import Path + +import pandas as pd + +from datacollective.logging_utils import get_logger +from datacollective.schema import DatasetSchema +from datacollective.schema_loaders.base import BaseSchemaLoader + +logger = get_logger(__name__) + + +class PairedGlobLoader(BaseSchemaLoader): + """Load a dataset where each audio file is paired with a sidecar file. + + Two variants exist, selected by ``schema.format``: + + - ``format: "json"``: each audio file has a JSON sidecar (matched via + ``file_pattern``); column mappings are required and are applied to the + normalised JSON records. + - otherwise: each audio file has a matching text sidecar containing the + transcription; requires ``file_pattern`` and ``audio_extension``. + """ + + def __init__(self, schema: DatasetSchema, extract_dir: Path) -> None: + super().__init__(schema, extract_dir) + if not schema.file_pattern: + raise ValueError("paired_glob schema must specify 'file_pattern'") + if self._is_json_variant(): + if not schema.columns: + raise ValueError( + "paired_glob schema with 'format: json' must specify column " + "mappings (e.g. for audio and transcription)" + ) + elif not schema.audio_extension: + raise ValueError( + "paired_glob schema must specify 'audio_extension' " + "(or 'format: json' for JSON sidecar files)" + ) + + def _is_json_variant(self) -> bool: + return (self.schema.format or "").casefold() == "json" + + def load(self) -> pd.DataFrame: + if self._is_json_variant(): + return self._load_json_sidecars() + return self._load_text_sidecars() + + def _load_json_sidecars(self) -> pd.DataFrame: + """ + Load a dataset where each audio file is paired with a JSON sidecar + (matched via ``file_pattern``) instead of a central index file. + + When ``record_path`` is set, the JSON key it names must hold a list of + records (e.g. time-aligned utterances) and each record becomes one row; + the remaining top-level keys are flattened with dot notation + (audio.filename, ...) and repeated on + every row of that file. Without ``record_path`` each JSON file yields + a single row. Column mappings are then applied as for index files, so + ``file_path`` columns (typically sourced from a filename field inside + the JSON) resolve through the usual audio-path machinery. + """ + assert self.schema.file_pattern is not None + + json_files = sorted(self.extract_dir.rglob(self.schema.file_pattern)) + json_files = [p for p in json_files if not p.name.startswith("._")] + if not json_files: + raise FileNotFoundError( + f"No files matching '{self.schema.file_pattern}' " + f"found under '{self.extract_dir}'" + ) + + logger.debug( + f"Found {len(json_files)} JSON files matching '{self.schema.file_pattern}'" + ) + + record_path = self.schema.record_path + frames: list[pd.DataFrame] = [] + for path in json_files: + data = json.loads(path.read_text(encoding=self.schema.encoding)) + if record_path: + if record_path not in data: + raise KeyError( + f"record_path '{record_path}' not found in '{path}'. " + f"Available keys: {list(data)}" + ) + frame = pd.json_normalize(data, record_path=record_path) + meta = pd.json_normalize( + {key: value for key, value in data.items() if key != record_path} + ) + for column in meta.columns: + frame[column] = meta[column].iloc[0] + else: + frame = pd.json_normalize(data) + frames.append(frame) + + raw_df = pd.concat(frames, ignore_index=True) + return self._apply_column_mappings(raw_df) + + def _load_text_sidecars(self) -> pd.DataFrame: + """ + Load a dataset where each audio file has a matching text file (e.g. + ``.txt``) containing the transcription. The loader searches recursively + for all text files matching the specified `file_pattern`, reads their + contents, and pairs them with the corresponding audio files based on + the same filename stem. The parent directory name of each text/audio + pair is captured as a `split` column in the resulting DataFrame. + """ + assert self.schema.file_pattern is not None + assert self.schema.audio_extension is not None + + text_files = sorted(self.extract_dir.rglob(self.schema.file_pattern)) + if not text_files: + raise FileNotFoundError( + f"No files matching '{self.schema.file_pattern}' " + f"found under '{self.extract_dir}'" + ) + + logger.debug( + f"Found {len(text_files)} text files matching '{self.schema.file_pattern}'" + ) + + audio_ext = self.schema.audio_extension + rows: list[dict[str, str]] = [] + + for txt_path in text_files: + audio_path = txt_path.with_suffix(audio_ext) + if not audio_path.exists(): + logger.debug( + f"No matching audio file for '{txt_path.name}' — skipping." + ) + continue + + transcription = txt_path.read_text(encoding=self.schema.encoding).strip() + row: dict[str, str] = { + "audio_path": str(audio_path), + "transcription": transcription, + } + + # Derive domain / split from parent directory name if present + parent_name = txt_path.parent.name + if parent_name: + row["split"] = parent_name + + rows.append(row) + + if not rows: + raise FileNotFoundError( + f"No paired (text + {audio_ext}) files found under '{self.extract_dir}'" + ) + + return pd.DataFrame(rows) diff --git a/src/datacollective/schema_loaders/tasks/asr.py b/src/datacollective/schema_loaders/tasks/asr.py deleted file mode 100644 index 642977a..0000000 --- a/src/datacollective/schema_loaders/tasks/asr.py +++ /dev/null @@ -1,141 +0,0 @@ -from __future__ import annotations - -import json -from pathlib import Path - -import pandas as pd - -from datacollective.logging_utils import get_logger -from datacollective.schema import DatasetSchema -from datacollective.schema_loaders.base import BaseSchemaLoader, Strategy - -logger = get_logger(__name__) - - -class ASRLoader(BaseSchemaLoader): - """Load an ASR dataset described by a `DatasetSchema`.""" - - def __init__(self, schema: DatasetSchema, extract_dir: Path) -> None: - super().__init__(schema, extract_dir) - if schema.root_strategy == Strategy.MULTI_SPLIT: - if not schema.splits: - raise ValueError( - "ASR multi_split schema must specify 'splits' (list of split names)" - ) - elif schema.root_strategy == Strategy.PAIRED_GLOB: - if (schema.format or "").casefold() != "json": - raise ValueError("ASR paired_glob schema only supports 'format: json'") - if not schema.file_pattern: - raise ValueError("ASR paired_glob schema must specify 'file_pattern'") - if not schema.columns: - raise ValueError( - "ASR paired_glob schema must specify at least two column mappings " - "for audio and transcription" - ) - else: - if not schema.index_file: - raise ValueError("ASR schema must specify 'index_file'") - if not schema.columns: - raise ValueError( - "ASR schema must specify at least two column mappings for audio and transcription" - ) - - def load(self) -> pd.DataFrame: - if self.schema.root_strategy == Strategy.MULTI_SPLIT: - return self._load_multi_split() - if self.schema.root_strategy == Strategy.PAIRED_GLOB: - return self._load_paired_glob_json() - raw_df = self._load_index_file() - return self._apply_column_mappings(raw_df) - - def _load_paired_glob_json(self) -> pd.DataFrame: - """ - Load an ASR dataset where each audio file is paired with a JSON sidecar - (matched via ``file_pattern``) instead of a central index file. - - When ``record_path`` is set, the JSON key it names must hold a list of - records (e.g. time-aligned utterances) and each record becomes one row; - the remaining top-level keys are flattened with dot notation - (audio.filename, ...) and repeated on - every row of that file. Without ``record_path`` each JSON file yields - a single row. Column mappings are then applied as for index files, so - ``file_path`` columns (typically sourced from a filename field inside - the JSON) resolve through the usual audio-path machinery. - """ - assert self.schema.file_pattern is not None - - json_files = sorted(self.extract_dir.rglob(self.schema.file_pattern)) - json_files = [p for p in json_files if not p.name.startswith("._")] - if not json_files: - raise FileNotFoundError( - f"No files matching '{self.schema.file_pattern}' " - f"found under '{self.extract_dir}'" - ) - - logger.debug( - f"Found {len(json_files)} JSON files matching '{self.schema.file_pattern}'" - ) - - record_path = self.schema.record_path - frames: list[pd.DataFrame] = [] - for path in json_files: - data = json.loads(path.read_text(encoding=self.schema.encoding)) - if record_path: - if record_path not in data: - raise KeyError( - f"record_path '{record_path}' not found in '{path}'. " - f"Available keys: {list(data)}" - ) - frame = pd.json_normalize(data, record_path=record_path) - meta = pd.json_normalize( - {key: value for key, value in data.items() if key != record_path} - ) - for column in meta.columns: - frame[column] = meta[column].iloc[0] - else: - frame = pd.json_normalize(data) - frames.append(frame) - - raw_df = pd.concat(frames, ignore_index=True) - return self._apply_column_mappings(raw_df) - - def _load_multi_split(self) -> pd.DataFrame: - """ - Load all split TSV/CSV files whose stems match the ``splits`` list, - add a ``split`` column to each, apply column mappings, and concatenate. - """ - assert self.schema.splits is not None - - pattern = self.schema.splits_file_pattern or "**/*.tsv" - allowed_splits = set(self.schema.splits) - - split_files: dict[str, Path] = {} - for path in self.extract_dir.rglob(pattern): - if path.stem in allowed_splits: - # Prefer the shallowest match per split name - if path.stem not in split_files or len(path.parts) < len( - split_files[path.stem].parts - ): - split_files[path.stem] = path - - if not split_files: - raise RuntimeError( - f"No split files matching pattern '{pattern}' with stems in " - f"{sorted(allowed_splits)} found under '{self.extract_dir}'" - ) - - frames: list[pd.DataFrame] = [] - - for split_name, file_path in sorted(split_files.items()): - logger.debug(f"Reading split '{split_name}' from {file_path}") - raw_df = self._read_delimited_file(file_path) - raw_df["split"] = split_name - - if self.schema.columns: - mapped = self._apply_column_mappings(raw_df) - mapped["split"] = split_name - frames.append(mapped) - else: - frames.append(raw_df) - - return pd.concat(frames, ignore_index=True) diff --git a/src/datacollective/schema_loaders/tasks/tts.py b/src/datacollective/schema_loaders/tasks/tts.py deleted file mode 100644 index 84d3bad..0000000 --- a/src/datacollective/schema_loaders/tasks/tts.py +++ /dev/null @@ -1,99 +0,0 @@ -from __future__ import annotations - -from pathlib import Path - -import pandas as pd - -from datacollective.logging_utils import get_logger -from datacollective.schema import DatasetSchema -from datacollective.schema_loaders.base import BaseSchemaLoader, Strategy - -logger = get_logger(__name__) - - -class TTSLoader(BaseSchemaLoader): - """Load a TTS dataset described by a `DatasetSchema`. - - See docs/loaders/tts.md for details on supported loading strategies and schema fields. - """ - - def __init__(self, schema: DatasetSchema, extract_dir: Path) -> None: - super().__init__(schema, extract_dir) - - def load(self) -> pd.DataFrame: - if self.schema.root_strategy == Strategy.PAIRED_GLOB: - return self._load_paired_glob() - elif self.schema.root_strategy == Strategy.MULTI_SECTIONS: - return self._load_multi_sections() - return self._load_based_on_index() - - def _load_based_on_index(self) -> pd.DataFrame: - """ - Load a TTS dataset using the "index" strategy, where an index file (e.g. CSV) maps audio paths to transcriptions. - """ - if not self.schema.index_file: - raise ValueError("TTS index-based schema must specify 'index_file'") - - raw_df = self._load_index_file() - - if not self.schema.columns: - # No column mapping -> return the raw dataframe as-is - return raw_df - - return self._apply_column_mappings(raw_df) - - def _load_paired_glob(self) -> pd.DataFrame: - """ - Load a TTS dataset using the "paired_glob" strategy, where each audio file has a - matching `.txt` file containing the transcription. The loader searches - recursively for all text files matching the specified `file_pattern`, - reads their contents, and pairs them with the corresponding audio files based - on the same filename stem. The parent directory name of each text/audio pair - is captured as a `split` column in the resulting DataFrame. - """ - if not self.schema.file_pattern: - raise ValueError("TTS paired_glob schema must specify 'file_pattern'") - if not self.schema.audio_extension: - raise ValueError("TTS paired_glob schema must specify 'audio_extension'") - - text_files = sorted(self.extract_dir.rglob(self.schema.file_pattern)) - if not text_files: - raise FileNotFoundError( - f"No files matching '{self.schema.file_pattern}' " - f"found under '{self.extract_dir}'" - ) - - logger.debug( - f"Found {len(text_files)} text files matching '{self.schema.file_pattern}'" - ) - - audio_ext = self.schema.audio_extension - rows: list[dict[str, str]] = [] - - for txt_path in text_files: - audio_path = txt_path.with_suffix(audio_ext) - if not audio_path.exists(): - logger.debug( - f"No matching audio file for '{txt_path.name}' — skipping." - ) - continue - - transcription = txt_path.read_text(encoding=self.schema.encoding).strip() - row: dict[str, str] = { - "audio_path": str(audio_path), - "transcription": transcription, - } - - # Derive domain / split from parent directory name if present - parent_name = txt_path.parent.name - if parent_name: - row["split"] = parent_name - - rows.append(row) - - if not rows: - raise FileNotFoundError( - f"No paired (text + {audio_ext}) files found under '{self.extract_dir}'" - ) - - return pd.DataFrame(rows) diff --git a/src/datacollective/schema_loaders/tasks/__init__.py b/tests/schema_loaders/strategies/__init__.py similarity index 100% rename from src/datacollective/schema_loaders/tasks/__init__.py rename to tests/schema_loaders/strategies/__init__.py diff --git a/tests/schema_loaders/strategies/test_glob_loader.py b/tests/schema_loaders/strategies/test_glob_loader.py new file mode 100644 index 0000000..19c335e --- /dev/null +++ b/tests/schema_loaders/strategies/test_glob_loader.py @@ -0,0 +1,72 @@ +from __future__ import annotations + +from pathlib import Path + +import pytest + +from datacollective.schema import DatasetSchema +from datacollective.schema_loaders.strategies.glob import GlobLoader + + +def _touch(path: Path) -> None: + path.parent.mkdir(parents=True, exist_ok=True) + path.write_bytes(b"\x00") + + +class TestGlobValidation: + def test_requires_file_pattern(self, tmp_path: Path) -> None: + schema = DatasetSchema(dataset_id="ds", root_strategy="glob") + with pytest.raises(ValueError, match="file_pattern"): + GlobLoader(schema, tmp_path) + + +class TestGlobLoader: + def test_derives_metadata_from_path(self, tmp_path: Path) -> None: + _touch(tmp_path / "spk1" / "en" / "a.wav") + _touch(tmp_path / "spk2" / "fr" / "b.wav") + + schema = DatasetSchema( + dataset_id="ds", + root_strategy="glob", + file_pattern="**/*.wav", + ) + df = GlobLoader(schema, tmp_path).load() + assert len(df) == 2 + assert list(df.columns) == ["audio_path", "language", "speaker_id"] + assert set(df["language"]) == {"en", "fr"} + assert set(df["speaker_id"]) == {"spk1", "spk2"} + + def test_splits_add_split_column(self, tmp_path: Path) -> None: + _touch(tmp_path / "train" / "spk1" / "en" / "a.wav") + _touch(tmp_path / "dev" / "spk2" / "fr" / "b.wav") + + schema = DatasetSchema( + dataset_id="ds", + root_strategy="glob", + file_pattern="**/*.wav", + splits=["train", "dev"], + ) + df = GlobLoader(schema, tmp_path).load() + assert len(df) == 2 + assert set(df["split"]) == {"train", "dev"} + + def test_missing_split_directory_raises(self, tmp_path: Path) -> None: + _touch(tmp_path / "train" / "spk1" / "en" / "a.wav") + + schema = DatasetSchema( + dataset_id="ds", + root_strategy="glob", + file_pattern="**/*.wav", + splits=["train", "dev"], + ) + with pytest.raises(FileNotFoundError, match="dev"): + GlobLoader(schema, tmp_path).load() + + def test_no_matching_files_raises(self, tmp_path: Path) -> None: + schema = DatasetSchema( + dataset_id="ds", + root_strategy="glob", + file_pattern="**/*.wav", + ) + with pytest.raises(FileNotFoundError, match="No files matching"): + GlobLoader(schema, tmp_path).load() diff --git a/tests/schema_loaders/tasks/test_asr_loader.py b/tests/schema_loaders/strategies/test_index_loader.py similarity index 58% rename from tests/schema_loaders/tasks/test_asr_loader.py rename to tests/schema_loaders/strategies/test_index_loader.py index efe350a..e4a4874 100644 --- a/tests/schema_loaders/tasks/test_asr_loader.py +++ b/tests/schema_loaders/strategies/test_index_loader.py @@ -5,48 +5,34 @@ import pytest from datacollective.schema import ColumnMapping, DatasetSchema -from datacollective.schema_loaders.tasks.asr import ASRLoader +from datacollective.schema_loaders.strategies.index import IndexLoader -def _write_tsv(path: Path, content: str) -> None: +def _write(path: Path, content: str) -> None: path.parent.mkdir(parents=True, exist_ok=True) path.write_text(content, encoding="utf-8") -class TestASRLoaderValidation: - def test_index_requires_index_file(self, tmp_path: Path) -> None: +class TestIndexLoaderValidation: + def test_requires_index_file(self, tmp_path: Path) -> None: schema = DatasetSchema( dataset_id="ds", - task="ASR", format="tsv", columns={"a": ColumnMapping(source_column="x")}, ) with pytest.raises(ValueError, match="index_file"): - ASRLoader(schema, tmp_path) + IndexLoader(schema, tmp_path) - def test_index_requires_columns(self, tmp_path: Path) -> None: - schema = DatasetSchema( - dataset_id="ds", task="ASR", format="tsv", index_file="f.tsv" - ) - with pytest.raises(ValueError, match="column mapping"): - ASRLoader(schema, tmp_path) - - def test_multi_split_requires_splits(self, tmp_path: Path) -> None: - schema = DatasetSchema(dataset_id="ds", task="ASR", root_strategy="multi_split") - with pytest.raises(ValueError, match="splits"): - ASRLoader(schema, tmp_path) - -class TestASRIndexBased: +class TestIndexLoader: def test_load_tsv_without_format(self, tmp_path: Path) -> None: - _write_tsv( + _write( tmp_path / "train.tsv", "path\tsentence\nclip1.mp3\thello\nclip2.mp3\tworld\n", ) schema = DatasetSchema( dataset_id="ds", - task="ASR", index_file="train.tsv", columns={ "audio_path": ColumnMapping(source_column="path", dtype="file_path"), @@ -55,19 +41,18 @@ def test_load_tsv_without_format(self, tmp_path: Path) -> None: ), }, ) - df = ASRLoader(schema, tmp_path).load() + df = IndexLoader(schema, tmp_path).load() assert len(df) == 2 assert list(df.columns) == ["audio_path", "transcription"] def test_load_tsv(self, tmp_path: Path) -> None: - _write_tsv( + _write( tmp_path / "train.tsv", "path\tsentence\nclip1.mp3\thello\nclip2.mp3\tworld\n", ) schema = DatasetSchema( dataset_id="ds", - task="ASR", format="tsv", index_file="train.tsv", columns={ @@ -77,17 +62,16 @@ def test_load_tsv(self, tmp_path: Path) -> None: ), }, ) - df = ASRLoader(schema, tmp_path).load() + df = IndexLoader(schema, tmp_path).load() assert len(df) == 2 assert list(df.columns) == ["audio_path", "transcription"] assert df["transcription"].iloc[0] == "hello" def test_load_csv(self, tmp_path: Path) -> None: - _write_tsv(tmp_path / "data.csv", "path,sentence\nc1.mp3,hi\n") + _write(tmp_path / "data.csv", "path,sentence\nc1.mp3,hi\n") schema = DatasetSchema( dataset_id="ds", - task="ASR", format="csv", index_file="data.csv", columns={ @@ -95,22 +79,59 @@ def test_load_csv(self, tmp_path: Path) -> None: "text": ColumnMapping(source_column="sentence"), }, ) - df = ASRLoader(schema, tmp_path).load() + df = IndexLoader(schema, tmp_path).load() assert len(df) == 1 assert "audio" in df.columns + def test_load_pipe_delimited_headerless(self, tmp_path: Path) -> None: + _write(tmp_path / "meta.csv", "clip1.mp3|hello world\nclip2.mp3|goodbye\n") + + schema = DatasetSchema( + dataset_id="ds", + format="pipe", + separator="|", + has_header=False, + index_file="meta.csv", + base_audio_path="wavs/", + columns={ + "audio_path": ColumnMapping(source_column=0, dtype="file_path"), + "transcription": ColumnMapping(source_column=1, dtype="string"), + }, + ) + df = IndexLoader(schema, tmp_path).load() + assert len(df) == 2 + assert df["transcription"].iloc[0] == "hello world" + # file_path dtype -> absolute path with base_audio_path + assert "wavs" in df["audio_path"].iloc[0] + + def test_no_columns_returns_raw(self, tmp_path: Path) -> None: + _write(tmp_path / "meta.csv", "a,b\n1,2\n") + + schema = DatasetSchema( + dataset_id="ds", + format="csv", + index_file="meta.csv", + ) + df = IndexLoader(schema, tmp_path).load() + assert list(df.columns) == ["a", "b"] + + def test_missing_format_uses_index_file_extension(self, tmp_path: Path) -> None: + _write(tmp_path / "meta.csv", "a,b\n1,2\n") + schema = DatasetSchema(dataset_id="ds", index_file="meta.csv") + df = IndexLoader(schema, tmp_path).load() + assert list(df.columns) == ["a", "b"] + def test_file_path_dtype_resolves_absolute(self, tmp_path: Path) -> None: - _write_tsv(tmp_path / "index.tsv", "path\nclip.mp3\n") + _write(tmp_path / "index.tsv", "path\nclip.mp3\n") schema = DatasetSchema( dataset_id="ds", - task="ASR", format="tsv", index_file="index.tsv", base_audio_path="clips/", columns={"audio": ColumnMapping(source_column="path", dtype="file_path")}, ) - df = ASRLoader(schema, tmp_path).load() + df = IndexLoader(schema, tmp_path).load() expected = str(tmp_path / "clips" / "clip.mp3") assert df["audio"].iloc[0] == expected @@ -118,11 +139,10 @@ def test_file_path_dtype_resolves_absolute_from_relative_extract_dir( self, tmp_path: Path, monkeypatch: pytest.MonkeyPatch ) -> None: dataset_dir = tmp_path / "dataset" - _write_tsv(dataset_dir / "index.tsv", "path\nclip.mp3\n") + _write(dataset_dir / "index.tsv", "path\nclip.mp3\n") schema = DatasetSchema( dataset_id="ds", - task="ASR", format="tsv", index_file="index.tsv", base_audio_path="clips/", @@ -130,7 +150,7 @@ def test_file_path_dtype_resolves_absolute_from_relative_extract_dir( ) monkeypatch.chdir(tmp_path) - df = ASRLoader(schema, Path("dataset")).load() + df = IndexLoader(schema, Path("dataset")).load() assert Path(df["audio"].iloc[0]).is_absolute() assert df["audio"].iloc[0] == str(dataset_dir / "clips" / "clip.mp3") @@ -138,33 +158,31 @@ def test_file_path_dtype_resolves_absolute_from_relative_extract_dir( def test_file_path_uses_first_existing_base_audio_path( self, tmp_path: Path ) -> None: - _write_tsv(tmp_path / "index.tsv", "path\nclip.wav\n") + _write(tmp_path / "index.tsv", "path\nclip.wav\n") audio_path = tmp_path / "secondary" / "clip.wav" audio_path.parent.mkdir(parents=True, exist_ok=True) audio_path.write_bytes(b"\x00") schema = DatasetSchema( dataset_id="ds", - task="ASR", format="tsv", index_file="index.tsv", base_audio_path=["primary/", "secondary/"], columns={"audio": ColumnMapping(source_column="path", dtype="file_path")}, ) - df = ASRLoader(schema, tmp_path).load() + df = IndexLoader(schema, tmp_path).load() assert df["audio"].iloc[0] == str(audio_path) def test_file_path_exact_search_uses_extension_and_recurses( self, tmp_path: Path ) -> None: - _write_tsv(tmp_path / "index.tsv", "clip_id\nclip_001\n") + _write(tmp_path / "index.tsv", "clip_id\nclip_001\n") audio_path = tmp_path / "audio" / "nested" / "clip_001.wav" audio_path.parent.mkdir(parents=True, exist_ok=True) audio_path.write_bytes(b"\x00") schema = DatasetSchema( dataset_id="ds", - task="ASR", format="tsv", index_file="index.tsv", base_audio_path="audio/", @@ -177,18 +195,17 @@ def test_file_path_exact_search_uses_extension_and_recurses( ) }, ) - df = ASRLoader(schema, tmp_path).load() + df = IndexLoader(schema, tmp_path).load() assert df["audio"].iloc[0] == str(audio_path) def test_file_path_contains_search_matches_substring(self, tmp_path: Path) -> None: - _write_tsv(tmp_path / "index.tsv", "clip_fragment\nclip_001\n") + _write(tmp_path / "index.tsv", "clip_fragment\nclip_001\n") audio_path = tmp_path / "audio" / "nested" / "speaker_clip_001_take2.wav" audio_path.parent.mkdir(parents=True, exist_ok=True) audio_path.write_bytes(b"\x00") schema = DatasetSchema( dataset_id="ds", - task="ASR", format="tsv", index_file="index.tsv", base_audio_path="audio/", @@ -201,30 +218,29 @@ def test_file_path_contains_search_matches_substring(self, tmp_path: Path) -> No ) }, ) - df = ASRLoader(schema, tmp_path).load() + df = IndexLoader(schema, tmp_path).load() assert df["audio"].iloc[0] == str(audio_path) def test_file_path_value_already_includes_base_path(self, tmp_path: Path) -> None: - _write_tsv(tmp_path / "index.tsv", "path\ndata/recipes/clip.wav\n") + _write(tmp_path / "index.tsv", "path\ndata/recipes/clip.wav\n") audio_path = tmp_path / "data" / "recipes" / "clip.wav" audio_path.parent.mkdir(parents=True, exist_ok=True) audio_path.write_bytes(b"\x00") schema = DatasetSchema( dataset_id="ds", - task="ASR", format="tsv", index_file="index.tsv", base_audio_path="data/recipes/", columns={"audio": ColumnMapping(source_column="path", dtype="file_path")}, ) - df = ASRLoader(schema, tmp_path).load() + df = IndexLoader(schema, tmp_path).load() assert df["audio"].iloc[0] == str(audio_path) def test_file_path_template_builds_name_from_multiple_columns( self, tmp_path: Path ) -> None: - _write_tsv( + _write( tmp_path / "dataset" / "data" / "metadata.csv", "Speaker ID,Sentence ID,Sentences\n" "f-adt1-0001,recipes_01_0001_0001,hello\n", @@ -241,7 +257,6 @@ def test_file_path_template_builds_name_from_multiple_columns( schema = DatasetSchema( dataset_id="ds", - task="ASR", index_file="data/metadata.csv", base_audio_path=["data/recipes/", "data/giving_gift/"], columns={ @@ -254,13 +269,13 @@ def test_file_path_template_builds_name_from_multiple_columns( "text": ColumnMapping(source_column="Sentences"), }, ) - df = ASRLoader(schema, tmp_path).load() + df = IndexLoader(schema, tmp_path).load() assert df["audio"].iloc[0] == str(audio_path) def test_file_path_template_renders_dynamic_audio_root_from_metadata( self, tmp_path: Path ) -> None: - _write_tsv( + _write( tmp_path / "dataset" / "data" / "metadata.csv", "Split,Speaker ID,Sentence ID,Sentences\n" "recipes,f-adt1-0001,recipes_01_0001_0001,hello\n", @@ -277,7 +292,6 @@ def test_file_path_template_renders_dynamic_audio_root_from_metadata( schema = DatasetSchema( dataset_id="ds", - task="ASR", index_file="data/metadata.csv", base_audio_path="data/${Split}/", columns={ @@ -290,12 +304,12 @@ def test_file_path_template_renders_dynamic_audio_root_from_metadata( "text": ColumnMapping(source_column="Sentences"), }, ) - df = ASRLoader(schema, tmp_path).load() + df = IndexLoader(schema, tmp_path).load() assert df["audio"].iloc[0] == str(audio_path) assert df["text"].iloc[0] == "hello" def test_contains_search_raises_on_ambiguous_matches(self, tmp_path: Path) -> None: - _write_tsv(tmp_path / "index.tsv", "clip_fragment\nclip_001\n") + _write(tmp_path / "index.tsv", "clip_fragment\nclip_001\n") audio_path_1 = tmp_path / "audio" / "nested" / "speaker_clip_001_take1.wav" audio_path_2 = tmp_path / "audio" / "nested" / "speaker_clip_001_take2.wav" audio_path_1.parent.mkdir(parents=True, exist_ok=True) @@ -304,7 +318,6 @@ def test_contains_search_raises_on_ambiguous_matches(self, tmp_path: Path) -> No schema = DatasetSchema( dataset_id="ds", - task="ASR", format="tsv", index_file="index.tsv", base_audio_path="audio/", @@ -318,16 +331,15 @@ def test_contains_search_raises_on_ambiguous_matches(self, tmp_path: Path) -> No }, ) with pytest.raises(ValueError, match="Ambiguous file_path value"): - ASRLoader(schema, tmp_path).load() + IndexLoader(schema, tmp_path).load() def test_category_dtype(self, tmp_path: Path) -> None: - _write_tsv( + _write( tmp_path / "i.tsv", "path\tsentence\tspk\nc.mp3\thi\tA\nc2.mp3\tbye\tA\n" ) schema = DatasetSchema( dataset_id="ds", - task="ASR", format="tsv", index_file="i.tsv", columns={ @@ -336,17 +348,14 @@ def test_category_dtype(self, tmp_path: Path) -> None: "speaker": ColumnMapping(source_column="spk", dtype="category"), }, ) - df = ASRLoader(schema, tmp_path).load() + df = IndexLoader(schema, tmp_path).load() assert df["speaker"].dtype.name == "category" def test_int_and_float_dtypes(self, tmp_path: Path) -> None: - _write_tsv( - tmp_path / "i.tsv", "path\tsentence\tdur\tscore\nc.mp3\thi\t100\t0.95\n" - ) + _write(tmp_path / "i.tsv", "path\tsentence\tdur\tscore\nc.mp3\thi\t100\t0.95\n") schema = DatasetSchema( dataset_id="ds", - task="ASR", format="tsv", index_file="i.tsv", columns={ @@ -356,16 +365,15 @@ def test_int_and_float_dtypes(self, tmp_path: Path) -> None: "score": ColumnMapping(source_column="score", dtype="float"), }, ) - df = ASRLoader(schema, tmp_path).load() + df = IndexLoader(schema, tmp_path).load() assert df["duration"].iloc[0] == 100 assert df["score"].iloc[0] == pytest.approx(0.95) def test_optional_column_missing(self, tmp_path: Path) -> None: - _write_tsv(tmp_path / "i.tsv", "path\tsentence\nc.mp3\thi\n") + _write(tmp_path / "i.tsv", "path\tsentence\nc.mp3\thi\n") schema = DatasetSchema( dataset_id="ds", - task="ASR", format="tsv", index_file="i.tsv", columns={ @@ -376,15 +384,14 @@ def test_optional_column_missing(self, tmp_path: Path) -> None: ), }, ) - df = ASRLoader(schema, tmp_path).load() + df = IndexLoader(schema, tmp_path).load() assert "speaker" not in df.columns # silently skipped def test_required_column_missing_raises(self, tmp_path: Path) -> None: - _write_tsv(tmp_path / "i.tsv", "path\tsentence\nc.mp3\thi\n") + _write(tmp_path / "i.tsv", "path\tsentence\nc.mp3\thi\n") schema = DatasetSchema( dataset_id="ds", - task="ASR", format="tsv", index_file="i.tsv", columns={ @@ -393,25 +400,23 @@ def test_required_column_missing_raises(self, tmp_path: Path) -> None: }, ) with pytest.raises(KeyError, match="nonexistent"): - ASRLoader(schema, tmp_path).load() + IndexLoader(schema, tmp_path).load() def test_index_file_not_found_raises(self, tmp_path: Path) -> None: schema = DatasetSchema( dataset_id="ds", - task="ASR", format="tsv", index_file="missing.tsv", columns={"a": ColumnMapping(source_column="x")}, ) with pytest.raises(FileNotFoundError, match="missing.tsv"): - ASRLoader(schema, tmp_path).load() + IndexLoader(schema, tmp_path).load() def test_explicit_separator_overrides_format(self, tmp_path: Path) -> None: - _write_tsv(tmp_path / "d.csv", "path|sentence\nc.mp3|hi\n") + _write(tmp_path / "d.csv", "path|sentence\nc.mp3|hi\n") schema = DatasetSchema( dataset_id="ds", - task="ASR", format="csv", separator="|", index_file="d.csv", @@ -420,12 +425,12 @@ def test_explicit_separator_overrides_format(self, tmp_path: Path) -> None: "text": ColumnMapping(source_column="sentence"), }, ) - df = ASRLoader(schema, tmp_path).load() + df = IndexLoader(schema, tmp_path).load() assert len(df) == 1 assert df["text"].iloc[0] == "hi" def test_sniffed_separator_and_trimmed_headers(self, tmp_path: Path) -> None: - _write_tsv( + _write( tmp_path / "metadata.csv", "Topic; Sentence ID ; Sentences \nFood; clip.wav; hello\n", ) @@ -433,7 +438,6 @@ def test_sniffed_separator_and_trimmed_headers(self, tmp_path: Path) -> None: schema = DatasetSchema( dataset_id="ds", - task="ASR", format="csv", index_file="metadata.csv", columns={ @@ -441,18 +445,17 @@ def test_sniffed_separator_and_trimmed_headers(self, tmp_path: Path) -> None: "text": ColumnMapping(source_column="Sentences"), }, ) - df = ASRLoader(schema, tmp_path).load() + df = IndexLoader(schema, tmp_path).load() assert df["audio"].iloc[0] == str(tmp_path / "clip.wav") assert df["text"].iloc[0] == "hello" def test_nested_index_file_found(self, tmp_path: Path) -> None: """Index file inside a subdirectory should be located via rglob.""" nested = tmp_path / "sub" / "deep" - _write_tsv(nested / "train.tsv", "path\tsentence\nc.mp3\thi\n") + _write(nested / "train.tsv", "path\tsentence\nc.mp3\thi\n") schema = DatasetSchema( dataset_id="ds", - task="ASR", format="tsv", index_file="train.tsv", columns={ @@ -460,11 +463,28 @@ def test_nested_index_file_found(self, tmp_path: Path) -> None: "text": ColumnMapping(source_column="sentence"), }, ) - df = ASRLoader(schema, tmp_path).load() + df = IndexLoader(schema, tmp_path).load() assert len(df) == 1 + def test_custom_encoding(self, tmp_path: Path) -> None: + content = "audio\ttext\nc1.wav\tgrüezi\n" + (tmp_path / "meta.tsv").write_text(content, encoding="utf-8-sig") + + schema = DatasetSchema( + dataset_id="ds", + format="tsv", + index_file="meta.tsv", + encoding="utf-8-sig", + columns={ + "audio": ColumnMapping(source_column="audio", dtype="file_path"), + "text": ColumnMapping(source_column="text"), + }, + ) + df = IndexLoader(schema, tmp_path).load() + assert df["text"].iloc[0] == "grüezi" + def test_file_content_dtype_reads_text_file(self, tmp_path: Path) -> None: - _write_tsv( + _write( tmp_path / "index.csv", "audio,transcript\nclip.wav,transcripts/clip.txt\n", ) @@ -474,7 +494,6 @@ def test_file_content_dtype_reads_text_file(self, tmp_path: Path) -> None: schema = DatasetSchema( dataset_id="ds", - task="ASR", format="csv", index_file="index.csv", columns={ @@ -482,11 +501,11 @@ def test_file_content_dtype_reads_text_file(self, tmp_path: Path) -> None: "text": ColumnMapping(source_column="transcript", dtype="file_content"), }, ) - df = ASRLoader(schema, tmp_path).load() + df = IndexLoader(schema, tmp_path).load() assert df["text"].iloc[0] == "hello world" def test_file_content_dtype_with_file_extension(self, tmp_path: Path) -> None: - _write_tsv( + _write( tmp_path / "index.csv", "audio,transcript\nclip.wav,transcripts/clip\n", ) @@ -496,7 +515,6 @@ def test_file_content_dtype_with_file_extension(self, tmp_path: Path) -> None: schema = DatasetSchema( dataset_id="ds", - task="ASR", format="csv", index_file="index.csv", columns={ @@ -508,216 +526,5 @@ def test_file_content_dtype_with_file_extension(self, tmp_path: Path) -> None: ), }, ) - df = ASRLoader(schema, tmp_path).load() + df = IndexLoader(schema, tmp_path).load() assert df["text"].iloc[0] == "resolved with extension" - - -class TestASRMultiSplit: - def test_load_multiple_splits(self, tmp_path: Path) -> None: - _write_tsv(tmp_path / "train.tsv", "path\tsentence\nc1.mp3\thello\n") - _write_tsv(tmp_path / "dev.tsv", "path\tsentence\nc2.mp3\tworld\n") - - schema = DatasetSchema( - dataset_id="ds", - task="ASR", - root_strategy="multi_split", - splits=["train", "dev"], - columns={ - "audio": ColumnMapping(source_column="path", dtype="file_path"), - "text": ColumnMapping(source_column="sentence"), - }, - ) - df = ASRLoader(schema, tmp_path).load() - assert len(df) == 2 - assert set(df["split"]) == {"train", "dev"} - assert "audio" in df.columns - assert "text" in df.columns - - def test_multi_split_without_columns(self, tmp_path: Path) -> None: - """When no column mappings, raw columns + split should be returned.""" - _write_tsv(tmp_path / "train.tsv", "path\tsentence\nc1.mp3\thello\n") - - schema = DatasetSchema( - dataset_id="ds", - task="ASR", - root_strategy="multi_split", - splits=["train"], - ) - df = ASRLoader(schema, tmp_path).load() - assert "split" in df.columns - assert "path" in df.columns # raw column name - assert df["split"].iloc[0] == "train" - - def test_multi_split_custom_pattern(self, tmp_path: Path) -> None: - _write_tsv(tmp_path / "train.csv", "path,sentence\nc1.mp3,hello\n") - - schema = DatasetSchema( - dataset_id="ds", - task="ASR", - root_strategy="multi_split", - splits=["train"], - splits_file_pattern="**/*.csv", - format="csv", - ) - df = ASRLoader(schema, tmp_path).load() - assert len(df) == 1 - - def test_multi_split_ignores_unlisted_splits(self, tmp_path: Path) -> None: - _write_tsv(tmp_path / "train.tsv", "path\tsentence\nc1.mp3\thello\n") - _write_tsv(tmp_path / "other.tsv", "path\tsentence\nc2.mp3\tbye\n") - - schema = DatasetSchema( - dataset_id="ds", - task="ASR", - root_strategy="multi_split", - splits=["train"], # only train, not "other" - ) - df = ASRLoader(schema, tmp_path).load() - assert len(df) == 1 - assert df["split"].iloc[0] == "train" - - def test_multi_split_no_matching_files_raises(self, tmp_path: Path) -> None: - schema = DatasetSchema( - dataset_id="ds", - task="ASR", - root_strategy="multi_split", - splits=["nonexistent"], - ) - with pytest.raises(RuntimeError, match="No split files"): - ASRLoader(schema, tmp_path).load() - - -def _write_json_sidecar(path: Path, filename: str, n_utts: int = 2) -> None: - import json - - path.parent.mkdir(parents=True, exist_ok=True) - data = { - "audio": {"filename": filename, "duration_sec": 12.5, "sample_rate_hz": 44100}, - "metadata": {"gender": "male", "id": "abc123"}, - "transcriptions": [ - { - "utt_id": f"{Path(filename).stem}_{i:04d}", - "speaker": f"SPEAKER{i % 2 + 1}", - "start_time": i * 2.0, - "end_time": i * 2.0 + 1.5, - "text": f"utterance {i}", - } - for i in range(1, n_utts + 1) - ], - } - path.write_text(json.dumps(data), encoding="utf-8") - - -def _paired_glob_json_schema(**overrides) -> DatasetSchema: - fields = { - "dataset_id": "ds", - "task": "ASR", - "root_strategy": "paired_glob", - "format": "json", - "file_pattern": "**/*.merged.json", - "record_path": "transcriptions", - "columns": { - "audio_path": ColumnMapping( - source_column="audio.filename", - dtype="file_path", - path_match_strategy="exact", - ), - "transcription": ColumnMapping(source_column="text"), - "speaker_id": ColumnMapping( - source_column="speaker", dtype="category", optional=True - ), - "start_time": ColumnMapping( - source_column="start_time", dtype="float", optional=True - ), - "gender": ColumnMapping( - source_column="metadata.gender", - dtype="category", - optional=True, - ), - }, - } - fields.update(overrides) - return DatasetSchema(**fields) - - -class TestASRPairedGlobJSONValidation: - def test_requires_json_format(self, tmp_path: Path) -> None: - schema = _paired_glob_json_schema(format="tsv") - with pytest.raises(ValueError, match="format: json"): - ASRLoader(schema, tmp_path) - - def test_requires_file_pattern(self, tmp_path: Path) -> None: - schema = _paired_glob_json_schema(file_pattern=None) - with pytest.raises(ValueError, match="file_pattern"): - ASRLoader(schema, tmp_path) - - def test_requires_columns(self, tmp_path: Path) -> None: - schema = _paired_glob_json_schema(columns={}) - with pytest.raises(ValueError, match="column mapping"): - ASRLoader(schema, tmp_path) - - -class TestASRPairedGlobJSON: - def test_one_row_per_record_with_flattened_meta(self, tmp_path: Path) -> None: - _write_json_sidecar(tmp_path / "rec1.merged.json", "rec1.wav", n_utts=3) - _write_json_sidecar(tmp_path / "rec2.merged.json", "rec2.wav", n_utts=2) - (tmp_path / "rec1.wav").touch() - (tmp_path / "rec2.wav").touch() - - df = ASRLoader(_paired_glob_json_schema(), tmp_path).load() - - assert len(df) == 5 - assert list(df.columns) == [ - "audio_path", - "transcription", - "speaker_id", - "start_time", - "gender", - ] - # Per-recording fields repeat on every utterance row - assert set(Path(p).name for p in df["audio_path"]) == {"rec1.wav", "rec2.wav"} - assert (df["gender"] == "male").all() - assert df["start_time"].dtype == "float64" - assert df["transcription"].iloc[0] == "utterance 1" - - def test_audio_resolved_via_exact_search(self, tmp_path: Path) -> None: - """Audio referenced by bare filename resolves even in nested layouts.""" - _write_json_sidecar(tmp_path / "inner" / "rec1.merged.json", "rec1.wav") - (tmp_path / "inner" / "rec1.wav").touch() - - df = ASRLoader(_paired_glob_json_schema(), tmp_path).load() - assert Path(df["audio_path"].iloc[0]).exists() - - def test_without_record_path_one_row_per_file(self, tmp_path: Path) -> None: - _write_json_sidecar(tmp_path / "rec1.merged.json", "rec1.wav") - (tmp_path / "rec1.wav").touch() - - schema = _paired_glob_json_schema( - record_path=None, - columns={ - "audio_path": ColumnMapping( - source_column="audio.filename", - dtype="file_path", - path_match_strategy="exact", - ), - "gender": ColumnMapping( - source_column="metadata.gender", dtype="category" - ), - }, - ) - df = ASRLoader(schema, tmp_path).load() - assert len(df) == 1 - - def test_missing_record_path_key_raises(self, tmp_path: Path) -> None: - import json - - (tmp_path / "rec1.merged.json").write_text( - json.dumps({"audio": {"filename": "rec1.wav"}}), encoding="utf-8" - ) - - with pytest.raises(KeyError, match="record_path 'transcriptions'"): - ASRLoader(_paired_glob_json_schema(), tmp_path).load() - - def test_no_matching_files_raises(self, tmp_path: Path) -> None: - with pytest.raises(FileNotFoundError, match="No files matching"): - ASRLoader(_paired_glob_json_schema(), tmp_path).load() diff --git a/tests/schema_loaders/strategies/test_multi_sections_loader.py b/tests/schema_loaders/strategies/test_multi_sections_loader.py new file mode 100644 index 0000000..f6d9300 --- /dev/null +++ b/tests/schema_loaders/strategies/test_multi_sections_loader.py @@ -0,0 +1,127 @@ +from __future__ import annotations + +from pathlib import Path + +import pytest + +from datacollective.schema import ColumnMapping, DatasetSchema +from datacollective.schema_loaders.strategies.multi_sections import MultiSectionsLoader + + +def _write(path: Path, content: str) -> None: + path.parent.mkdir(parents=True, exist_ok=True) + path.write_text(content, encoding="utf-8") + + +def _setup_sections(root: Path, sections: list[str]) -> None: + """Create a multi-sections dataset structure under root.""" + for section in sections: + _write( + root / "dataset" / section / "metadata.tsv", + f"audio\ttext\n{section.lower()}.wav\tHello from {section}\n", + ) + + +class TestMultiSectionsValidation: + def test_requires_sections(self, tmp_path: Path) -> None: + schema = DatasetSchema( + dataset_id="ds", + root_strategy="multi_sections", + section_root="dataset", + index_file="metadata.tsv", + ) + with pytest.raises(ValueError, match="sections"): + MultiSectionsLoader(schema, tmp_path) + + def test_requires_section_root(self, tmp_path: Path) -> None: + schema = DatasetSchema( + dataset_id="ds", + root_strategy="multi_sections", + sections=["General"], + index_file="metadata.tsv", + ) + with pytest.raises(ValueError, match="section_root"): + MultiSectionsLoader(schema, tmp_path) + + def test_requires_index_file(self, tmp_path: Path) -> None: + schema = DatasetSchema( + dataset_id="ds", + root_strategy="multi_sections", + sections=["General"], + section_root="dataset", + ) + with pytest.raises(ValueError, match="index_file"): + MultiSectionsLoader(schema, tmp_path) + + +class TestMultiSectionsLoader: + def test_load_multiple_sections(self, tmp_path: Path) -> None: + _setup_sections(tmp_path, ["General", "Chat"]) + + schema = DatasetSchema( + dataset_id="ds", + root_strategy="multi_sections", + section_root="dataset", + sections=["General", "Chat"], + index_file="metadata.tsv", + format="tsv", + ) + df = MultiSectionsLoader(schema, tmp_path).load() + assert len(df) == 2 + assert "section" in df.columns + assert set(df["section"]) == {"General", "Chat"} + assert set(df["text"]) == {"Hello from General", "Hello from Chat"} + + def test_multi_sections_ignores_unlisted_sections(self, tmp_path: Path) -> None: + _setup_sections(tmp_path, ["General", "Chat", "Other"]) + + schema = DatasetSchema( + dataset_id="ds", + root_strategy="multi_sections", + section_root="dataset", + sections=["General", "Chat"], + index_file="metadata.tsv", + format="tsv", + ) + df = MultiSectionsLoader(schema, tmp_path).load() + assert len(df) == 2 + assert set(df["section"]) == {"General", "Chat"} + + def test_multi_sections_missing_index_file_raises(self, tmp_path: Path) -> None: + _write( + tmp_path / "dataset" / "General" / "metadata.tsv", + "audio\ttext\ngeneral.wav\tHello from General\n", + ) + + schema = DatasetSchema( + dataset_id="ds", + root_strategy="multi_sections", + section_root="dataset", + sections=["General", "Chat"], + index_file="metadata.tsv", + format="tsv", + ) + with pytest.raises(FileNotFoundError, match="Chat"): + MultiSectionsLoader(schema, tmp_path).load() + + def test_multi_sections_applies_column_mappings(self, tmp_path: Path) -> None: + """Declared column mappings are applied and the section column is kept.""" + _setup_sections(tmp_path, ["General", "Chat"]) + + schema = DatasetSchema( + dataset_id="ds", + root_strategy="multi_sections", + section_root="dataset", + sections=["General", "Chat"], + index_file="metadata.tsv", + format="tsv", + columns={ + "audio_path": ColumnMapping(source_column="audio", dtype="file_path"), + "transcription": ColumnMapping(source_column="text"), + }, + ) + df = MultiSectionsLoader(schema, tmp_path).load() + assert len(df) == 2 + assert list(df.columns) == ["audio_path", "transcription", "section"] + assert set(df["section"]) == {"General", "Chat"} + assert set(df["transcription"]) == {"Hello from General", "Hello from Chat"} diff --git a/tests/schema_loaders/strategies/test_multi_split_loader.py b/tests/schema_loaders/strategies/test_multi_split_loader.py new file mode 100644 index 0000000..cb7e7f1 --- /dev/null +++ b/tests/schema_loaders/strategies/test_multi_split_loader.py @@ -0,0 +1,90 @@ +from __future__ import annotations + +from pathlib import Path + +import pytest + +from datacollective.schema import ColumnMapping, DatasetSchema +from datacollective.schema_loaders.strategies.multi_split import MultiSplitLoader + + +def _write_tsv(path: Path, content: str) -> None: + path.parent.mkdir(parents=True, exist_ok=True) + path.write_text(content, encoding="utf-8") + + +class TestMultiSplitValidation: + def test_requires_splits(self, tmp_path: Path) -> None: + schema = DatasetSchema(dataset_id="ds", root_strategy="multi_split") + with pytest.raises(ValueError, match="splits"): + MultiSplitLoader(schema, tmp_path) + + +class TestMultiSplitLoader: + def test_load_multiple_splits(self, tmp_path: Path) -> None: + _write_tsv(tmp_path / "train.tsv", "path\tsentence\nc1.mp3\thello\n") + _write_tsv(tmp_path / "dev.tsv", "path\tsentence\nc2.mp3\tworld\n") + + schema = DatasetSchema( + dataset_id="ds", + root_strategy="multi_split", + splits=["train", "dev"], + columns={ + "audio": ColumnMapping(source_column="path", dtype="file_path"), + "text": ColumnMapping(source_column="sentence"), + }, + ) + df = MultiSplitLoader(schema, tmp_path).load() + assert len(df) == 2 + assert set(df["split"]) == {"train", "dev"} + assert "audio" in df.columns + assert "text" in df.columns + + def test_multi_split_without_columns(self, tmp_path: Path) -> None: + """When no column mappings, raw columns + split should be returned.""" + _write_tsv(tmp_path / "train.tsv", "path\tsentence\nc1.mp3\thello\n") + + schema = DatasetSchema( + dataset_id="ds", + root_strategy="multi_split", + splits=["train"], + ) + df = MultiSplitLoader(schema, tmp_path).load() + assert "split" in df.columns + assert "path" in df.columns # raw column name + assert df["split"].iloc[0] == "train" + + def test_multi_split_custom_pattern(self, tmp_path: Path) -> None: + _write_tsv(tmp_path / "train.csv", "path,sentence\nc1.mp3,hello\n") + + schema = DatasetSchema( + dataset_id="ds", + root_strategy="multi_split", + splits=["train"], + splits_file_pattern="**/*.csv", + format="csv", + ) + df = MultiSplitLoader(schema, tmp_path).load() + assert len(df) == 1 + + def test_multi_split_ignores_unlisted_splits(self, tmp_path: Path) -> None: + _write_tsv(tmp_path / "train.tsv", "path\tsentence\nc1.mp3\thello\n") + _write_tsv(tmp_path / "other.tsv", "path\tsentence\nc2.mp3\tbye\n") + + schema = DatasetSchema( + dataset_id="ds", + root_strategy="multi_split", + splits=["train"], # only train, not "other" + ) + df = MultiSplitLoader(schema, tmp_path).load() + assert len(df) == 1 + assert df["split"].iloc[0] == "train" + + def test_multi_split_no_matching_files_raises(self, tmp_path: Path) -> None: + schema = DatasetSchema( + dataset_id="ds", + root_strategy="multi_split", + splits=["nonexistent"], + ) + with pytest.raises(RuntimeError, match="No split files"): + MultiSplitLoader(schema, tmp_path).load() diff --git a/tests/schema_loaders/strategies/test_paired_glob_loader.py b/tests/schema_loaders/strategies/test_paired_glob_loader.py new file mode 100644 index 0000000..0758a58 --- /dev/null +++ b/tests/schema_loaders/strategies/test_paired_glob_loader.py @@ -0,0 +1,280 @@ +from __future__ import annotations + +import json +from pathlib import Path + +import pytest + +from datacollective.schema import ColumnMapping, DatasetSchema +from datacollective.schema_loaders.strategies.paired_glob import PairedGlobLoader + + +def _write(path: Path, content: str) -> None: + path.parent.mkdir(parents=True, exist_ok=True) + path.write_text(content, encoding="utf-8") + + +class TestPairedGlobValidation: + def test_missing_file_pattern_raises(self, tmp_path: Path) -> None: + schema = DatasetSchema( + dataset_id="ds", + root_strategy="paired_glob", + audio_extension=".webm", + ) + with pytest.raises(ValueError, match="file_pattern"): + PairedGlobLoader(schema, tmp_path) + + def test_text_variant_missing_audio_extension_raises(self, tmp_path: Path) -> None: + schema = DatasetSchema( + dataset_id="ds", + root_strategy="paired_glob", + file_pattern="**/*.txt", + ) + with pytest.raises(ValueError, match="audio_extension"): + PairedGlobLoader(schema, tmp_path) + + def test_json_variant_missing_columns_raises(self, tmp_path: Path) -> None: + schema = DatasetSchema( + dataset_id="ds", + root_strategy="paired_glob", + format="json", + file_pattern="**/*.json", + ) + with pytest.raises(ValueError, match="column mapping"): + PairedGlobLoader(schema, tmp_path) + + +class TestPairedGlobText: + def _setup_paired(self, root: Path) -> None: + """Create a paired-glob dataset structure under root.""" + for split in ("split_a", "split_b"): + d = root / split + d.mkdir(parents=True) + _write(d / "001.txt", f"Hello from {split}") + # Create matching audio files + (d / "001.webm").write_bytes(b"\x00") + + def test_load_paired_glob(self, tmp_path: Path) -> None: + self._setup_paired(tmp_path) + + schema = DatasetSchema( + dataset_id="ds", + root_strategy="paired_glob", + file_pattern="**/*.txt", + audio_extension=".webm", + ) + df = PairedGlobLoader(schema, tmp_path).load() + assert len(df) == 2 + assert "audio_path" in df.columns + assert "transcription" in df.columns + assert "split" in df.columns + assert set(df["split"]) == {"split_a", "split_b"} + + def test_paired_glob_skips_missing_audio(self, tmp_path: Path) -> None: + d = tmp_path / "split" + d.mkdir() + _write(d / "001.txt", "hello") + # No matching .webm -> should be skipped + _write(d / "002.txt", "world") + (d / "002.webm").write_bytes(b"\x00") + + schema = DatasetSchema( + dataset_id="ds", + root_strategy="paired_glob", + file_pattern="**/*.txt", + audio_extension=".webm", + ) + df = PairedGlobLoader(schema, tmp_path).load() + assert len(df) == 1 + assert df["transcription"].iloc[0] == "world" + + def test_paired_glob_no_text_files_raises(self, tmp_path: Path) -> None: + schema = DatasetSchema( + dataset_id="ds", + root_strategy="paired_glob", + file_pattern="**/*.txt", + audio_extension=".webm", + ) + with pytest.raises(FileNotFoundError, match="No files matching"): + PairedGlobLoader(schema, tmp_path).load() + + def test_paired_glob_no_matching_audio_raises(self, tmp_path: Path) -> None: + """Text files exist but none have matching audio -> error.""" + d = tmp_path / "split" + d.mkdir() + _write(d / "001.txt", "hello") + + schema = DatasetSchema( + dataset_id="ds", + root_strategy="paired_glob", + file_pattern="**/*.txt", + audio_extension=".webm", + ) + with pytest.raises(FileNotFoundError, match="No paired"): + PairedGlobLoader(schema, tmp_path).load() + + def test_paired_glob_reads_transcription_stripped(self, tmp_path: Path) -> None: + d = tmp_path / "s" + d.mkdir() + _write(d / "001.txt", " hello world \n") + (d / "001.wav").write_bytes(b"\x00") + + schema = DatasetSchema( + dataset_id="ds", + root_strategy="paired_glob", + file_pattern="**/*.txt", + audio_extension=".wav", + ) + df = PairedGlobLoader(schema, tmp_path).load() + assert df["transcription"].iloc[0] == "hello world" + + def test_paired_glob_audio_path_is_absolute(self, tmp_path: Path) -> None: + d = tmp_path / "s" + d.mkdir() + _write(d / "001.txt", "hi") + (d / "001.wav").write_bytes(b"\x00") + + schema = DatasetSchema( + dataset_id="ds", + root_strategy="paired_glob", + file_pattern="**/*.txt", + audio_extension=".wav", + ) + df = PairedGlobLoader(schema, tmp_path).load() + assert Path(df["audio_path"].iloc[0]).is_absolute() + + def test_paired_glob_audio_path_is_absolute_from_relative_extract_dir( + self, tmp_path: Path, monkeypatch: pytest.MonkeyPatch + ) -> None: + dataset_dir = tmp_path / "dataset" + d = dataset_dir / "s" + d.mkdir(parents=True) + _write(d / "001.txt", "hi") + (d / "001.wav").write_bytes(b"\x00") + + schema = DatasetSchema( + dataset_id="ds", + root_strategy="paired_glob", + file_pattern="**/*.txt", + audio_extension=".wav", + ) + + monkeypatch.chdir(tmp_path) + df = PairedGlobLoader(schema, Path("dataset")).load() + + assert Path(df["audio_path"].iloc[0]).is_absolute() + assert df["audio_path"].iloc[0] == str(dataset_dir / "s" / "001.wav") + + +def _write_json_sidecar(path: Path, filename: str, n_utts: int = 2) -> None: + path.parent.mkdir(parents=True, exist_ok=True) + data = { + "audio": {"filename": filename, "duration_sec": 12.5, "sample_rate_hz": 44100}, + "metadata": {"gender": "male", "id": "abc123"}, + "transcriptions": [ + { + "utt_id": f"{Path(filename).stem}_{i:04d}", + "speaker": f"SPEAKER{i % 2 + 1}", + "start_time": i * 2.0, + "end_time": i * 2.0 + 1.5, + "text": f"utterance {i}", + } + for i in range(1, n_utts + 1) + ], + } + path.write_text(json.dumps(data), encoding="utf-8") + + +def _paired_glob_json_schema(**overrides) -> DatasetSchema: + fields = { + "dataset_id": "ds", + "root_strategy": "paired_glob", + "format": "json", + "file_pattern": "**/*.merged.json", + "record_path": "transcriptions", + "columns": { + "audio_path": ColumnMapping( + source_column="audio.filename", + dtype="file_path", + path_match_strategy="exact", + ), + "transcription": ColumnMapping(source_column="text"), + "speaker_id": ColumnMapping( + source_column="speaker", dtype="category", optional=True + ), + "start_time": ColumnMapping( + source_column="start_time", dtype="float", optional=True + ), + "gender": ColumnMapping( + source_column="metadata.gender", + dtype="category", + optional=True, + ), + }, + } + fields.update(overrides) + return DatasetSchema(**fields) + + +class TestPairedGlobJSON: + def test_one_row_per_record_with_flattened_meta(self, tmp_path: Path) -> None: + _write_json_sidecar(tmp_path / "rec1.merged.json", "rec1.wav", n_utts=3) + _write_json_sidecar(tmp_path / "rec2.merged.json", "rec2.wav", n_utts=2) + (tmp_path / "rec1.wav").touch() + (tmp_path / "rec2.wav").touch() + + df = PairedGlobLoader(_paired_glob_json_schema(), tmp_path).load() + + assert len(df) == 5 + assert list(df.columns) == [ + "audio_path", + "transcription", + "speaker_id", + "start_time", + "gender", + ] + # Per-recording fields repeat on every utterance row + assert set(Path(p).name for p in df["audio_path"]) == {"rec1.wav", "rec2.wav"} + assert (df["gender"] == "male").all() + assert df["start_time"].dtype == "float64" + assert df["transcription"].iloc[0] == "utterance 1" + + def test_audio_resolved_via_exact_search(self, tmp_path: Path) -> None: + """Audio referenced by bare filename resolves even in nested layouts.""" + _write_json_sidecar(tmp_path / "inner" / "rec1.merged.json", "rec1.wav") + (tmp_path / "inner" / "rec1.wav").touch() + + df = PairedGlobLoader(_paired_glob_json_schema(), tmp_path).load() + assert Path(df["audio_path"].iloc[0]).exists() + + def test_without_record_path_one_row_per_file(self, tmp_path: Path) -> None: + _write_json_sidecar(tmp_path / "rec1.merged.json", "rec1.wav") + (tmp_path / "rec1.wav").touch() + + schema = _paired_glob_json_schema( + record_path=None, + columns={ + "audio_path": ColumnMapping( + source_column="audio.filename", + dtype="file_path", + path_match_strategy="exact", + ), + "gender": ColumnMapping( + source_column="metadata.gender", dtype="category" + ), + }, + ) + df = PairedGlobLoader(schema, tmp_path).load() + assert len(df) == 1 + + def test_missing_record_path_key_raises(self, tmp_path: Path) -> None: + (tmp_path / "rec1.merged.json").write_text( + json.dumps({"audio": {"filename": "rec1.wav"}}), encoding="utf-8" + ) + + with pytest.raises(KeyError, match="record_path 'transcriptions'"): + PairedGlobLoader(_paired_glob_json_schema(), tmp_path).load() + + def test_no_matching_files_raises(self, tmp_path: Path) -> None: + with pytest.raises(FileNotFoundError, match="No files matching"): + PairedGlobLoader(_paired_glob_json_schema(), tmp_path).load() diff --git a/tests/schema_loaders/tasks/__init__.py b/tests/schema_loaders/tasks/__init__.py deleted file mode 100644 index e69de29..0000000 diff --git a/tests/schema_loaders/tasks/test_tts_loader.py b/tests/schema_loaders/tasks/test_tts_loader.py deleted file mode 100644 index ceee82f..0000000 --- a/tests/schema_loaders/tasks/test_tts_loader.py +++ /dev/null @@ -1,337 +0,0 @@ -from __future__ import annotations - -from pathlib import Path - -import pytest - -from datacollective.schema import ColumnMapping, DatasetSchema -from datacollective.schema_loaders.tasks.tts import TTSLoader - - -def _write(path: Path, content: str) -> None: - path.parent.mkdir(parents=True, exist_ok=True) - path.write_text(content, encoding="utf-8") - - -class TestTTSIndexBased: - def test_load_pipe_delimited(self, tmp_path: Path) -> None: - _write(tmp_path / "meta.csv", "clip1.mp3|hello world\nclip2.mp3|goodbye\n") - - schema = DatasetSchema( - dataset_id="ds", - task="TTS", - format="pipe", - separator="|", - has_header=False, - index_file="meta.csv", - base_audio_path="wavs/", - columns={ - "audio_path": ColumnMapping(source_column=0, dtype="file_path"), - "transcription": ColumnMapping(source_column=1, dtype="string"), - }, - ) - df = TTSLoader(schema, tmp_path).load() - assert len(df) == 2 - assert df["transcription"].iloc[0] == "hello world" - # file_path dtype -> absolute path with base_audio_path - assert "wavs" in df["audio_path"].iloc[0] - - def test_load_tsv_with_header(self, tmp_path: Path) -> None: - _write(tmp_path / "meta.tsv", "audio\ttext\nc1.wav\thi\nc2.wav\tbye\n") - - schema = DatasetSchema( - dataset_id="ds", - task="TTS", - format="tsv", - index_file="meta.tsv", - columns={ - "audio": ColumnMapping(source_column="audio", dtype="file_path"), - "text": ColumnMapping(source_column="text"), - }, - ) - df = TTSLoader(schema, tmp_path).load() - assert len(df) == 2 - assert df["text"].iloc[1] == "bye" - - def test_no_columns_returns_raw(self, tmp_path: Path) -> None: - _write(tmp_path / "meta.csv", "a,b\n1,2\n") - - schema = DatasetSchema( - dataset_id="ds", - task="TTS", - format="csv", - index_file="meta.csv", - ) - df = TTSLoader(schema, tmp_path).load() - assert list(df.columns) == ["a", "b"] - - def test_missing_index_file_raises(self, tmp_path: Path) -> None: - schema = DatasetSchema(dataset_id="ds", task="TTS") - with pytest.raises(ValueError, match="index_file"): - TTSLoader(schema, tmp_path).load() - - def test_missing_format_uses_index_file_extension(self, tmp_path: Path) -> None: - _write(tmp_path / "meta.csv", "a,b\n1,2\n") - schema = DatasetSchema(dataset_id="ds", task="TTS", index_file="meta.csv") - df = TTSLoader(schema, tmp_path).load() - assert list(df.columns) == ["a", "b"] - - def test_custom_encoding(self, tmp_path: Path) -> None: - content = "audio\ttext\nc1.wav\tgrüezi\n" - (tmp_path / "meta.tsv").write_text(content, encoding="utf-8-sig") - - schema = DatasetSchema( - dataset_id="ds", - task="TTS", - format="tsv", - index_file="meta.tsv", - encoding="utf-8-sig", - columns={ - "audio": ColumnMapping(source_column="audio", dtype="file_path"), - "text": ColumnMapping(source_column="text"), - }, - ) - df = TTSLoader(schema, tmp_path).load() - assert df["text"].iloc[0] == "grüezi" - - def test_file_path_template_renders_dynamic_audio_root_from_metadata( - self, tmp_path: Path - ) -> None: - _write( - tmp_path / "dataset" / "metadata.tsv", - "split\tspeaker_id\tsentence_id\ttext\nrecipes\tspk-01\tsent-01\thello\n", - ) - audio_path = tmp_path / "dataset" / "recipes" / "spk-01_khm_sent-01.wav" - audio_path.parent.mkdir(parents=True, exist_ok=True) - audio_path.write_bytes(b"\x00") - - schema = DatasetSchema( - dataset_id="ds", - task="TTS", - format="tsv", - index_file="metadata.tsv", - base_audio_path="${split}/", - columns={ - "audio": ColumnMapping( - source_column="sentence_id", - dtype="file_path", - file_extension=".wav", - path_template="${speaker_id}_khm_${value}", - ), - "text": ColumnMapping(source_column="text"), - }, - ) - df = TTSLoader(schema, tmp_path / "dataset").load() - assert df["audio"].iloc[0] == str(audio_path) - assert df["text"].iloc[0] == "hello" - - -class TestTTSPairedGlob: - def _setup_paired(self, root: Path) -> None: - """Create a paired-glob dataset structure under root.""" - for split in ("split_a", "split_b"): - d = root / split - d.mkdir(parents=True) - _write(d / "001.txt", f"Hello from {split}") - # Create matching audio files - (d / "001.webm").write_bytes(b"\x00") - - def test_load_paired_glob(self, tmp_path: Path) -> None: - self._setup_paired(tmp_path) - - schema = DatasetSchema( - dataset_id="ds", - task="TTS", - root_strategy="paired_glob", - file_pattern="**/*.txt", - audio_extension=".webm", - ) - df = TTSLoader(schema, tmp_path).load() - assert len(df) == 2 - assert "audio_path" in df.columns - assert "transcription" in df.columns - assert "split" in df.columns - assert set(df["split"]) == {"split_a", "split_b"} - - def test_paired_glob_skips_missing_audio(self, tmp_path: Path) -> None: - d = tmp_path / "split" - d.mkdir() - _write(d / "001.txt", "hello") - # No matching .webm -> should be skipped - _write(d / "002.txt", "world") - (d / "002.webm").write_bytes(b"\x00") - - schema = DatasetSchema( - dataset_id="ds", - task="TTS", - root_strategy="paired_glob", - file_pattern="**/*.txt", - audio_extension=".webm", - ) - df = TTSLoader(schema, tmp_path).load() - assert len(df) == 1 - assert df["transcription"].iloc[0] == "world" - - def test_paired_glob_no_text_files_raises(self, tmp_path: Path) -> None: - schema = DatasetSchema( - dataset_id="ds", - task="TTS", - root_strategy="paired_glob", - file_pattern="**/*.txt", - audio_extension=".webm", - ) - with pytest.raises(FileNotFoundError, match="No files matching"): - TTSLoader(schema, tmp_path).load() - - def test_paired_glob_no_matching_audio_raises(self, tmp_path: Path) -> None: - """Text files exist but none have matching audio -> error.""" - d = tmp_path / "split" - d.mkdir() - _write(d / "001.txt", "hello") - - schema = DatasetSchema( - dataset_id="ds", - task="TTS", - root_strategy="paired_glob", - file_pattern="**/*.txt", - audio_extension=".webm", - ) - with pytest.raises(FileNotFoundError, match="No paired"): - TTSLoader(schema, tmp_path).load() - - def test_paired_glob_missing_file_pattern_raises(self, tmp_path: Path) -> None: - schema = DatasetSchema( - dataset_id="ds", - task="TTS", - root_strategy="paired_glob", - audio_extension=".webm", - ) - with pytest.raises(ValueError, match="file_pattern"): - TTSLoader(schema, tmp_path).load() - - def test_paired_glob_missing_audio_extension_raises(self, tmp_path: Path) -> None: - schema = DatasetSchema( - dataset_id="ds", - task="TTS", - root_strategy="paired_glob", - file_pattern="**/*.txt", - ) - with pytest.raises(ValueError, match="audio_extension"): - TTSLoader(schema, tmp_path).load() - - def test_paired_glob_reads_transcription_stripped(self, tmp_path: Path) -> None: - d = tmp_path / "s" - d.mkdir() - _write(d / "001.txt", " hello world \n") - (d / "001.wav").write_bytes(b"\x00") - - schema = DatasetSchema( - dataset_id="ds", - task="TTS", - root_strategy="paired_glob", - file_pattern="**/*.txt", - audio_extension=".wav", - ) - df = TTSLoader(schema, tmp_path).load() - assert df["transcription"].iloc[0] == "hello world" - - def test_paired_glob_audio_path_is_absolute(self, tmp_path: Path) -> None: - d = tmp_path / "s" - d.mkdir() - _write(d / "001.txt", "hi") - (d / "001.wav").write_bytes(b"\x00") - - schema = DatasetSchema( - dataset_id="ds", - task="TTS", - root_strategy="paired_glob", - file_pattern="**/*.txt", - audio_extension=".wav", - ) - df = TTSLoader(schema, tmp_path).load() - assert Path(df["audio_path"].iloc[0]).is_absolute() - - def test_paired_glob_audio_path_is_absolute_from_relative_extract_dir( - self, tmp_path: Path, monkeypatch: pytest.MonkeyPatch - ) -> None: - dataset_dir = tmp_path / "dataset" - d = dataset_dir / "s" - d.mkdir(parents=True) - _write(d / "001.txt", "hi") - (d / "001.wav").write_bytes(b"\x00") - - schema = DatasetSchema( - dataset_id="ds", - task="TTS", - root_strategy="paired_glob", - file_pattern="**/*.txt", - audio_extension=".wav", - ) - - monkeypatch.chdir(tmp_path) - df = TTSLoader(schema, Path("dataset")).load() - - assert Path(df["audio_path"].iloc[0]).is_absolute() - assert df["audio_path"].iloc[0] == str(dataset_dir / "s" / "001.wav") - - -class TestTTSMultiSections: - def _setup_sections(self, root: Path, sections: list[str]) -> None: - """Create a multi-sections dataset structure under root.""" - for section in sections: - _write( - root / "dataset" / section / "metadata.tsv", - f"audio\ttext\n{section.lower()}.wav\tHello from {section}\n", - ) - - def test_load_multiple_sections(self, tmp_path: Path) -> None: - self._setup_sections(tmp_path, ["General", "Chat"]) - - schema = DatasetSchema( - dataset_id="ds", - task="TTS", - root_strategy="multi_sections", - section_root="dataset", - sections=["General", "Chat"], - index_file="metadata.tsv", - format="tsv", - ) - df = TTSLoader(schema, tmp_path).load() - assert len(df) == 2 - assert "section" in df.columns - assert set(df["section"]) == {"General", "Chat"} - assert set(df["text"]) == {"Hello from General", "Hello from Chat"} - - def test_multi_sections_ignores_unlisted_sections(self, tmp_path: Path) -> None: - self._setup_sections(tmp_path, ["General", "Chat", "Other"]) - - schema = DatasetSchema( - dataset_id="ds", - task="TTS", - root_strategy="multi_sections", - section_root="dataset", - sections=["General", "Chat"], - index_file="metadata.tsv", - format="tsv", - ) - df = TTSLoader(schema, tmp_path).load() - assert len(df) == 2 - assert set(df["section"]) == {"General", "Chat"} - - def test_multi_sections_missing_index_file_raises(self, tmp_path: Path) -> None: - _write( - tmp_path / "dataset" / "General" / "metadata.tsv", - "audio\ttext\ngeneral.wav\tHello from General\n", - ) - - schema = DatasetSchema( - dataset_id="ds", - task="TTS", - root_strategy="multi_sections", - section_root="dataset", - sections=["General", "Chat"], - index_file="metadata.tsv", - format="tsv", - ) - with pytest.raises(FileNotFoundError, match="Chat"): - TTSLoader(schema, tmp_path).load() diff --git a/tests/schema_loaders/test_registry.py b/tests/schema_loaders/test_registry.py index ae8c07c..9977460 100644 --- a/tests/schema_loaders/test_registry.py +++ b/tests/schema_loaders/test_registry.py @@ -4,6 +4,7 @@ import pytest +from datacollective.errors import TaskValidationError from datacollective.schema import ColumnMapping, DatasetSchema from datacollective.schema_loaders.registry import _load_dataset_from_schema @@ -13,9 +14,9 @@ def _write(path: Path, content: str) -> None: path.write_text(content, encoding="utf-8") -class TestLoadDatasetFromSchema: - def test_dispatches_asr(self, tmp_path: Path) -> None: - """An ASR index-based schema should be dispatched to ASRLoader and produce a DataFrame.""" +class TestStrategyDispatch: + def test_default_strategy_is_index(self, tmp_path: Path) -> None: + """Without a root_strategy the schema is loaded via the index strategy.""" _write( tmp_path / "train.tsv", "path\tsentence\nclip1.mp3\thello\nclip2.mp3\tworld\n", @@ -23,7 +24,6 @@ def test_dispatches_asr(self, tmp_path: Path) -> None: schema = DatasetSchema( dataset_id="test", - task="ASR", format="tsv", index_file="train.tsv", columns={ @@ -37,34 +37,14 @@ def test_dispatches_asr(self, tmp_path: Path) -> None: assert len(df) == 2 assert list(df.columns) == ["audio_path", "transcription"] - def test_dispatches_tts_index(self, tmp_path: Path) -> None: - """A TTS index-based schema should be dispatched to TTSLoader.""" - _write(tmp_path / "meta.csv", "audio,text\nc1.wav,hello\n") - - schema = DatasetSchema( - dataset_id="test-tts", - task="TTS", - format="csv", - index_file="meta.csv", - columns={ - "audio": ColumnMapping(source_column="audio", dtype="file_path"), - "text": ColumnMapping(source_column="text"), - }, - ) - df = _load_dataset_from_schema(schema, tmp_path) - assert len(df) == 1 - assert "audio" in df.columns - - def test_dispatches_tts_paired_glob(self, tmp_path: Path) -> None: - """A TTS paired-glob schema should be dispatched to TTSLoader.""" + def test_dispatches_paired_glob(self, tmp_path: Path) -> None: d = tmp_path / "split" d.mkdir() _write(d / "001.txt", "hello") (d / "001.wav").write_bytes(b"\x00") schema = DatasetSchema( - dataset_id="test-tts-pg", - task="TTS", + dataset_id="test-pg", root_strategy="paired_glob", file_pattern="**/*.txt", audio_extension=".wav", @@ -74,8 +54,147 @@ def test_dispatches_tts_paired_glob(self, tmp_path: Path) -> None: assert "audio_path" in df.columns assert "transcription" in df.columns - def test_dispatches_oth_index(self, tmp_path: Path) -> None: - """An OTH index-based schema should load via OTHLoader and apply column mappings.""" + def test_dispatches_multi_split(self, tmp_path: Path) -> None: + _write(tmp_path / "train.tsv", "path\tsentence\nc1.mp3\thello\n") + _write(tmp_path / "dev.tsv", "path\tsentence\nc2.mp3\tworld\n") + + schema = DatasetSchema( + dataset_id="test-ms", + root_strategy="multi_split", + splits=["train", "dev"], + ) + df = _load_dataset_from_schema(schema, tmp_path) + assert len(df) == 2 + assert set(df["split"]) == {"train", "dev"} + + def test_dispatches_multi_sections(self, tmp_path: Path) -> None: + for section in ("General", "Chat"): + _write( + tmp_path / "dataset" / section / "metadata.tsv", + f"audio\ttext\n{section.lower()}.wav\tHello from {section}\n", + ) + + schema = DatasetSchema( + dataset_id="test-msec", + root_strategy="multi_sections", + section_root="dataset", + sections=["General", "Chat"], + index_file="metadata.tsv", + format="tsv", + ) + df = _load_dataset_from_schema(schema, tmp_path) + assert len(df) == 2 + assert set(df["section"]) == {"General", "Chat"} + + def test_dispatches_glob(self, tmp_path: Path) -> None: + _write(tmp_path / "spk1" / "en" / "a.wav", "") + _write(tmp_path / "spk2" / "fr" / "b.wav", "") + + schema = DatasetSchema( + dataset_id="test-glob", + root_strategy="glob", + file_pattern="**/*.wav", + ) + df = _load_dataset_from_schema(schema, tmp_path) + assert len(df) == 2 + assert list(df.columns) == ["audio_path", "language", "speaker_id"] + assert set(df["language"]) == {"en", "fr"} + assert set(df["speaker_id"]) == {"spk1", "spk2"} + + def test_any_strategy_works_with_any_task(self, tmp_path: Path) -> None: + """Strategies are task-agnostic: e.g. TTS + multi_split is expressible.""" + _write(tmp_path / "train.tsv", "path\tsentence\nc1.mp3\thello\n") + + schema = DatasetSchema( + dataset_id="test-tts-ms", + task="TTS", + root_strategy="multi_split", + splits=["train"], + columns={ + "audio_path": ColumnMapping(source_column="path", dtype="file_path"), + "transcription": ColumnMapping(source_column="sentence"), + }, + ) + df = _load_dataset_from_schema(schema, tmp_path) + assert len(df) == 1 + assert {"audio_path", "transcription", "split"} <= set(df.columns) + + def test_glob_requires_file_pattern(self, tmp_path: Path) -> None: + schema = DatasetSchema(dataset_id="test-glob", root_strategy="glob") + with pytest.raises(ValueError, match="must specify 'file_pattern'"): + _load_dataset_from_schema(schema, tmp_path) + + def test_index_requires_index_file(self, tmp_path: Path) -> None: + schema = DatasetSchema(dataset_id="test-idx") + with pytest.raises(ValueError, match="index_file"): + _load_dataset_from_schema(schema, tmp_path) + + def test_unknown_strategy_raises(self, tmp_path: Path) -> None: + schema = DatasetSchema(dataset_id="ds", root_strategy="unknown_strategy") + with pytest.raises(ValueError, match="Unknown root_strategy"): + _load_dataset_from_schema(schema, tmp_path) + + +class TestTaskContracts: + def test_asr_contract_satisfied(self, tmp_path: Path) -> None: + _write(tmp_path / "train.tsv", "path\tsentence\nc1.mp3\thello\n") + + schema = DatasetSchema( + dataset_id="test-asr", + task="ASR", + format="tsv", + index_file="train.tsv", + columns={ + "audio_path": ColumnMapping(source_column="path", dtype="file_path"), + "transcription": ColumnMapping(source_column="sentence"), + }, + ) + df = _load_dataset_from_schema(schema, tmp_path) + assert {"audio_path", "transcription"} <= set(df.columns) + + def test_declared_columns_violating_contract_fail_fast( + self, tmp_path: Path + ) -> None: + """Misconfigured mappings are rejected before any file resolution.""" + schema = DatasetSchema( + dataset_id="test-asr-bad", + task="ASR", + format="tsv", + index_file="missing.tsv", # never touched: validation fails first + columns={ + "audio": ColumnMapping(source_column="path", dtype="file_path"), + "text": ColumnMapping(source_column="sentence"), + }, + ) + with pytest.raises(TaskValidationError, match="audio_path"): + _load_dataset_from_schema(schema, tmp_path) + + def test_raw_load_violating_contract_raises_post_load(self, tmp_path: Path) -> None: + """A columns-less index schema loads raw, then fails the task contract.""" + _write(tmp_path / "train.tsv", "path\tsentence\nc1.mp3\thello\n") + + schema = DatasetSchema( + dataset_id="test-asr-raw", + task="ASR", + format="tsv", + index_file="train.tsv", + ) + with pytest.raises(TaskValidationError, match="ASR"): + _load_dataset_from_schema(schema, tmp_path) + + def test_unknown_task_loads_without_validation(self, tmp_path: Path) -> None: + _write(tmp_path / "data.tsv", "a\tb\n1\t2\n") + + schema = DatasetSchema( + dataset_id="test-unknown", + task="BRAND_NEW_TASK", + format="tsv", + index_file="data.tsv", + ) + df = _load_dataset_from_schema(schema, tmp_path) + assert list(df.columns) == ["a", "b"] + + def test_oth_task_has_no_contract(self, tmp_path: Path) -> None: _write( tmp_path / "data.tsv", "id\tsentence\tlang\n1\thello\ten\n2\tbonjour\tfr\n", @@ -97,36 +216,42 @@ def test_dispatches_oth_index(self, tmp_path: Path) -> None: assert list(df.columns) == ["id", "sentence", "lang"] assert df["lang"].dtype.name == "category" - def test_dispatches_oth_glob(self, tmp_path: Path) -> None: - """An OTH glob schema should load via OTHLoader and derive metadata from the path.""" + def test_no_task_loads_without_validation(self, tmp_path: Path) -> None: + _write(tmp_path / "data.csv", "a,b\n1,2\n") + + schema = DatasetSchema( + dataset_id="test-no-task", + format="csv", + index_file="data.csv", + ) + df = _load_dataset_from_schema(schema, tmp_path) + assert list(df.columns) == ["a", "b"] + + def test_glob_strategy_skips_declared_contract_check(self, tmp_path: Path) -> None: + """Glob output is fixed (audio_path/...), so declared columns are not checked.""" _write(tmp_path / "spk1" / "en" / "a.wav", "") - _write(tmp_path / "spk2" / "fr" / "b.wav", "") schema = DatasetSchema( - dataset_id="test-oth", + dataset_id="test-oth-glob", task="OTH", root_strategy="glob", file_pattern="**/*.wav", ) df = _load_dataset_from_schema(schema, tmp_path) - assert len(df) == 2 - assert list(df.columns) == ["audio_path", "language", "speaker_id"] - assert set(df["language"]) == {"en", "fr"} - assert set(df["speaker_id"]) == {"spk1", "spk2"} - - def test_oth_glob_requires_file_pattern(self, tmp_path: Path) -> None: - schema = DatasetSchema(dataset_id="test-oth", task="OTH", root_strategy="glob") - with pytest.raises(ValueError, match="must specify 'file_pattern'"): - _load_dataset_from_schema(schema, tmp_path) + assert len(df) == 1 - def test_oth_requires_index_file(self, tmp_path: Path) -> None: - schema = DatasetSchema(dataset_id="test-oth", task="OTH") - with pytest.raises( - ValueError, match="'root_strategy: glob'.*or an 'index_file'" - ): - _load_dataset_from_schema(schema, tmp_path) + def test_paired_glob_text_variant_satisfies_contract(self, tmp_path: Path) -> None: + d = tmp_path / "split" + d.mkdir() + _write(d / "001.txt", "hello") + (d / "001.wav").write_bytes(b"\x00") - def test_unknown_task_raises(self, tmp_path: Path) -> None: - schema = DatasetSchema(dataset_id="ds", task="UNKNOWN_TASK") - with pytest.raises(ValueError, match="No schema loader registered"): - _load_dataset_from_schema(schema, tmp_path) + schema = DatasetSchema( + dataset_id="test-tts-pg", + task="TTS", + root_strategy="paired_glob", + file_pattern="**/*.txt", + audio_extension=".wav", + ) + df = _load_dataset_from_schema(schema, tmp_path) + assert {"audio_path", "transcription"} <= set(df.columns) diff --git a/tests/schema_loaders/test_schema_loading_e2e.py b/tests/schema_loaders/test_schema_loading_e2e.py index 1fba9f8..1c49cab 100644 --- a/tests/schema_loaders/test_schema_loading_e2e.py +++ b/tests/schema_loaders/test_schema_loading_e2e.py @@ -14,6 +14,7 @@ import pandas as pd import pytest +from datacollective.errors import TaskValidationError from datacollective.schema import DatasetSchema, _parse_schema from datacollective.schema_loaders.registry import _load_dataset_from_schema @@ -28,7 +29,7 @@ def _schema_from_dict(d: dict) -> DatasetSchema: class TestASRIndexE2E: - """Full pipeline: TSV/CSV index -> ASRLoader -> DataFrame.""" + """Full pipeline: TSV/CSV index -> IndexLoader (+ ASR contract) -> DataFrame.""" def test_tsv_basic(self, tmp_path: Path) -> None: _write( @@ -69,8 +70,8 @@ def test_csv_with_optional_column(self, tmp_path: Path) -> None: "format": "csv", "index_file": "data.csv", "columns": { - "audio": {"source_column": "path", "dtype": "file_path"}, - "text": {"source_column": "sentence"}, + "audio_path": {"source_column": "path", "dtype": "file_path"}, + "transcription": {"source_column": "sentence"}, "missing": { "source_column": "nonexistent", "dtype": "string", @@ -96,8 +97,8 @@ def test_nested_index_file(self, tmp_path: Path) -> None: "format": "tsv", "index_file": "meta.tsv", "columns": { - "audio": {"source_column": "path", "dtype": "file_path"}, - "text": {"source_column": "sentence"}, + "audio_path": {"source_column": "path", "dtype": "file_path"}, + "transcription": {"source_column": "sentence"}, }, } ) @@ -116,8 +117,8 @@ def test_int_and_float_columns(self, tmp_path: Path) -> None: "format": "tsv", "index_file": "meta.tsv", "columns": { - "audio": {"source_column": "path", "dtype": "file_path"}, - "text": {"source_column": "sentence"}, + "audio_path": {"source_column": "path", "dtype": "file_path"}, + "transcription": {"source_column": "sentence"}, "age": {"source_column": "age", "dtype": "int"}, "score": {"source_column": "score", "dtype": "float"}, }, @@ -228,8 +229,8 @@ def test_three_splits(self, tmp_path: Path) -> None: "root_strategy": "multi_split", "splits": ["train", "dev", "test"], "columns": { - "audio": {"source_column": "path", "dtype": "file_path"}, - "text": {"source_column": "sentence"}, + "audio_path": {"source_column": "path", "dtype": "file_path"}, + "transcription": {"source_column": "sentence"}, }, } ) @@ -245,7 +246,6 @@ def test_subset_of_splits(self, tmp_path: Path) -> None: schema = _schema_from_dict( { "dataset_id": "asr-ms2", - "task": "ASR", "root_strategy": "multi_split", "splits": ["train", "dev"], } @@ -259,7 +259,6 @@ def test_custom_file_pattern(self, tmp_path: Path) -> None: schema = _schema_from_dict( { "dataset_id": "asr-ms-csv", - "task": "ASR", "root_strategy": "multi_split", "splits": ["train"], "splits_file_pattern": "**/*.csv", @@ -307,8 +306,8 @@ def test_tsv_with_header(self, tmp_path: Path) -> None: "format": "tsv", "index_file": "meta.tsv", "columns": { - "audio": {"source_column": "audio", "dtype": "file_path"}, - "text": {"source_column": "text"}, + "audio_path": {"source_column": "audio", "dtype": "file_path"}, + "transcription": {"source_column": "text"}, }, } ) @@ -321,7 +320,6 @@ def test_no_column_mappings_returns_raw(self, tmp_path: Path) -> None: schema = _schema_from_dict( { "dataset_id": "tts-raw", - "task": "TTS", "format": "csv", "index_file": "meta.csv", } @@ -342,13 +340,13 @@ def test_custom_encoding(self, tmp_path: Path) -> None: "index_file": "meta.tsv", "encoding": "utf-8-sig", "columns": { - "audio": {"source_column": "audio", "dtype": "file_path"}, - "text": {"source_column": "text"}, + "audio_path": {"source_column": "audio", "dtype": "file_path"}, + "transcription": {"source_column": "text"}, }, } ) df = _load_dataset_from_schema(schema, tmp_path) - assert df["text"].iloc[0] == "grüezi" + assert df["transcription"].iloc[0] == "grüezi" class TestTTSPairedGlobE2E: @@ -454,7 +452,6 @@ def test_basic(self, tmp_path: Path) -> None: schema = _schema_from_dict( { "dataset_id": "tts-ms", - "task": "TTS", "root_strategy": "multi_sections", "section_root": "dataset", "sections": ["General", "Chat"], @@ -467,13 +464,36 @@ def test_basic(self, tmp_path: Path) -> None: assert set(df["section"]) == {"General", "Chat"} assert df["text"].tolist() == ["Text General", "Text Chat"] + def test_with_column_mappings_and_task_contract(self, tmp_path: Path) -> None: + """multi_sections applies declared mappings and satisfies the TTS contract.""" + self._create_multi_sections_dataset(tmp_path, ["General", "Chat"]) + + schema = _schema_from_dict( + { + "dataset_id": "tts-ms-mapped", + "task": "TTS", + "root_strategy": "multi_sections", + "section_root": "dataset", + "sections": ["General", "Chat"], + "index_file": "metadata.tsv", + "format": "tsv", + "columns": { + "audio_path": {"source_column": "audio", "dtype": "file_path"}, + "transcription": {"source_column": "text"}, + }, + } + ) + df = _load_dataset_from_schema(schema, tmp_path) + assert len(df) == 2 + assert list(df.columns) == ["audio_path", "transcription", "section"] + assert set(df["section"]) == {"General", "Chat"} + def test_ignores_unlisted_sections(self, tmp_path: Path) -> None: self._create_multi_sections_dataset(tmp_path, ["General", "Chat", "Other"]) schema = _schema_from_dict( { "dataset_id": "tts-ms-subset", - "task": "TTS", "root_strategy": "multi_sections", "section_root": "dataset", "sections": ["General", "Chat"], @@ -487,9 +507,9 @@ def test_ignores_unlisted_sections(self, tmp_path: Path) -> None: class TestErrorPaths: - def test_unknown_task_raises(self, tmp_path: Path) -> None: - schema = DatasetSchema(dataset_id="ds", task="UNKNOWN_TASK") - with pytest.raises(ValueError, match="No schema loader registered"): + def test_unknown_strategy_raises(self, tmp_path: Path) -> None: + schema = DatasetSchema(dataset_id="ds", root_strategy="unknown_strategy") + with pytest.raises(ValueError, match="Unknown root_strategy"): _load_dataset_from_schema(schema, tmp_path) def test_asr_missing_index_file_raises(self, tmp_path: Path) -> None: @@ -499,7 +519,10 @@ def test_asr_missing_index_file_raises(self, tmp_path: Path) -> None: "task": "ASR", "format": "tsv", "index_file": "missing.tsv", - "columns": {"a": {"source_column": "x"}}, + "columns": { + "audio_path": {"source_column": "path", "dtype": "file_path"}, + "transcription": {"source_column": "sentence"}, + }, } ) with pytest.raises(FileNotFoundError, match="missing.tsv"): @@ -513,12 +536,43 @@ def test_asr_missing_required_column_raises(self, tmp_path: Path) -> None: "task": "ASR", "format": "tsv", "index_file": "d.tsv", - "columns": {"audio": {"source_column": "nonexistent"}}, + "columns": { + "audio_path": {"source_column": "nonexistent"}, + "transcription": {"source_column": "sentence"}, + }, } ) with pytest.raises(KeyError, match="nonexistent"): _load_dataset_from_schema(schema, tmp_path) + def test_asr_contract_violation_fails_fast(self, tmp_path: Path) -> None: + """Declared mappings that cannot satisfy the ASR contract are rejected.""" + schema = _schema_from_dict( + { + "dataset_id": "asr-contract", + "task": "ASR", + "format": "tsv", + "index_file": "missing.tsv", + "columns": {"a": {"source_column": "x"}}, + } + ) + with pytest.raises(TaskValidationError, match="audio_path"): + _load_dataset_from_schema(schema, tmp_path) + + def test_asr_raw_load_contract_violation_raises(self, tmp_path: Path) -> None: + """A columns-less ASR schema loads raw, then fails the task contract.""" + _write(tmp_path / "d.tsv", "path\tsentence\nc.mp3\thi\n") + schema = _schema_from_dict( + { + "dataset_id": "asr-raw", + "task": "ASR", + "format": "tsv", + "index_file": "d.tsv", + } + ) + with pytest.raises(TaskValidationError, match="ASR"): + _load_dataset_from_schema(schema, tmp_path) + def test_tts_paired_glob_no_text_files_raises(self, tmp_path: Path) -> None: schema = _schema_from_dict( { @@ -537,7 +591,6 @@ def test_tts_index_no_format_uses_index_extension(self, tmp_path: Path) -> None: schema = _schema_from_dict( { "dataset_id": "tts-nf", - "task": "TTS", "index_file": "meta.csv", } ) diff --git a/tests/test_schema.py b/tests/test_schema.py index 14a0a02..ffc25c2 100644 --- a/tests/test_schema.py +++ b/tests/test_schema.py @@ -155,9 +155,9 @@ def test_missing_dataset_id_raises(self) -> None: with pytest.raises(ValueError, match="dataset_id"): _parse_schema({"task": "ASR"}) - def test_missing_task_raises(self) -> None: - with pytest.raises(ValueError, match="dataset_id.*task"): - _parse_schema({"dataset_id": "ds1"}) + def test_missing_task_is_allowed(self) -> None: + s = _parse_schema({"dataset_id": "ds1"}) + assert s.task is None def test_invalid_yaml_type_raises(self) -> None: with pytest.raises(ValueError, match="Expected a dict"):