142 lines
3.8 KiB
Python
142 lines
3.8 KiB
Python
"""Dataset validation utilities for HuggingFace uploads."""
|
|
|
|
from __future__ import annotations
|
|
|
|
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
|
|
from pandas import read_parquet
|
|
|
|
|
|
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,
|
|
) -> str | None:
|
|
"""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:
|
|
Download path if download succeeds, None 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:
|
|
download_path = 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 None
|
|
raise
|
|
except Exception:
|
|
return None
|
|
|
|
return download_path
|
|
|
|
|
|
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.
|
|
"""
|
|
reader = read_parquet_func or read_parquet
|
|
|
|
try:
|
|
reader(file_path)
|
|
except Exception:
|
|
return False
|
|
|
|
return True
|
|
|
|
REQUIRED_COLUMNS = (
|
|
"subject",
|
|
"predicate",
|
|
"object",
|
|
"object_type",
|
|
"object_datatype",
|
|
"object_language",
|
|
)
|
|
|
|
def dataset_has_required_columns(file_path: Path | str) -> bool:
|
|
"""Check whether a parquet file has the required RDF columns."""
|
|
try:
|
|
dataframe = read_parquet(file_path)
|
|
except Exception:
|
|
return False
|
|
|
|
columns = set(dataframe.columns)
|
|
return all(column in columns for column in REQUIRED_COLUMNS)
|
|
|