themis-eval 0.1.0__py3-none-any.whl → 0.2.0__py3-none-any.whl

themis/datasets/piqa.py

@@ -0,0 +1,146 @@
+"""Helpers for working with the ybisk/piqa dataset."""
+
+from __future__ import annotations
+
+import json
+from pathlib import Path
+from typing import Any, Iterable, Iterator, List, Sequence
+
+from pydantic import BaseModel, Field, field_validator
+
+_DATASET_NAME = "ybisk/piqa"
+
+
+class PiqaSample(BaseModel):
+    unique_id: str
+    goal: str
+    choices: list[str]
+    answer: str
+    metadata: dict[str, Any] = Field(default_factory=dict)
+
+    @field_validator("choices", mode="before")
+    @classmethod
+    def _ensure_choices(cls, value: Any) -> list[str]:
+        if value is None:
+            return []
+        if isinstance(value, (list, tuple)):
+            return [str(item) for item in value]
+        raise TypeError("choices must be a sequence")
+
+    def to_generation_example(self) -> dict[str, Any]:
+        return {
+            "unique_id": self.unique_id,
+            "goal": self.goal,
+            "choices": list(self.choices),
+            "answer": self.answer,
+            "metadata": dict(self.metadata),
+        }
+
+
+def load_piqa(
+    *,
+    split: str = "validation", # Test set usually has no labels
+    limit: int | None = None,
+    source: str = "huggingface",
+    data_dir: str | Path | None = None,
+) -> List[PiqaSample]:
+    """Load PIQA samples from Hugging Face or a local directory."""
+
+    if source not in {"huggingface", "local"}:
+        raise ValueError(
+            f"Unsupported source '{source}'. Expected one of: 'huggingface', 'local'."
+        )
+
+    if source == "huggingface":
+        rows = _load_from_huggingface(split=split)
+    else:
+        if data_dir is None:
+            raise ValueError(
+                "data_dir must be provided when source='local'. "
+                "Pass dataset.data_dir in configs or --data-dir on the CLI."
+            )
+        rows = _load_from_local(Path(data_dir))
+
+    samples: list[PiqaSample] = []
+    for index, row in enumerate(rows, start=1):
+        sample = _row_to_sample(row, index=index)
+        samples.append(sample)
+        if limit is not None and len(samples) >= limit:
+            break
+    return samples
+
+
+def _row_to_sample(row: dict[str, Any], *, index: int) -> PiqaSample:
+    unique_id = (
+        row.get("id")
+        or row.get("unique_id")
+        or f"piqa-{index:05d}"
+    )
+    goal = row.get("goal") or ""
+    
+    # PIQA has 'sol1', 'sol2'
+    choices = [
+        str(row.get("sol1") or ""),
+        str(row.get("sol2") or ""),
+    ]
+    
+    # label is integer 0 or 1
+    label = row.get("label")
+    answer = ""
+    if label is not None:
+        try:
+            label_int = int(label)
+            if 0 <= label_int < len(choices):
+                answer = choices[label_int]
+        except (ValueError, TypeError):
+            pass
+
+    metadata_keys = {
+        "goal", "sol1", "sol2", "label", "id"
+    }
+    metadata = {key: value for key, value in row.items() if key not in metadata_keys}
+    
+    return PiqaSample(
+        unique_id=str(unique_id),
+        goal=str(goal),
+        choices=choices,
+        answer=answer,
+        metadata=metadata,
+    )
+
+
+def _load_from_huggingface(*, split: str) -> Iterable[dict[str, Any]]:
+    try:
+        from datasets import load_dataset
+    except ImportError as exc:  # pragma: no cover - optional dependency
+        raise RuntimeError(
+            "datasets is required to load PIQA from Hugging Face. Install it via `uv pip install '.[hf]'`."
+        ) from exc
+
+    dataset = load_dataset(_DATASET_NAME, split=split)
+    for row in dataset:
+        yield dict(row)
+
+
+def _load_from_local(root: Path) -> Iterator[dict[str, Any]]:
+    if not root.exists():
+        raise FileNotFoundError(f"Local dataset directory not found: {root}")
+
+    for path in root.rglob("*"):
+        if path.suffix.lower() == ".json":
+            with path.open("r", encoding="utf-8") as handle:
+                row = json.load(handle)
+                row.setdefault("id", path.stem)
+                yield row
+        elif path.suffix.lower() in {".jsonl", ".ndjson"}:
+            with path.open("r", encoding="utf-8") as handle:
+                for line_num, line in enumerate(handle, start=1):
+                    line = line.strip()
+                    if not line:
+                        continue
+                    row = json.loads(line)
+                    row.setdefault("id", f"{path.stem}-{line_num}")
+                    yield row
+
+
+__all__ = ["PiqaSample", "load_piqa"]

themis/datasets/registry.py

@@ -0,0 +1,201 @@
+"""Registry for dataset loaders.
+
+This module provides a plugin-based registry system for datasets, allowing
+users to register custom datasets without modifying core Themis code.
+
+Example:
+    ```python
+    from themis.datasets import register_dataset
+
+    def create_my_dataset(**options):
+        from my_module import MyDataset
+        return MyDataset(path=options.get('path'))
+
+    register_dataset('my-dataset', create_my_dataset)
+    ```
+"""
+
+from __future__ import annotations
+
+from typing import Any, Callable
+
+# Factory type: takes config options, returns list of samples
+DatasetFactory = Callable[[dict[str, Any]], list[dict[str, Any]]]
+
+
+class DatasetRegistry:
+    """Registry for dataset loaders.
+
+    Maintains a mapping from dataset names to factory functions that
+    load and return dataset samples.
+    """
+
+    def __init__(self):
+        self._datasets: dict[str, DatasetFactory] = {}
+
+    def register(self, name: str, factory: DatasetFactory) -> None:
+        """Register a dataset factory.
+
+        Args:
+            name: Unique identifier for the dataset (e.g., 'math500', 'my-dataset')
+            factory: Callable that takes config options and returns list of samples
+
+        Raises:
+            ValueError: If dataset name is already registered
+        """
+        if name in self._datasets:
+            raise ValueError(
+                f"Dataset '{name}' is already registered. "
+                f"Use a different name or unregister the existing dataset first."
+            )
+        self._datasets[name] = factory
+
+    def unregister(self, name: str) -> None:
+        """Unregister a dataset.
+
+        Args:
+            name: Dataset identifier to remove
+
+        Raises:
+            ValueError: If dataset name is not registered
+        """
+        if name not in self._datasets:
+            raise ValueError(f"Dataset '{name}' is not registered")
+        del self._datasets[name]
+
+    def create(self, name: str, **options) -> list[dict[str, Any]]:
+        """Create a dataset instance by loading samples.
+
+        Args:
+            name: Registered dataset identifier
+            **options: Configuration options passed to the factory function
+                Common options include:
+                - source: 'huggingface', 'local', or custom source
+                - data_dir: Path for local datasets
+                - split: Dataset split (e.g., 'train', 'test')
+                - limit: Maximum number of samples to load
+                - subjects: List of subjects to filter
+
+        Returns:
+            List of sample dictionaries ready for generation
+
+        Raises:
+            ValueError: If dataset name is not registered
+        """
+        if name not in self._datasets:
+            available = list(self._datasets.keys())
+            raise ValueError(
+                f"Unknown dataset: '{name}'. "
+                f"Available datasets: {', '.join(sorted(available)) or 'none'}"
+            )
+        factory = self._datasets[name]
+        return factory(options)
+
+    def list_datasets(self) -> list[str]:
+        """List all registered dataset names.
+
+        Returns:
+            Sorted list of registered dataset identifiers
+        """
+        return sorted(self._datasets.keys())
+
+    def is_registered(self, name: str) -> bool:
+        """Check if a dataset is registered.
+
+        Args:
+            name: Dataset identifier to check
+
+        Returns:
+            True if the dataset is registered, False otherwise
+        """
+        return name in self._datasets
+
+
+# Global registry instance
+_REGISTRY = DatasetRegistry()
+
+
+def register_dataset(name: str, factory: DatasetFactory) -> None:
+    """Register a dataset factory in the global registry.
+
+    Args:
+        name: Unique identifier for the dataset
+        factory: Callable that takes config options and returns samples
+
+    Example:
+        ```python
+        def create_my_dataset(options):
+            from my_module import load_data
+            return load_data(
+                path=options.get('path'),
+                limit=options.get('limit')
+            )
+
+        register_dataset('my-dataset', create_my_dataset)
+        ```
+    """
+    _REGISTRY.register(name, factory)
+
+
+def unregister_dataset(name: str) -> None:
+    """Unregister a dataset from the global registry.
+
+    Args:
+        name: Dataset identifier to remove
+    """
+    _REGISTRY.unregister(name)
+
+
+def create_dataset(name: str, **options) -> list[dict[str, Any]]:
+    """Create a dataset by loading samples from a registered factory.
+
+    Args:
+        name: Registered dataset identifier
+        **options: Configuration options for the dataset
+
+    Returns:
+        List of sample dictionaries
+
+    Example:
+        ```python
+        samples = create_dataset(
+            'math500',
+            source='huggingface',
+            split='test',
+            limit=10
+        )
+        ```
+    """
+    return _REGISTRY.create(name, **options)
+
+
+def list_datasets() -> list[str]:
+    """List all registered datasets.
+
+    Returns:
+        Sorted list of dataset names
+    """
+    return _REGISTRY.list_datasets()
+
+
+def is_dataset_registered(name: str) -> bool:
+    """Check if a dataset is registered.
+
+    Args:
+        name: Dataset identifier
+
+    Returns:
+        True if registered, False otherwise
+    """
+    return _REGISTRY.is_registered(name)
+
+
+__all__ = [
+    "DatasetFactory",
+    "DatasetRegistry",
+    "register_dataset",
+    "unregister_dataset",
+    "create_dataset",
+    "list_datasets",
+    "is_dataset_registered",
+]

themis/datasets/schema.py

@@ -0,0 +1,245 @@
+"""Dataset schema and metadata definitions.
+
+This module provides enhanced dataset abstractions with schema validation,
+metadata, and filtering capabilities while maintaining backward compatibility.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from typing import Any, Callable, Iterable, Protocol, runtime_checkable
+
+
+@dataclass
+class DatasetSchema:
+    """Describes the structure and validation rules for dataset samples.
+
+    Examples:
+        # Basic schema
+        schema = DatasetSchema(
+            id_field="unique_id",
+            reference_field="answer",
+            required_fields={"unique_id", "problem", "answer"},
+        )
+
+        # Schema with validation
+        def validate_problem(sample: dict) -> None:
+            if len(sample.get("problem", "")) < 10:
+                raise ValueError("Problem text too short")
+
+        schema = DatasetSchema(
+            id_field="id",
+            reference_field="expected",
+            required_fields={"id", "problem", "expected"},
+            validators=[validate_problem],
+        )
+    """
+
+    id_field: str
+    reference_field: str | None
+    required_fields: set[str] = field(default_factory=set)
+    optional_fields: set[str] = field(default_factory=set)
+    metadata_fields: set[str] = field(default_factory=set)
+    validators: list[Callable[[dict], None]] = field(default_factory=list)
+
+    def validate_sample(self, sample: dict[str, Any]) -> None:
+        """Validate a single sample against this schema.
+
+        Args:
+            sample: Sample to validate
+
+        Raises:
+            ValueError: If validation fails
+        """
+        # Check required fields
+        for field_name in self.required_fields:
+            if field_name not in sample:
+                raise ValueError(
+                    f"Missing required field '{field_name}' in sample {sample.get(self.id_field)}"
+                )
+
+        # Run custom validators
+        for validator in self.validators:
+            validator(sample)
+
+    def get_all_fields(self) -> set[str]:
+        """Get all known fields (required + optional + metadata)."""
+        return self.required_fields | self.optional_fields | self.metadata_fields
+
+
+@dataclass
+class DatasetMetadata:
+    """Metadata about the entire dataset.
+
+    This provides information useful for experiment planning, reporting,
+    and understanding dataset characteristics.
+
+    Examples:
+        metadata = DatasetMetadata(
+            name="MATH-500",
+            version="1.0",
+            total_samples=500,
+            categories={
+                "subject": ["algebra", "geometry", "number_theory"],
+                "difficulty": ["easy", "medium", "hard"],
+            },
+            difficulty_distribution={
+                "easy": 100,
+                "medium": 250,
+                "hard": 150,
+            },
+            description="Math problems from competition mathematics",
+        )
+    """
+
+    name: str
+    version: str = "1.0"
+    total_samples: int | None = None
+    categories: dict[str, list[str]] = field(default_factory=dict)
+    difficulty_distribution: dict[str, int] | None = None
+    description: str = ""
+    source_url: str | None = None
+    license: str | None = None
+    citation: str | None = None
+    custom_metadata: dict[str, Any] = field(default_factory=dict)
+
+    def get_category_values(self, category: str) -> list[str]:
+        """Get all possible values for a category."""
+        return self.categories.get(category, [])
+
+    def has_category(self, category: str) -> bool:
+        """Check if dataset has a specific category."""
+        return category in self.categories
+
+
+@runtime_checkable
+class EnhancedDatasetAdapter(Protocol):
+    """Extended dataset interface with schema and metadata support.
+
+    This protocol extends the basic DatasetAdapter with additional
+    capabilities for schema validation, filtering, and stratification.
+    """
+
+    def iter_samples(self) -> Iterable[dict[str, Any]]:
+        """Iterate over dataset samples."""
+        ...
+
+    def get_schema(self) -> DatasetSchema:
+        """Get the dataset schema."""
+        ...
+
+    def get_metadata(self) -> DatasetMetadata:
+        """Get dataset metadata."""
+        ...
+
+    def filter(
+        self, predicate: Callable[[dict[str, Any]], bool]
+    ) -> EnhancedDatasetAdapter:
+        """Return filtered view of dataset.
+
+        Args:
+            predicate: Function that returns True for samples to keep
+
+        Returns:
+            New dataset adapter with filtered samples
+        """
+        ...
+
+    def limit(self, n: int) -> EnhancedDatasetAdapter:
+        """Return dataset limited to first n samples.
+
+        Args:
+            n: Maximum number of samples
+
+        Returns:
+            New dataset adapter with limited samples
+        """
+        ...
+
+    def stratify(
+        self, field: str, distribution: dict[str, float]
+    ) -> EnhancedDatasetAdapter:
+        """Return stratified sample of dataset.
+
+        Args:
+            field: Field to stratify by
+            distribution: Desired distribution (values should sum to 1.0)
+
+        Returns:
+            New dataset adapter with stratified samples
+        """
+        ...
+
+
+# Common validators
+def validate_non_empty_field(field_name: str) -> Callable[[dict], None]:
+    """Create validator that ensures field is non-empty.
+
+    Args:
+        field_name: Name of field to validate
+
+    Returns:
+        Validator function
+    """
+
+    def validator(sample: dict) -> None:
+        value = sample.get(field_name)
+        if not value:
+            raise ValueError(f"Field '{field_name}' cannot be empty")
+
+    return validator
+
+
+def validate_field_type(field_name: str, expected_type: type) -> Callable[[dict], None]:
+    """Create validator that ensures field has correct type.
+
+    Args:
+        field_name: Name of field to validate
+        expected_type: Expected type
+
+    Returns:
+        Validator function
+    """
+
+    def validator(sample: dict) -> None:
+        value = sample.get(field_name)
+        if value is not None and not isinstance(value, expected_type):
+            raise ValueError(
+                f"Field '{field_name}' expected type {expected_type.__name__}, "
+                f"got {type(value).__name__}"
+            )
+
+    return validator
+
+
+def validate_field_in_choices(
+    field_name: str, choices: set[str]
+) -> Callable[[dict], None]:
+    """Create validator that ensures field value is in allowed choices.
+
+    Args:
+        field_name: Name of field to validate
+        choices: Set of allowed values
+
+    Returns:
+        Validator function
+    """
+
+    def validator(sample: dict) -> None:
+        value = sample.get(field_name)
+        if value is not None and value not in choices:
+            raise ValueError(
+                f"Field '{field_name}' value '{value}' not in allowed choices: {choices}"
+            )
+
+    return validator
+
+
+__all__ = [
+    "DatasetSchema",
+    "DatasetMetadata",
+    "EnhancedDatasetAdapter",
+    "validate_non_empty_field",
+    "validate_field_type",
+    "validate_field_in_choices",
+]