feat(datasets): add dataset materialisation utility

Author: fin-griffin

cspell/library-words.txt

@@ -0,0 +1 @@
+huggingface

pyproject.toml

@@ -7,6 +7,7 @@ requires-python = ">=3.12"
 dependencies = [
     "datasets>=4.5.0",
     "dotenv>=0.9.9",
+    "huggingface-hub>=1.4.1",
 ]
 
 [build-system]
@@ -23,6 +24,7 @@ dev = [
     "flake8-cognitive-complexity>=0.1.0",
     "isort>=7.0.0",
     "mypy>=1.19.1",
+    "notebook>=7.5.3",
     "pre-commit>=4.5.1",
     "pydoclint>=0.8.3",
     "pytest>=9.0.2",
@@ -83,7 +85,7 @@ module = ["tests.*"]
 ignore_errors = true
 
 [[tool.mypy.overrides]]
-module = ["datasets", "datasets.*"]
+module = ["datasets", "datasets.*", "huggingface_hub", "huggingface_hub.*"]
 ignore_missing_imports = true
 
 [tool.isort]

src/voice/__init__.py

@@ -5,7 +5,7 @@
 LLMs for stylistic fidelity.
 """
 
-from voice.datasets import DatasetSpec
+from voice.datasets import DatasetSpec, get_dataset
 from voice.stylometry import get_metrics
 
-__all__: list[str] = ["DatasetSpec", "get_metrics"]
+__all__: list[str] = ["DatasetSpec", "get_metrics", "get_dataset"]

src/voice/datasets/__init__.py

@@ -1,7 +1,15 @@
-"""Placeholder docstring."""
+"""
+Datasets package for VOICE.
 
-from voice.datasets.dataset import DatasetSpec
+This package contains user-facing abstractions for dataset management,
+in particular we expose:
 
-__all__: list[str] = [
-    "DatasetSpec",
-]
+- `DatasetSpec` - Lightweight, declarative specification of a dataset.
+- `VoiceDataset` - Split-aware wrapper around a Hugging Face dataset.
+- `get_dataset` - Entrypoint for materialising a dataset from a specification.
+"""
+
+from voice.datasets.dataset import DatasetSpec, VoiceDataset
+from voice.datasets.get_dataset import get_dataset
+
+__all__: list[str] = ["DatasetSpec", "VoiceDataset", "get_dataset"]

src/voice/datasets/_schema.py

@@ -1,4 +1,11 @@
-"""Placeholder docstring."""
+"""
+Dataset schema primitives.
+
+This module defines lightweight schema-level abstractions shared across
+the dataset loading pipeline.
+
+This module is private by convention and not part of the public API.
+"""
 
 from __future__ import annotations
 

src/voice/datasets/_specs.py

@@ -0,0 +1,16 @@
+"""
+Pre-defined dataset specifications for testing.
+
+All specifications are of type `voice.datasets.DatasetSpec`.
+
+This module is private by convention and not part of the public API.
+"""
+
+from voice.datasets import DatasetSpec
+from voice.datasets._schema import Split
+
+BUSH_LATEST = DatasetSpec(
+    repo_id="AccelerateScience/bush-dataset",
+    revision="main",
+    splits=(Split.TRAIN, Split.VALIDATION, Split.TEST),
+)

src/voice/datasets/dataset.py

@@ -1,4 +1,17 @@
-"""Placeholder docstring."""
+"""
+Core dataset abstractions for the `voice.datasets` package.
+
+This module defines:
+
+- `DatasetSpec` - A lightweight declarative specification of a dataset
+  hosted on Hugging Face.
+- `_PinnedDatasetSpec` - An internal revision-pinned specification
+  resolved to a specific git commit SHA.
+- `Example` - A representation of a single question-answer example
+  with provenance metadata.
+- `VoiceDataset` — A split-aware wrapper providing structured access
+  to examples across dataset splits.
+"""
 
 from __future__ import annotations
 
@@ -24,9 +37,6 @@ class DatasetSpec:
     """
     Lightweight specification of a Hugging Face dataset.
 
-    `_pin()` is used to resolve symbolic revisions to a commit hash.
-    It is private by convention and not part of the public API
-
     .. attribute :: repo_id
 
         Hugging Face dataset repository id, i.e. namespace/name

src/voice/datasets/get_dataset.py

@@ -0,0 +1,181 @@
+"""
+Dataset loading and canonicalisation utilities.
+
+This module implements the dataset materialisation pipeline
+used by the public entrypoint `get_dataset`.
+
+Responsibilities include:
+
+- Resolving symbolic revisions (branches/tags) to commit SHAs.
+- Loading splits from Hugging Face via `datasets.load_dataset`.
+- Canonicalising dataset schemas into a uniform
+  (system, question, answer) format.
+- Validating dataset structure.
+- Constructing a split-aware `VoiceDataset` instance.
+"""
+
+from __future__ import annotations
+
+from datasets import Dataset, load_dataset
+from huggingface_hub import HfApi
+from huggingface_hub.utils import HfHubHTTPError
+
+from voice.datasets._schema import Split
+from voice.datasets.dataset import (
+    DatasetSpec,
+    VoiceDataset,
+    _PinnedDatasetSpec,
+)
+
+# -----------------------------------------------------------------------------
+# Utilities (revision resolution, schema canonicalisation .etc)
+# -----------------------------------------------------------------------------
+
+
+def _resolve_revision(repo_id: str, revision: str) -> str:
+    """
+    Resolve a symbolic revision (branch or tag) to a commit SHA.
+
+    :param repo_id: Hugging Face dataset repository id
+    :param revision: Branch name, tag, or commit hash
+    :return: Commit SHA corresponding to the revision
+    :raises RuntimeError: If dataset or revision does not exist
+    """
+    api = HfApi()
+
+    try:
+        info = api.dataset_info(repo_id=repo_id, revision=revision)
+    except HfHubHTTPError as e:
+        raise RuntimeError(
+            f"Failed to resolve revision '{revision}' for dataset '{repo_id}'."
+        ) from e
+
+    return str(info.sha)
+
+
+def _extract_chat_style(
+    column: dict[str, list[dict[str, str]]],
+) -> dict[str, str]:
+    """
+    Extract (system, question, answer) dict from chat style example.
+
+    :param column: column containing chat-style messages
+    :return: canonical dict
+    """
+    messages = column["messages"]
+
+    system = ""
+    question = ""
+    answer = ""
+
+    for msg in messages:
+        role = msg.get("role")
+        content = msg.get("content", "")
+
+        if role == "system":
+            system = content
+        elif role == "user":
+            question = content
+        elif role == "assistant":
+            answer = content
+
+    return {
+        "system": system,
+        "question": question,
+        "answer": answer,
+    }
+
+
+def _canonicalise_dataset(ds: Dataset) -> Dataset:
+    """
+    Convert a HF dataset into canonical (system, question, answer) format.
+
+    Currently supported schemas:
+    - Already canonical: columns {system, question, answer}
+    - Chat style: column 'messages' with list[{role, content}]
+
+    :param ds: Raw Hugging Face dataset
+    :return: Canonicalised QA dataset
+    :raises NotImplementedError: If schema not recognised
+    """
+    cols = set(ds.column_names)
+
+    if {"system", "question", "answer"}.issubset(cols):
+        return ds
+
+    if "messages" in cols and len(cols) == 1:
+        return ds.map(
+            _extract_chat_style,
+            remove_columns=ds.column_names,
+        )
+
+    raise NotImplementedError(
+        f"Unsupported dataset schema. Columns found: {sorted(cols)}"
+    )
+
+
+def _validate_canonical(ds: Dataset) -> None:
+    """
+    Ensure dataset contains required canonical columns.
+
+    :param ds: Canonicalised dataset
+    :return: None
+    :raises ValueError: If required columns missing
+    """
+    required = {"system", "question", "answer"}
+    missing = required - set(ds.column_names)
+    if not required.issubset(ds.column_names):
+        raise ValueError(f"Missing required columns: {sorted(missing)}")
+
+
+# -----------------------------------------------------------------------------
+# Public API
+# -----------------------------------------------------------------------------
+
+
+def get_dataset(spec: DatasetSpec) -> VoiceDataset:
+    """
+    Materialise a Hugging Face dataset as a VoiceDataset.
+
+    Steps:
+        1. Resolve DatasetSpec to _PinnedDatasetSpec.
+        2. Load requested splits from Hugging Face.
+        3. Canonicalise each split into (system, question, answer).
+        4. Validate canonical structure.
+        5. Construct VoiceDataset.
+
+    :param spec: Dataset specification
+    :return: Split-aware VoiceDataset
+    :raises RuntimeError: If a split cannot be loaded
+    """
+    # Resolve dataset revision
+    pinned: _PinnedDatasetSpec = spec._pin(_resolve_revision)
+
+    # Load dataset splits
+    datasets: dict[Split, Dataset] = {}
+
+    for split in pinned.splits:
+        try:
+            ds = load_dataset(
+                pinned.repo_id,
+                split=split.value,
+                revision=pinned.revision,
+            )
+        except Exception as e:
+            raise RuntimeError(
+                f"Failed to load split '{split.value}' "
+                f"from {pinned.repo_id}@{pinned.revision}"
+            ) from e
+
+        # Canonicalise and validate dataset
+        ds = _canonicalise_dataset(ds)
+
+        _validate_canonical(ds)
+
+        datasets[split] = ds
+
+    # Build VoiceDataset object
+    return VoiceDataset(
+        datasets=datasets,
+        spec=pinned,
+    )