Pipelines

The three high-level entry points behind the ecgbench CLI. Each is a plain keyword-argument function, so everything the CLI does is importable:

from ecgbench import run_splits, run_croissant, run_upload

See the CLI page for the equivalent commands and their flags.

run_splits

splits

Full pipeline subcommand: validate + split + export + Croissant.

run_splits

run_splits(dataset: str, data_path: Path | str | None = None, output_dir: Path | str | None = None, sampling_rate: int | None = None, n_folds: int | None = None, max_workers: int = 4, skip_validation: bool = False, skip_croissant: bool = False) -> dict

Run the full pipeline: validate + split + export + Croissant.

n_folds defaults to the dataset's own config.n_folds, which is 10 for every dataset but szdb — see that field for why a small release cannot take ten. Passing an explicit value overrides it.

Returns the stats dict produced by export_splits plus output_dir, dataset and dataset_name keys for convenience.

Source code in ecgbench/cli/splits.py

def run_splits(
    dataset: str,
    data_path: Path | str | None = None,
    output_dir: Path | str | None = None,
    sampling_rate: int | None = None,
    n_folds: int | None = None,
    max_workers: int = 4,
    skip_validation: bool = False,
    skip_croissant: bool = False,
) -> dict:
    """Run the full pipeline: validate + split + export + Croissant.

    ``n_folds`` defaults to the dataset's own ``config.n_folds``, which is 10 for
    every dataset but ``szdb`` — see that field for why a small release cannot
    take ten. Passing an explicit value overrides it.

    Returns the stats dict produced by ``export_splits`` plus ``output_dir``,
    ``dataset`` and ``dataset_name`` keys for convenience.
    """
    from ecgbench.config import load_config
    from ecgbench.download import resolve_data_path
    from ecgbench.splitting import export_splits, get_splitter, split_dataset

    logger.info("Loading config for '%s'", dataset)
    config = load_config(dataset)
    if n_folds is None:
        n_folds = config.n_folds

    resolved_data_path = resolve_data_path(Path(data_path) if data_path else None, config)
    logger.info("Data path: %s", resolved_data_path)

    splitter = get_splitter(dataset)
    logger.info("Using splitter: %s", type(splitter).__name__)

    df = splitter.load_metadata(resolved_data_path, config)
    labels = splitter.get_stratification_labels(df, config)

    if not skip_validation:
        from ecgbench.validation import validate_dataset

        logger.info("Validating dataset...")
        val_result = validate_dataset(
            resolved_data_path,
            config,
            sampling_rate=sampling_rate,
            max_workers=max_workers,
        )
        logger.info(
            "Validation: %d total, %d valid, %d excluded",
            val_result.total_records,
            val_result.valid_records,
            val_result.excluded_records,
        )
    else:
        from ecgbench.validation.engine import ValidationResult

        original_df = df.copy()
        original_df["is_valid"] = True
        original_df["quality_issues"] = ""
        val_result = ValidationResult(
            original_df=original_df,
            clean_df=df.copy(),
            record_validations=[],
            summary={},
            total_records=len(df),
            valid_records=len(df),
            excluded_records=0,
        )
        logger.info("Validation skipped — all %d records marked as valid", len(df))

    logger.info("Splitting dataset into %d folds...", n_folds)
    split_result = split_dataset(df, labels, config, n_folds=n_folds)
    logger.info(
        "Split complete: %d folds, train=%s, val=%s, test=%s",
        split_result.n_folds,
        split_result.default_train_folds,
        split_result.default_val_folds,
        split_result.default_test_folds,
    )

    resolved_output = Path(output_dir) if output_dir else Path("output") / dataset
    logger.info("Exporting to %s", resolved_output)
    stats = export_splits(split_result, val_result, resolved_output, config)

    manifest = _write_manifest(config, resolved_data_path, resolved_output, split_result, n_folds)

    if not skip_croissant:
        try:
            from ecgbench.croissant import save_croissant

            for version in ("clean", "original"):
                version_dir = resolved_output / version
                if version_dir.exists():
                    croissant_path = version_dir / "croissant.json"
                    save_croissant(config, version_dir, croissant_path, version=version)
                    logger.info("Generated Croissant metadata: %s", croissant_path)
        except ImportError:
            logger.warning(
                "mlcroissant not installed — skipping Croissant generation. "
                "Install with: pip install ecgbench[croissant]"
            )

    return {
        "dataset": config.slug,
        "dataset_name": config.name,
        "output_dir": resolved_output,
        "manifest": manifest,
        **stats,
    }

run_croissant

croissant

Standalone Croissant metadata generation subcommand.

run_croissant

run_croissant(dataset: str, splits_dir: Path | str, output: Path | str | None = None, version: str = 'clean', validate: bool = False) -> Path

Generate (and optionally validate) Croissant 1.1 JSON-LD for a dataset.

Returns the path to the saved croissant.json file.

Source code in ecgbench/cli/croissant.py

def run_croissant(
    dataset: str,
    splits_dir: Path | str,
    output: Path | str | None = None,
    version: str = "clean",
    validate: bool = False,
) -> Path:
    """Generate (and optionally validate) Croissant 1.1 JSON-LD for a dataset.

    Returns the path to the saved ``croissant.json`` file.
    """
    if version not in ("clean", "original"):
        raise ValueError(f"version must be 'clean' or 'original', got {version!r}")

    from ecgbench.config import load_config
    from ecgbench.croissant import save_croissant, validate_croissant

    config = load_config(dataset)
    splits_path = Path(splits_dir)
    output_path = Path(output) if output else None

    logger.info("Generating Croissant metadata for '%s' (%s)", config.name, version)
    saved_path = save_croissant(config, splits_path, output_path, version=version)
    logger.info("Saved to %s", saved_path)

    if validate:
        logger.info("Validating...")
        is_valid, errors = validate_croissant(saved_path)
        if is_valid:
            logger.info("Validation passed")
        else:
            logger.error("Validation failed:")
            for err in errors:
                logger.error("  %s", err)
            raise RuntimeError(f"Croissant validation failed: {errors}")

    return saved_path

run_upload

run_upload refuses before any network call for a dataset whose config sets publish_fold_csvs: False.

upload

Upload fold CSVs and metadata to HuggingFace Hub.

run_upload

run_upload(data_dir: Path | str, datasets: list[str], hf_repo_id: str = 'vlbthambawita/ECGBench', dry_run: bool = False, token: str | None = None) -> dict[str, int]

Upload per-dataset fold CSVs and metadata to HuggingFace Hub.

Returns a mapping of dataset slug -> number of files uploaded (or that would have been uploaded, for dry_run=True).

Source code in ecgbench/cli/upload.py

def run_upload(
    data_dir: Path | str,
    datasets: list[str],
    hf_repo_id: str = "vlbthambawita/ECGBench",
    dry_run: bool = False,
    token: str | None = None,
) -> dict[str, int]:
    """Upload per-dataset fold CSVs and metadata to HuggingFace Hub.

    Returns a mapping of dataset slug -> number of files uploaded (or that would
    have been uploaded, for ``dry_run=True``).
    """
    from huggingface_hub import HfApi

    data_root = Path(data_dir)
    resolved_token = _resolve_hf_token(token)
    api = HfApi(token=resolved_token)

    uploaded: dict[str, int] = {}

    for dataset_slug in datasets:
        # Refuse before touching the network. Fold CSVs are identifiers only, but
        # for a credentialed or restricted source those identifiers are still
        # derived from data under a use agreement, and this repo is public and
        # ungated. Publication is effectively irreversible, so the config's
        # policy is enforced here rather than left to the operator to remember.
        try:
            from ecgbench.config import load_config

            dataset_config = load_config(dataset_slug)
        except Exception:  # unknown slug: fall through to the directory check
            dataset_config = None
        if dataset_config is not None and not dataset_config.publish_fold_csvs:
            raise PermissionError(
                f"Refusing to upload '{dataset_slug}': its config sets "
                f"publish_fold_csvs: false.\n{dataset_config.no_publish_reason.strip()}\n"
                "If this is genuinely intended, change the config and say why there."
            )

        dataset_dir = data_root / dataset_slug
        if not dataset_dir.exists():
            logger.warning("Directory not found: %s, skipping", dataset_dir)
            uploaded[dataset_slug] = 0
            continue

        logger.info("Processing %s...", dataset_slug)

        files_to_upload: list[tuple[Path, str]] = []
        for version in ("original", "clean"):
            version_dir = dataset_dir / version
            if not version_dir.exists():
                logger.warning("  %s/ not found, skipping", version)
                continue
            # export.py writes croissant.json inside each version directory, not
            # at the dataset root, so it has to be collected here.
            for pattern in ("*.csv", "croissant.json"):
                for path in sorted(version_dir.rglob(pattern)):
                    rel_path = path.relative_to(data_root)
                    files_to_upload.append((path, str(rel_path)))

        for extra_file in ("validation_report.json", "croissant.json"):
            extra_path = dataset_dir / extra_file
            if extra_path.exists():
                rel_path = extra_path.relative_to(data_root)
                files_to_upload.append((extra_path, str(rel_path)))

        if not files_to_upload:
            logger.warning("  No files found to upload for %s", dataset_slug)
            uploaded[dataset_slug] = 0
            continue

        logger.info("  Found %d files to upload", len(files_to_upload))

        if dry_run:
            for local_path, remote_path in files_to_upload:
                size_kb = local_path.stat().st_size / 1024
                logger.info("  [DRY RUN] %s (%.1f KB)", remote_path, size_kb)
            uploaded[dataset_slug] = len(files_to_upload)
            continue

        for local_path, remote_path in files_to_upload:
            logger.info("  Uploading %s", remote_path)
            api.upload_file(
                path_or_fileobj=str(local_path),
                path_in_repo=remote_path,
                repo_id=hf_repo_id,
                repo_type="dataset",
            )

        logger.info("  Done uploading %s (%d files)", dataset_slug, len(files_to_upload))
        uploaded[dataset_slug] = len(files_to_upload)

    logger.info("Upload complete!")
    return uploaded