235 lines
7.5 KiB
Python
Executable File
235 lines
7.5 KiB
Python
Executable File
#!/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 pathlib import Path
|
|
from typing import TYPE_CHECKING
|
|
|
|
import click
|
|
from rich.console import Console
|
|
from rich.table import Table
|
|
|
|
# Local imports from scripts directory
|
|
from dataset_downloader import DatasetDownloader
|
|
from dataset_registry import DATASET_REGISTRY, get_dataset_info, list_datasets
|
|
|
|
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", # noqa: ARG001 # Keep for CLI compatibility
|
|
force: bool = False, # noqa: ARG001 # Keep for CLI compatibility
|
|
skip_extraction: bool = False, # noqa: ARG001 # Keep for CLI compatibility
|
|
cleanup: bool = False, # noqa: ARG001 # 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 (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()
|