210 lines
8.7 KiB
Python
210 lines
8.7 KiB
Python
"""Dataset downloader for pre-processed RDF knowledge graphs from HuggingFace Hub.
|
|
|
|
This module provides functionality to download pre-processed RDF datasets
|
|
in HuggingFace format for GISM training. These datasets have been converted
|
|
from raw RDF to an optimized format suitable for training.
|
|
|
|
For downloading raw RDF datasets from their original sources (for dataset
|
|
preparation/conversion), see scripts/rdf_dataset_downloader.py instead.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
from pathlib import Path
|
|
from typing import TYPE_CHECKING
|
|
|
|
from datasets import load_dataset
|
|
from rich.console import Console
|
|
from rich.progress import Progress, SpinnerColumn, TextColumn
|
|
from rich.table import Table
|
|
|
|
from dataset_registry import DATASET_REGISTRY, get_dataset_config
|
|
|
|
if TYPE_CHECKING:
|
|
from datasets import Dataset, DatasetDict
|
|
|
|
|
|
class DatasetDownloader:
|
|
"""Handles downloading pre-processed datasets from HuggingFace Hub."""
|
|
|
|
def __init__(
|
|
self,
|
|
output_dir: Path | str,
|
|
organization: str | None = None,
|
|
console: Console | None = None,
|
|
):
|
|
"""Initialize the dataset downloader.
|
|
|
|
Args:
|
|
output_dir: Directory to store downloaded datasets
|
|
organization: HuggingFace organization name (defaults to config value)
|
|
console: Rich console for output (created if not provided)
|
|
"""
|
|
self.output_dir = Path(output_dir)
|
|
self.output_dir.mkdir(parents=True, exist_ok=True)
|
|
|
|
# Use organization from config if not provided
|
|
if organization is None:
|
|
config = get_dataset_config()
|
|
organization = config.get("organization", "CleverThis")
|
|
|
|
self.organization = organization
|
|
self.console = console or Console()
|
|
|
|
def download_dataset(
|
|
self,
|
|
dataset_id: str,
|
|
split: str | None = None,
|
|
cache_dir: str | None = None,
|
|
token: str | None = None,
|
|
) -> Dataset | DatasetDict | None: # type: ignore[return-value]
|
|
"""Download a pre-processed dataset from HuggingFace Hub.
|
|
|
|
Args:
|
|
dataset_id: Dataset identifier (must exist in registry)
|
|
split: Specific split to load ('train', 'validation', 'test', 'data')
|
|
If None, returns all splits as DatasetDict
|
|
cache_dir: Custom cache directory for HuggingFace datasets
|
|
token: HuggingFace API token (for private datasets)
|
|
|
|
Returns:
|
|
Dataset or DatasetDict object, or None on failure
|
|
"""
|
|
# Verify dataset exists in registry
|
|
dataset_info = DATASET_REGISTRY.get(dataset_id)
|
|
if not dataset_info:
|
|
self.console.print(f"[red]Error: Dataset '{dataset_id}' not found in registry[/red]")
|
|
self.console.print("[yellow]Use 'cleverernie list-datasets' to see available datasets[/yellow]")
|
|
return None
|
|
|
|
# Check if dataset is available
|
|
if not dataset_info.available:
|
|
self.console.print(f"[yellow]Warning: Dataset '{dataset_id}' is marked as unavailable[/yellow]")
|
|
self.console.print("[yellow]This dataset may not have been uploaded to HuggingFace yet.[/yellow]")
|
|
return None
|
|
|
|
# Check if dataset has custom HuggingFace repo
|
|
if hasattr(dataset_info, 'huggingface_repo') and dataset_info.huggingface_repo:
|
|
hf_dataset_name = dataset_info.huggingface_repo
|
|
else:
|
|
# Construct HuggingFace dataset name using organization
|
|
hf_dataset_name = f"{self.organization}/{dataset_id}"
|
|
|
|
# Display dataset info
|
|
self.console.print(f"\n[bold cyan]Downloading: {dataset_info.name}[/bold cyan]")
|
|
self.console.print(f"HuggingFace Dataset: {hf_dataset_name}")
|
|
self.console.print(f"Description: {dataset_info.description}")
|
|
self.console.print(f"License: {dataset_info.license or 'See dataset card'}")
|
|
|
|
# Download dataset with progress
|
|
try:
|
|
with Progress(
|
|
SpinnerColumn(),
|
|
TextColumn("[progress.description]{task.description}"),
|
|
console=self.console,
|
|
) as progress:
|
|
progress.add_task(
|
|
f"Downloading {hf_dataset_name}...",
|
|
total=None,
|
|
)
|
|
|
|
dataset = load_dataset(
|
|
hf_dataset_name,
|
|
split=split,
|
|
cache_dir=cache_dir or str(self.output_dir / ".cache"),
|
|
token=token,
|
|
)
|
|
|
|
# Display success message
|
|
if isinstance(dataset, dict): # DatasetDict
|
|
self.console.print(f"\n[green]✓ Downloaded {len(dataset)} splits[/green]")
|
|
for split_name, ds in dataset.items():
|
|
self.console.print(f" • {split_name}: {len(ds):,} rows") # type: ignore[arg-type]
|
|
else: # Single Dataset
|
|
self.console.print(f"\n[green]✓ Downloaded dataset with {len(dataset):,} rows[/green]") # type: ignore[arg-type]
|
|
|
|
self.console.print(f"\n[cyan]Dataset cached at: {self.output_dir}[/cyan]")
|
|
return dataset # type: ignore[return-value]
|
|
|
|
except Exception as e:
|
|
error_msg = str(e)
|
|
self.console.print(f"[red]Error downloading dataset: {error_msg}[/red]")
|
|
|
|
# Provide helpful context
|
|
if "404" in error_msg or "not found" in error_msg.lower():
|
|
self.console.print(f"[yellow]→ Dataset '{hf_dataset_name}' not found on HuggingFace Hub[/yellow]")
|
|
self.console.print("[yellow] • The dataset may not have been uploaded yet[/yellow]")
|
|
self.console.print(f"[yellow] • Check https://huggingface.co/datasets/{hf_dataset_name}[/yellow]")
|
|
self.console.print("[yellow] • Use scripts/upload_all_datasets.py to upload datasets[/yellow]")
|
|
elif "authentication" in error_msg.lower() or "401" in error_msg:
|
|
self.console.print("[yellow]→ Authentication required[/yellow]")
|
|
self.console.print("[yellow] • Provide a HuggingFace token with --token[/yellow]")
|
|
self.console.print("[yellow] • Or run: huggingface-cli login[/yellow]")
|
|
elif "connection" in error_msg.lower() or "network" in error_msg.lower():
|
|
self.console.print("[yellow]→ Network connection error[/yellow]")
|
|
self.console.print("[yellow] • Check your internet connection[/yellow]")
|
|
self.console.print("[yellow] • Try again in a moment[/yellow]")
|
|
|
|
return None
|
|
|
|
def list_available_datasets(self, show_unavailable: bool = False) -> None:
|
|
"""List datasets available in the registry with a nice colored table.
|
|
|
|
Args:
|
|
show_unavailable: Include datasets marked as unavailable
|
|
"""
|
|
self.console.print("\n[bold cyan]Available Datasets:[/bold cyan]\n")
|
|
|
|
# Create table
|
|
table = Table(show_header=True)
|
|
table.add_column("ID", style="cyan")
|
|
table.add_column("Name", style="green")
|
|
table.add_column("Category", style="yellow")
|
|
table.add_column("Size", style="blue")
|
|
|
|
# Collect all datasets (optionally filtered)
|
|
datasets_to_show = []
|
|
for dataset_id, info in DATASET_REGISTRY.items():
|
|
if not show_unavailable and not info.available:
|
|
continue
|
|
datasets_to_show.append((dataset_id, info))
|
|
|
|
# Sort by category, then by ID
|
|
category_order = {"small": 0, "medium": 1, "large": 2, "xlarge": 3}
|
|
datasets_to_show.sort(key=lambda x: (category_order.get(x[1].category, 99), x[0]))
|
|
|
|
# Add rows to table
|
|
for dataset_id, info in datasets_to_show:
|
|
table.add_row(dataset_id, info.name, info.category, f"{info.size_gb} GB")
|
|
|
|
self.console.print(table)
|
|
|
|
# Count statistics
|
|
self.console.print(f"\n[bold]Total: {len(datasets_to_show)} datasets[/bold]")
|
|
self.console.print() # Extra blank line
|
|
|
|
def get_downloaded_datasets(self) -> list[str]:
|
|
"""Get list of datasets cached locally.
|
|
|
|
Returns:
|
|
List of dataset IDs that have been downloaded
|
|
"""
|
|
cached = []
|
|
cache_dir = self.output_dir / ".cache"
|
|
|
|
if not cache_dir.exists():
|
|
return cached
|
|
|
|
# HuggingFace datasets cache structure
|
|
for org_dir in cache_dir.iterdir():
|
|
if not org_dir.is_dir():
|
|
continue
|
|
for dataset_dir in org_dir.iterdir():
|
|
if not dataset_dir.is_dir():
|
|
continue
|
|
dataset_id = dataset_dir.name
|
|
if dataset_id in DATASET_REGISTRY:
|
|
cached.append(dataset_id)
|
|
|
|
return sorted(set(cached))
|