omniloader 1.0.0__py3-none-any.whl

omniloader/__init__.py

@@ -0,0 +1,141 @@
+"""OmniLoader: a PyTorch meta data loader for disjoint, multi-task datasets.
+
+OmniLoader unifies several datasets with heterogeneous annotations into a single
+masked sample scheme so a model can be trained jointly across all of them. It is
+modality- and model-agnostic: it reasons only about vectors and sequences. See
+:class:`~omniloader.loader.OmniLoader` for the entry point.
+"""
+
+from __future__ import annotations
+
+from importlib.metadata import PackageNotFoundError, version
+
+from omniloader.collate import DynamicCollator, unified_collate
+from omniloader.config import OmniConfig, seed_everything
+from omniloader.data.datasets import DictTensorDataset, HDF5Dataset
+from omniloader.data.factory import build_datasets
+from omniloader.data.npy import NpyFolderDataset
+from omniloader.data.splits import load_split_info, save_split_info, split_indices
+from omniloader.introspection import Report, describe, validate
+from omniloader.loader import OmniLoader
+from omniloader.sampling.bucketing import LengthBucketBatchSampler
+from omniloader.sampling.sampler import OmniSampler
+from omniloader.sampling.strategies import (
+    AnnealedTemperatureStrategy,
+    FixedWeightStrategy,
+    MixingStrategy,
+    ProportionalStrategy,
+    RoundRobinStrategy,
+    SubsampleConfig,
+    TemperatureStrategy,
+)
+from omniloader.sampling.subsamplers import ExhaustionPolicy, IndexPool
+from omniloader.sampling.weights import (
+    class_histogram,
+    class_weights_for_loss,
+    class_weights_for_sampler,
+)
+from omniloader.schema.spec import (
+    DatasetSchema,
+    TensorSpec,
+    UnifiedSchema,
+)
+from omniloader.schema.unify import SampleUnifier
+from omniloader.templates import config_template_path
+from omniloader.transforms import (
+    CenterCrop,
+    Compose,
+    FeatureDropout,
+    FeatureMasking,
+    GaussianNoise,
+    InstanceNormalize,
+    MinMaxNormalize,
+    MixupCollator,
+    Normalize,
+    PerDatasetNormalize,
+    RandomCrop,
+    RobustNormalize,
+    SpanMasking,
+    TimeWarp,
+    Transform,
+    build_transform,
+    compute_dataset_stats,
+    compute_stats,
+    load_stats,
+    save_stats,
+)
+from omniloader.transforms.stats import compute_feature_stats
+from omniloader.utils.padding import (
+    fill_gaps_with_repeat,
+    pad_or_crop_time_dim,
+    repeat_pad_time_dim,
+)
+from omniloader.utils.seeding import seed_worker
+
+try:
+    __version__ = version("omniloader")
+except PackageNotFoundError:  # pragma: no cover - only during local, uninstalled use
+    __version__ = "1.0.0"
+
+__all__ = [
+    "AnnealedTemperatureStrategy",
+    "CenterCrop",
+    "Compose",
+    "DatasetSchema",
+    "DictTensorDataset",
+    "DynamicCollator",
+    "ExhaustionPolicy",
+    "FeatureDropout",
+    "FeatureMasking",
+    "FixedWeightStrategy",
+    "GaussianNoise",
+    "HDF5Dataset",
+    "IndexPool",
+    "InstanceNormalize",
+    "LengthBucketBatchSampler",
+    "MinMaxNormalize",
+    "MixingStrategy",
+    "MixupCollator",
+    "Normalize",
+    "NpyFolderDataset",
+    "OmniConfig",
+    "OmniLoader",
+    "OmniSampler",
+    "PerDatasetNormalize",
+    "ProportionalStrategy",
+    "RandomCrop",
+    "Report",
+    "RobustNormalize",
+    "RoundRobinStrategy",
+    "SampleUnifier",
+    "SpanMasking",
+    "SubsampleConfig",
+    "TemperatureStrategy",
+    "TensorSpec",
+    "TimeWarp",
+    "Transform",
+    "UnifiedSchema",
+    "__version__",
+    "build_datasets",
+    "build_transform",
+    "class_histogram",
+    "class_weights_for_loss",
+    "class_weights_for_sampler",
+    "compute_dataset_stats",
+    "compute_feature_stats",
+    "compute_stats",
+    "config_template_path",
+    "describe",
+    "fill_gaps_with_repeat",
+    "load_split_info",
+    "load_stats",
+    "pad_or_crop_time_dim",
+    "repeat_pad_time_dim",
+    "save_split_info",
+    "save_stats",
+    "seed_everything",
+    "seed_worker",
+    "split_indices",
+    "unified_collate",
+    "validate",
+]

omniloader/cli.py

@@ -0,0 +1,156 @@
+"""Command-line interface for inspecting a dataset mix from a config file.
+
+Subcommands (all take a JSON/YAML config whose ``datasets`` section declares the
+adapters and schemas — see :mod:`omniloader.data.factory`):
+
+* ``omniloader describe CONFIG`` — print the coverage/statistics report.
+* ``omniloader validate CONFIG`` — check datasets against their specs (exit 1 on issues).
+* ``omniloader compute-stats CONFIG -o STATS`` — compute and save normalization stats.
+* ``omniloader class-weights-for-loss CONFIG --target KEY -o WEIGHTS`` — compute and
+  save per-class loss weights (and counts) for a categorical target.
+"""
+
+from __future__ import annotations
+
+import argparse
+import json
+from pathlib import Path
+from typing import TYPE_CHECKING
+
+from omniloader.config import OmniConfig
+from omniloader.introspection import describe, validate
+from omniloader.loader import OmniLoader
+from omniloader.sampling.weights import class_histogram, class_weights_for_loss
+from omniloader.transforms.stats import compute_dataset_stats, compute_stats, save_stats
+
+if TYPE_CHECKING:
+    from collections.abc import Sequence
+
+
+def _load(config_path: str) -> tuple[list, list]:
+    """Load a config and build its declared datasets and schemas."""
+    config = OmniConfig.from_file(config_path)
+    datasets, schemas = config.build_datasets()
+    if not datasets:
+        raise SystemExit(f"Config {config_path!r} declares no datasets")
+    return datasets, schemas
+
+
+def _describe(args: argparse.Namespace) -> int:
+    """Print the dataset report."""
+    datasets, schemas = _load(args.config)
+    print(describe(datasets, schemas, max_samples=args.max_samples))
+    return 0
+
+
+def _validate(args: argparse.Namespace) -> int:
+    """Validate datasets against their specs; exit non-zero if issues are found."""
+    datasets, schemas = _load(args.config)
+    issues = validate(datasets, schemas, num_samples=args.num_samples)
+    if issues:
+        print("\n".join(issues))
+        return 1
+    print("OK: all datasets match their declared specs.")
+    return 0
+
+
+def _compute_stats(args: argparse.Namespace) -> int:
+    """Compute normalization statistics over the loader and save them to JSON.
+
+    With ``--per-dataset`` the stats are grouped by source dataset (for
+    :class:`~omniloader.transforms.normalize.PerDatasetNormalize`); otherwise they
+    are pooled over the union.
+    """
+    datasets, schemas = _load(args.config)
+    loader = OmniLoader(datasets, schemas)
+    keys = args.keys or loader.schema.feature_keys
+    samples = (loader[i] for i in range(len(loader)))
+    if args.per_dataset:
+        stats = compute_dataset_stats(samples, keys)
+        print(f"Saved per-dataset stats for {sorted(stats)} to {args.output}")
+    else:
+        stats = compute_stats(samples, keys)
+        print(f"Saved stats for {sorted(stats)} to {args.output}")
+    save_stats(stats, args.output)
+    return 0
+
+
+def _class_weights_for_loss(args: argparse.Namespace) -> int:
+    """Compute per-class counts and loss weights for a target and save them to JSON."""
+    datasets, schemas = _load(args.config)
+    loader = OmniLoader(datasets, schemas)
+    samples = [loader[i] for i in range(len(loader))]  # single pass, reused for both
+    hist = class_histogram(samples, args.target, num_classes=args.num_classes)
+    weights = class_weights_for_loss(
+        samples, args.target, num_classes=args.num_classes, scheme=args.scheme, beta=args.beta
+    )
+    payload = {
+        "target": args.target,
+        "scheme": args.scheme,
+        "num_classes": int(weights.numel()),
+        "counts": hist.tolist(),
+        "weights": weights.tolist(),
+    }
+    Path(args.output).write_text(json.dumps(payload, indent=2), encoding="utf-8")
+    print(
+        f"Saved class weights for {args.target!r} ({payload['num_classes']} classes) "
+        f"to {args.output}"
+    )
+    return 0
+
+
+def build_parser() -> argparse.ArgumentParser:
+    """Construct the argument parser for the ``omniloader`` CLI."""
+    parser = argparse.ArgumentParser(prog="omniloader", description=__doc__)
+    sub = parser.add_subparsers(dest="command", required=True)
+
+    p_desc = sub.add_parser("describe", help="print a dataset coverage/statistics report")
+    p_desc.add_argument("config", help="path to a JSON/YAML config")
+    p_desc.add_argument("--max-samples", type=int, default=64, dest="max_samples")
+    p_desc.set_defaults(func=_describe)
+
+    p_val = sub.add_parser("validate", help="check datasets against their declared specs")
+    p_val.add_argument("config", help="path to a JSON/YAML config")
+    p_val.add_argument("--num-samples", type=int, default=4, dest="num_samples")
+    p_val.set_defaults(func=_validate)
+
+    p_stats = sub.add_parser("compute-stats", help="compute and save normalization stats")
+    p_stats.add_argument("config", help="path to a JSON/YAML config")
+    p_stats.add_argument("-o", "--output", required=True, help="destination stats JSON")
+    p_stats.add_argument("--keys", nargs="*", help="feature keys (default: all features)")
+    p_stats.add_argument(
+        "--per-dataset",
+        action="store_true",
+        dest="per_dataset",
+        help="group stats by source dataset (for PerDatasetNormalize)",
+    )
+    p_stats.set_defaults(func=_compute_stats)
+
+    p_cw = sub.add_parser("class-weights-for-loss", help="compute and save per-class loss weights")
+    p_cw.add_argument("config", help="path to a JSON/YAML config")
+    p_cw.add_argument("--target", required=True, help="categorical target key")
+    p_cw.add_argument("-o", "--output", required=True, help="destination weights JSON")
+    p_cw.add_argument("--num-classes", type=int, default=None, dest="num_classes")
+    p_cw.add_argument("--scheme", choices=["inverse", "effective"], default="inverse")
+    p_cw.add_argument("--beta", type=float, default=0.999, help="beta for scheme=effective")
+    p_cw.set_defaults(func=_class_weights_for_loss)
+
+    return parser
+
+
+def main(argv: Sequence[str] | None = None) -> int:
+    """Entry point for the ``omniloader`` console script.
+
+    Args:
+        argv: Optional argument list (defaults to ``sys.argv``).
+
+    Returns:
+        The process exit code.
+
+    """
+    args = build_parser().parse_args(argv)
+    return int(args.func(args))
+
+
+if __name__ == "__main__":  # pragma: no cover
+    raise SystemExit(main())

omniloader/collate.py

@@ -0,0 +1,123 @@
+"""Collate functions for batching unified samples.
+
+Two collators are provided:
+
+* :func:`unified_collate` — the default: every sample already shares the same
+  tensor shapes (the unified schema, with sequences padded to a fixed ``time_dim``),
+  so it simply stacks tensors and gathers string metadata into lists.
+* :class:`DynamicCollator` — pads each sequence key only to the batch's longest
+  sample instead of a fixed ``time_dim``. Use it with
+  ``OmniLoader(pad_features=False)`` to avoid wasting memory/compute on padding
+  when sequence lengths vary a lot across the batch.
+
+Masks vs. cost: a ``<name>_mask`` guarantees *correctness* (padded steps never
+corrupt the model output or the loss), but a padded step still costs compute and
+memory (self-attention is ``O(T²)``). Dynamic padding targets that cost — so it is
+only useful in the native-length regime. Under the default fixed ``time_dim``
+(``pad_features=True``) every batch is already one length and :class:`DynamicCollator`
+is a no-op.
+"""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Any
+
+import torch
+
+from omniloader.utils.padding import pad_or_crop_time_dim
+
+if TYPE_CHECKING:
+    from collections.abc import Sequence
+
+    from omniloader.schema.spec import UnifiedSchema
+
+
+def unified_collate(batch: Sequence[dict[str, Any]]) -> dict[str, Any]:
+    """Collate a list of unified samples into a batched dict.
+
+    Args:
+        batch: Samples produced by :class:`~omniloader.loader.OmniLoader`.
+
+    Returns:
+        A dict where each tensor key holds a stacked tensor of shape
+        ``(B, ...)`` and each non-tensor key holds a list of length ``B``.
+
+    Raises:
+        ValueError: If the batch is empty.
+
+    """
+    if not batch:
+        raise ValueError("Cannot collate an empty batch")
+    keys = batch[0].keys()
+    out: dict[str, Any] = {}
+    for key in keys:
+        values = [sample[key] for sample in batch]
+        if isinstance(values[0], torch.Tensor):
+            out[key] = torch.stack(values, dim=0)
+        else:
+            out[key] = list(values)
+    return out
+
+
+class DynamicCollator:
+    """Collate with per-batch dynamic padding of sequence keys.
+
+    Sequence values (those whose spec sets ``time_dim``) are padded to the
+    longest sample in the batch — ``(B, L_batch, F)`` instead of a fixed
+    ``(B, time_dim, F)`` — and their ``<name>_mask`` is rebuilt accordingly.
+    Vectors, scalars and metadata are batched exactly as
+    :func:`unified_collate` does. Pair with ``OmniLoader(pad_features=False)`` so
+    samples reach the collator at their native length; under the default fixed
+    ``time_dim`` every sample already shares one length, so this collator has
+    nothing to shrink and behaves like :func:`unified_collate`.
+
+    Args:
+        schema: The unified schema, used to identify sequence keys and their
+            placeholder fill values.
+        keys: Optional subset of sequence keys to pad dynamically. When ``None``,
+            every sequence key in the schema is padded dynamically.
+
+    """
+
+    def __init__(self, schema: UnifiedSchema, keys: Sequence[str] | None = None) -> None:
+        seq = {spec.name: spec for spec in schema.specs if spec.is_sequence}
+        self.seq_specs = seq if keys is None else {k: seq[k] for k in keys if k in seq}
+
+    def __call__(self, batch: Sequence[dict[str, Any]]) -> dict[str, Any]:
+        """Collate ``batch``, dynamically padding the configured sequence keys.
+
+        Args:
+            batch: Samples to collate.
+
+        Returns:
+            The batched dict.
+
+        Raises:
+            ValueError: If the batch is empty.
+
+        """
+        if not batch:
+            raise ValueError("Cannot collate an empty batch")
+        out: dict[str, Any] = {}
+        for key in batch[0]:
+            # Sequence masks are produced alongside their value key; skip them.
+            if key.endswith("_mask") and key[:-5] in self.seq_specs:
+                continue
+            values = [sample[key] for sample in batch]
+            if key in self.seq_specs:
+                spec = self.seq_specs[key]
+                masks = [sample[f"{key}_mask"] for sample in batch]
+                length = max(v.shape[0] for v in values)  # batch max sequence length
+                # (T_i, ...) -> (L, ...) padded with the spec placeholder value.
+                out[key] = torch.stack(
+                    [pad_or_crop_time_dim(v, length, spec.placeholder)[0] for v in values]
+                )
+                # (T_i,) bool mask -> (L,) padded with False.
+                out[f"{key}_mask"] = torch.stack(
+                    [pad_or_crop_time_dim(m.to(torch.float32), length)[0].bool() for m in masks]
+                )
+            elif isinstance(values[0], torch.Tensor):
+                out[key] = torch.stack(values, dim=0)
+            else:
+                out[key] = list(values)
+        return out