"""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 dataset_registry import DATASET_REGISTRY, get_dataset_config from datasets import load_dataset from rich.console import Console from rich.progress import Progress, SpinnerColumn, TextColumn from rich.table import Table 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 " f"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 if isinstance(dataset, Dataset): self.console.print( f"\n[green]✓ Downloaded dataset with {len(dataset):,} " "rows[/green]" ) # type: ignore[arg-type] else: self.console.print( "\n[green]✓ Downloaded dataset" ) 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 " f"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))