124 lines
3.5 KiB
Python
124 lines
3.5 KiB
Python
"""Dataset validation utilities for HuggingFace uploads."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import importlib
|
|
from collections.abc import Callable
|
|
from pathlib import Path
|
|
|
|
from dataset_registry import get_dataset_config
|
|
from huggingface_hub import HfApi, snapshot_download
|
|
from huggingface_hub.errors import HfHubHTTPError
|
|
|
|
|
|
def resolve_organization(organization: str | None = None) -> str:
|
|
"""Resolve the HuggingFace organization name."""
|
|
if organization:
|
|
return organization
|
|
|
|
config = get_dataset_config()
|
|
return config.get("organization", "CleverThis")
|
|
|
|
|
|
def build_repo_id(dataset_id: str, organization: str | None = None) -> str:
|
|
"""Build the full HuggingFace repo ID for a dataset."""
|
|
org_name = resolve_organization(organization)
|
|
return f"{org_name}/{dataset_id}"
|
|
|
|
|
|
def dataset_exists_on_hf(
|
|
dataset_id: str,
|
|
organization: str | None = None,
|
|
api: HfApi | None = None,
|
|
) -> bool:
|
|
"""Check whether a dataset exists on HuggingFace Hub.
|
|
|
|
Args:
|
|
dataset_id: Dataset identifier (without organization prefix).
|
|
organization: HuggingFace organization name.
|
|
api: Optional HfApi instance for dependency injection.
|
|
|
|
Returns:
|
|
True if dataset exists, False if not found.
|
|
|
|
Raises:
|
|
HfHubHTTPError: For non-404 errors.
|
|
"""
|
|
repo_id = build_repo_id(dataset_id, organization)
|
|
client = api or HfApi()
|
|
|
|
try:
|
|
client.dataset_info(repo_id=repo_id)
|
|
except HfHubHTTPError as exc:
|
|
status_code = exc.response.status_code if exc.response else None
|
|
if status_code == 404:
|
|
return False
|
|
raise
|
|
|
|
return True
|
|
|
|
|
|
def dataset_can_download_from_hf(
|
|
dataset_id: str,
|
|
organization: str | None = None,
|
|
download_func: Callable[..., str] | None = None,
|
|
cache_dir: Path | str | None = None,
|
|
) -> bool:
|
|
"""Check whether a dataset can be downloaded from HuggingFace Hub.
|
|
|
|
Args:
|
|
dataset_id: Dataset identifier (without organization prefix).
|
|
organization: HuggingFace organization name.
|
|
download_func: Optional download function for dependency injection.
|
|
cache_dir: Optional cache directory for HuggingFace downloads.
|
|
|
|
Returns:
|
|
True if download succeeds, False if not found or download fails.
|
|
|
|
Raises:
|
|
HfHubHTTPError: For non-404 errors.
|
|
"""
|
|
repo_id = build_repo_id(dataset_id, organization)
|
|
downloader = download_func or snapshot_download
|
|
|
|
try:
|
|
downloader(
|
|
repo_id=repo_id,
|
|
repo_type="dataset",
|
|
local_dir=cache_dir,
|
|
)
|
|
except HfHubHTTPError as exc:
|
|
status_code = exc.response.status_code if exc.response else None
|
|
if status_code == 404:
|
|
return False
|
|
raise
|
|
except Exception:
|
|
return False
|
|
|
|
return True
|
|
|
|
|
|
def dataset_is_parquet_file(
|
|
file_path: Path | str,
|
|
read_parquet_func: Callable[[Path | str], object] | None = None,
|
|
) -> bool:
|
|
"""Check whether a file can be read by pandas.read_parquet().
|
|
|
|
Args:
|
|
file_path: Path to the file to validate.
|
|
read_parquet_func: Optional read_parquet function for dependency injection.
|
|
|
|
Returns:
|
|
True if the file can be read as parquet, False otherwise.
|
|
"""
|
|
if read_parquet_func is None:
|
|
pandas_module = importlib.import_module("pandas")
|
|
read_parquet_func = getattr(pandas_module, "read_parquet")
|
|
|
|
try:
|
|
read_parquet_func(file_path)
|
|
except Exception:
|
|
return False
|
|
|
|
return True
|