164 lines
4.9 KiB
Python
164 lines
4.9 KiB
Python
"""Dataset registry for RDF knowledge graphs.
|
|
|
|
This module provides a comprehensive registry of public RDF datasets
|
|
that can be downloaded and used for GISM training. The registry is loaded
|
|
from a JSON configuration file for easy maintenance.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import json
|
|
from dataclasses import dataclass
|
|
from pathlib import Path
|
|
from typing import TYPE_CHECKING
|
|
|
|
if TYPE_CHECKING:
|
|
pass
|
|
|
|
|
|
@dataclass
|
|
class DatasetInfo:
|
|
"""Information about a downloadable RDF dataset."""
|
|
|
|
id: str
|
|
name: str
|
|
description: str
|
|
url: str
|
|
format: str
|
|
size_gb: float
|
|
compressed_size_gb: float | None
|
|
entities: str | None
|
|
triples: str | None
|
|
category: str
|
|
license: str | None
|
|
recommended_for: str
|
|
extraction_cmd: str | None = None
|
|
notes: str | None = None
|
|
available: bool = True # Flag to mark if dataset is currently accessible
|
|
huggingface_repo: str | None = None # HuggingFace repository name
|
|
short_name: str | None = None # Short name for the dataset
|
|
|
|
|
|
def load_dataset_config() -> dict:
|
|
"""Load dataset configuration from JSON file.
|
|
|
|
Returns:
|
|
Dictionary containing dataset configuration
|
|
"""
|
|
config_path = Path(__file__).parent / "dataset_config.json"
|
|
if not config_path.exists():
|
|
raise FileNotFoundError(f"Dataset configuration file not found: {config_path}")
|
|
|
|
with open(config_path, "r") as f:
|
|
return json.load(f)
|
|
|
|
|
|
def build_dataset_registry() -> dict[str, DatasetInfo]:
|
|
"""Build dataset registry from JSON configuration.
|
|
|
|
Returns:
|
|
Dictionary mapping dataset IDs to DatasetInfo objects
|
|
"""
|
|
config = load_dataset_config()
|
|
registry = {}
|
|
|
|
for dataset_id, dataset_config in config["datasets"].items():
|
|
# Extract download and metadata information
|
|
download = dataset_config["download"]
|
|
metadata = dataset_config["metadata"]
|
|
huggingface = dataset_config.get("huggingface", {})
|
|
|
|
# Handle unavailable datasets
|
|
available = True
|
|
if metadata.get("notes"):
|
|
if any(phrase in metadata["notes"].lower() for phrase in [
|
|
"requires registration",
|
|
"multi-file ftp",
|
|
"hdt format",
|
|
"multiple files on figshare",
|
|
"databus collection"
|
|
]):
|
|
available = False
|
|
|
|
# Special cases for specific datasets
|
|
if dataset_id in ["babelnet", "pubchem", "bio2rdf", "lod-a-lot",
|
|
"opencitations", "dbpedia-latest"]:
|
|
available = False
|
|
elif dataset_id in ["conceptnet", "framenet", "opencyc"]:
|
|
available = True
|
|
|
|
registry[dataset_id] = DatasetInfo(
|
|
id=dataset_id,
|
|
name=dataset_config["full_name"],
|
|
description=dataset_config["description"],
|
|
url=download["url"],
|
|
format=download["format"],
|
|
size_gb=download["size_gb"],
|
|
compressed_size_gb=download.get("compressed_size_gb"),
|
|
entities=metadata.get("entities"),
|
|
triples=metadata.get("triples"),
|
|
category=dataset_config["category"],
|
|
license=metadata.get("license"),
|
|
recommended_for=metadata.get("recommended_for", ""),
|
|
extraction_cmd=download.get("extraction_command"),
|
|
notes=metadata.get("notes"),
|
|
available=available,
|
|
huggingface_repo=huggingface.get("repository"),
|
|
short_name=dataset_config.get("short_name")
|
|
)
|
|
|
|
return registry
|
|
|
|
|
|
# Build the registry on module import
|
|
try:
|
|
DATASET_REGISTRY = build_dataset_registry()
|
|
except FileNotFoundError:
|
|
# Fallback to empty registry if config file is missing
|
|
print("Warning: dataset_config.json not found, using empty registry")
|
|
DATASET_REGISTRY = {}
|
|
|
|
|
|
def get_dataset_info(dataset_id: str) -> DatasetInfo | None:
|
|
"""Get information about a dataset by ID.
|
|
|
|
Args:
|
|
dataset_id: Dataset identifier
|
|
|
|
Returns:
|
|
DatasetInfo object or None if not found
|
|
"""
|
|
return DATASET_REGISTRY.get(dataset_id)
|
|
|
|
|
|
def list_datasets(category: str | None = None) -> list[DatasetInfo]:
|
|
"""List all available datasets, optionally filtered by category.
|
|
|
|
Args:
|
|
category: Filter by category (small, medium, large, xlarge)
|
|
|
|
Returns:
|
|
List of DatasetInfo objects
|
|
"""
|
|
datasets = list(DATASET_REGISTRY.values())
|
|
if category:
|
|
datasets = [d for d in datasets if d.category == category]
|
|
return datasets
|
|
|
|
|
|
def get_dataset_categories() -> list[str]:
|
|
"""Get list of all dataset categories.
|
|
|
|
Returns:
|
|
List of category names
|
|
"""
|
|
return sorted({d.category for d in DATASET_REGISTRY.values()})
|
|
|
|
|
|
def get_dataset_config() -> dict:
|
|
"""Get the full dataset configuration.
|
|
|
|
Returns:
|
|
Dictionary containing the full dataset configuration
|
|
"""
|
|
return load_dataset_config() |