| from __future__ import annotations |
|
|
| from dataclasses import dataclass |
| from datetime import datetime |
| from pathlib import Path |
| from typing import Any |
|
|
| import pandas as pd |
| from huggingface_hub import HfApi, hf_hub_download |
| from huggingface_hub.utils import EntryNotFoundError, HfHubHTTPError |
|
|
| from config import AppConfig |
|
|
|
|
| README_CANDIDATES = ("README.md", "readme.md") |
|
|
|
|
| @dataclass(frozen=True) |
| class RepoRevision: |
| repo_id: str |
| revision: str |
| committed_at: datetime | None |
|
|
|
|
| class HubClient: |
| def __init__(self, config: AppConfig): |
| self.config = config |
| self.api = HfApi(token=config.hf_token) |
|
|
| def discover_dataset_repos(self) -> list[str]: |
| repos = set(self.config.repo_allowlist) |
|
|
| if self.config.discovery_enabled: |
| try: |
| datasets = self.api.list_datasets( |
| author=self.config.hf_owner, |
| search=self.config.repo_id_contains or None, |
| limit=max(self.config.max_repos * 3, self.config.max_repos), |
| full=False, |
| ) |
| for dataset in datasets: |
| repo_id = getattr(dataset, "id", None) |
| if not repo_id: |
| continue |
| if self._repo_matches_filters(repo_id): |
| repos.add(repo_id) |
| if len(repos) >= self.config.max_repos: |
| break |
| except HfHubHTTPError: |
| pass |
|
|
| filtered = [repo for repo in repos if self._repo_matches_filters(repo)] |
| return sorted(filtered)[: self.config.max_repos] |
|
|
| def list_dataset_commits(self, repo_id: str, max_commits: int) -> list[RepoRevision]: |
| revisions: list[RepoRevision] = [] |
| try: |
| commits = self.api.list_repo_commits(repo_id=repo_id, repo_type="dataset") |
| for commit in commits[:max_commits]: |
| revision = getattr(commit, "commit_id", None) or getattr(commit, "oid", None) |
| if not revision: |
| continue |
| committed_at = _coerce_datetime(getattr(commit, "created_at", None) or getattr(commit, "date", None)) |
| revisions.append( |
| RepoRevision( |
| repo_id=repo_id, |
| revision=revision, |
| committed_at=committed_at, |
| ) |
| ) |
| except HfHubHTTPError: |
| revisions = [] |
|
|
| if revisions: |
| return revisions |
|
|
| try: |
| info = self.api.dataset_info(repo_id=repo_id) |
| fallback_revision = getattr(info, "sha", None) or "main" |
| fallback_timestamp = _coerce_datetime(getattr(info, "lastModified", None)) |
| except HfHubHTTPError: |
| fallback_revision = "main" |
| fallback_timestamp = None |
|
|
| return [RepoRevision(repo_id=repo_id, revision=fallback_revision, committed_at=fallback_timestamp)] |
|
|
| def fetch_dataset_readme(self, repo_id: str, revision: str | None = None) -> str | None: |
| for filename in README_CANDIDATES: |
| try: |
| path = hf_hub_download( |
| repo_id=repo_id, |
| filename=filename, |
| repo_type="dataset", |
| revision=revision, |
| token=self.config.hf_token, |
| ) |
| return Path(path).read_text(encoding="utf-8", errors="replace") |
| except EntryNotFoundError: |
| continue |
| except HfHubHTTPError: |
| continue |
| return None |
|
|
| def _repo_matches_filters(self, repo_id: str) -> bool: |
| lowered_repo_id = repo_id.lower() |
|
|
| if self.config.repo_id_contains and self.config.repo_id_contains.lower() not in lowered_repo_id: |
| return False |
|
|
| if self.config.repo_prefixes and not any(repo_id.startswith(prefix) for prefix in self.config.repo_prefixes): |
| return False |
|
|
| return True |
|
|
|
|
| def _coerce_datetime(value: Any) -> datetime | None: |
| if value is None: |
| return None |
|
|
| timestamp = pd.to_datetime(value, errors="coerce", utc=True) |
| if pd.isna(timestamp): |
| return None |
| return timestamp.to_pydatetime() |
|
|
