dataset-uploader/scripts/download-dataset.py

#!/usr/bin/env python3
"""Standalone download-dataset command for downloading pre-processed RDF
datasets from HuggingFace Hub.

This script replaces the former cleverernie download-dataset CLI command.
It downloads pre-processed datasets that have been converted to an optimized
format for GISM training from HuggingFace Hub under the CleverThis organization.

USAGE:
    # List all available datasets
    python scripts/download-dataset.py --list

    # List datasets by category
    python scripts/download-dataset.py --list --category medium

    # Get information about a specific dataset
    python scripts/download-dataset.py --info yago-4.5

    # Download a pre-processed dataset from HuggingFace
    python scripts/download-dataset.py -d wordnet -o ./datasets

    # Download and cleanup compressed files
    python scripts/download-dataset.py -d dbpedia-core-en -o ./datasets --cleanup

    # Show already downloaded datasets
    python scripts/download-dataset.py --show-downloaded

AVAILABLE CATEGORIES:
    small   - Small datasets (< 1 GB) for testing
    medium  - Medium datasets (1-20 GB) for development
    large   - Large datasets (20-100 GB) for production
    xlarge  - Extra large datasets (> 100 GB) for research

EXAMPLES:
    # Quick start with WordNet
    python scripts/download-dataset.py -d wordnet
    python -m cleverernie rdf-train -d datasets/wordnet/wn31.ttl

    # Production training with YAGO
    python scripts/download-dataset.py -d yago-4.5 -o ./datasets
    python -m cleverernie rdf-train -d datasets/yago-4.5/yago-wd-facts.nt
"""

from __future__ import annotations

import sys
from typing import TYPE_CHECKING

import click

# Local imports from scripts directory
from dataset_downloader import DatasetDownloader
from dataset_registry import DATASET_REGISTRY, get_dataset_info
from rich.console import Console

if TYPE_CHECKING:
    pass


@click.command()
@click.option(
    "--dataset",
    "-d",
    type=str,
    help="Dataset ID to download (e.g., 'wordnet', 'yago-4.5')",
)
@click.option(
    "--output-dir",
    "-o",
    type=click.Path(),
    default="./datasets",
    help="Output directory for downloaded datasets (default: ./datasets)",
)
@click.option(
    "--list",
    "list_datasets_flag",
    is_flag=True,
    help="List all available datasets",
)
@click.option(
    "--info",
    type=str,
    help="Show detailed information about a specific dataset",
)
@click.option(
    "--category",
    type=click.Choice(["small", "medium", "large", "xlarge", "all"]),
    default="all",
    help="Filter datasets by category (default: all)",
)
@click.option(
    "--force",
    is_flag=True,
    help="Force re-download even if file exists",
)
@click.option(
    "--skip-extraction",
    is_flag=True,
    help="Skip extraction of compressed files",
)
@click.option(
    "--cleanup",
    is_flag=True,
    help="Remove compressed files after extraction",
)
@click.option(
    "--show-downloaded",
    is_flag=True,
    help="Show already downloaded datasets",
)
def main(
    dataset: str | None = None,
    output_dir: str = "./datasets",
    list_datasets_flag: bool = False,
    info: str | None = None,
    category: str = "all",  # Keep for CLI compatibility
    force: bool = False,  # Keep for CLI compatibility
    skip_extraction: bool = False,  # Keep for CLI compatibility
    cleanup: bool = False,  # Keep for CLI compatibility
    show_downloaded: bool = False,
) -> None:
    """Download pre-processed RDF datasets from HuggingFace Hub."""
    console = Console()

    # Initialize downloader
    downloader = DatasetDownloader(output_dir, console=console)

    # Handle --list flag
    if list_datasets_flag:
        downloader.list_available_datasets(show_unavailable=False)
        return

    # Handle --info flag
    if info:
        _show_dataset_info(console, info)
        return

    # Handle --show-downloaded flag
    if show_downloaded:
        downloaded = downloader.get_downloaded_datasets()
        if downloaded:
            console.print("\n[bold cyan]Downloaded Datasets:[/bold cyan]\n")
            for dataset_id in downloaded:
                ds_info = DATASET_REGISTRY.get(dataset_id)
                if ds_info:
                    console.print(f"  ✓ {dataset_id}: {ds_info.name}")
        else:
            console.print("\n[yellow]No datasets downloaded yet[/yellow]")
        return

    # Validate dataset argument
    if not dataset:
        console.print("[red]Error: No dataset specified[/red]")
        console.print("[yellow]Use --list to see available datasets[/yellow]")
        console.print(
            "[yellow]Use -d DATASET_ID to download a specific dataset[/yellow]"
        )
        sys.exit(1)

    # Download the dataset from HuggingFace
    console.print("\n[bold cyan]Dataset Download from HuggingFace Hub[/bold cyan]\n")

    result = downloader.download_dataset(
        dataset_id=dataset,
        split=None,  # Download all splits
    )

    if result:
        console.print("\n[bold green]✓ Dataset ready for training![/bold green]")
        console.print("\n[cyan]Next steps:[/cyan]")
        console.print("  1. Load the dataset in your training script:")
        console.print("     [dim]from datasets import load_dataset[/dim]")
        console.print(f"     [dim]ds = load_dataset('CleverThis/{dataset}')[/dim]")
        console.print("\n  2. Or use CleverErnie's training commands:")
        console.print(
            f"     [dim]python -m cleverernie rdf-train --dataset-id {dataset}[/dim]"
        )
    else:
        console.print("[red]Dataset download failed[/red]")
        sys.exit(1)


def _show_dataset_info(console: Console, dataset_id: str) -> None:
    """Display detailed information about a dataset.

    Args:
        console: Rich console for output
        dataset_id: Dataset identifier
    """
    dataset = get_dataset_info(dataset_id)

    if not dataset:
        console.print(f"[red]Error: Dataset '{dataset_id}' not found[/red]")
        console.print("[yellow]Use --list to see available datasets[/yellow]")
        sys.exit(1)

    console.print(f"\n[bold cyan]{dataset.name}[/bold cyan]\n")
    console.print(f"[bold]ID:[/bold] {dataset.id}")
    console.print(f"[bold]Description:[/bold] {dataset.description}")
    console.print()
    console.print(f"[bold]Category:[/bold] {dataset.category}")
    console.print(f"[bold]Format:[/bold] {dataset.format}")

    if dataset.compressed_size_gb:
        console.print(
            f"[bold]Download Size:[/bold] {dataset.compressed_size_gb:.1f} GB "
            f"(compressed)"
        )
        console.print(f"[bold]Extracted Size:[/bold] {dataset.size_gb:.1f} GB")
    else:
        console.print(f"[bold]Size:[/bold] ~{dataset.size_gb:.1f} GB")

    if dataset.entities:
        console.print(f"[bold]Entities:[/bold] {dataset.entities}")
    if dataset.triples:
        console.print(f"[bold]Triples:[/bold] {dataset.triples}")

    if dataset.license:
        console.print(f"[bold]License:[/bold] {dataset.license}")

    console.print()
    console.print("[bold]Recommended For:[/bold]")
    console.print(f"  {dataset.recommended_for}")

    if dataset.notes:
        console.print()
        console.print("[bold yellow]Notes:[/bold yellow]")
        console.print(f"  {dataset.notes}")

    console.print()
    console.print("[bold]Download URL:[/bold]")
    console.print(f"  {dataset.url}")

    console.print()
    console.print("[cyan]To download this dataset:[/cyan]")
    console.print(f"  [dim]python scripts/download-dataset.py -d {dataset.id}[/dim]")


if __name__ == "__main__":
    main()