cleverdatasets/dataset-uploader
0
0
Fork 0
Code Issues 31 Pull Requests 7 Actions 12 Packages Projects Releases Wiki Activity
Files
master
dataset-uploader/scripts/upload_all_datasets.py
T

2029 lines
90 KiB
Python
Executable File
Raw Permalink Blame History

#!/usr/bin/env python3
"""Download, convert, and upload all RDF datasets to HuggingFace Hub.
This script orchestrates the complete dataset preparation pipeline for RDF
knowledge graphs, automating download, conversion, and upload to HuggingFace Hub.
Part of the CleverErnie dataset preparation infrastructure.
PARALLEL PROCESSING:
The script supports parallel processing of multiple datasets to speed up
bulk operations. When using --parallel/-p with a value > 1, datasets are
processed concurrently with rich progress display showing all datasets
being processed simultaneously.
MEMORY OPTIMIZATION:
The script automatically uses a streaming converter for medium and larger datasets
(or datasets identified as memory-intensive like GeoNames) to prevent excessive
memory usage. The streaming converter:
- Processes RDF files in chunks instead of loading entirely into memory
- Uses parallel processing for suitable formats
- Writes data incrementally to Parquet files
- Typical memory usage: 2-4GB instead of 100GB+ for large datasets
USAGE:
# Process a single dataset
python scripts/upload_all_datasets.py --dataset <dataset_id>
# Process multiple specific datasets in parallel
python scripts/upload_all_datasets.py --dataset wordnet --dataset yago-4.5 --parallel 2
# Process all datasets in a category with 4 parallel workers
python scripts/upload_all_datasets.py --category small --parallel 4
# Dry run (preview without executing)
python scripts/upload_all_datasets.py --category small --dry-run
# List datasets that would be processed
python scripts/upload_all_datasets.py --category small --list
# Skip certain steps (use existing downloads/conversions)
python scripts/upload_all_datasets.py --dataset wordnet --skip-download
python scripts/upload_all_datasets.py --dataset wordnet --skip-convert
python scripts/upload_all_datasets.py --dataset wordnet --skip-upload
# Remove mode: delete downloads after processing (save disk space)
python scripts/upload_all_datasets.py --category small --rm
ARGUMENTS:
-d, --dataset ID Specific dataset ID(s) to process (repeatable)
-c, --category CAT Process datasets by category: small, medium, large, xlarge, all
-p, --parallel N Number of datasets to process in parallel (default: 1)
--base-dir DIR Base directory for operations (default: ./dataset_processing)
--skip-download Skip download step (use existing downloads)
--skip-convert Skip conversion step (use existing conversions)
--skip-upload Skip upload step (only download and convert)
--dry-run Show what would be done without executing
--list List datasets that would be processed
--rm Remove mode: wipe downloads at start, delete each after processing
WORKFLOW:
For each dataset, the script executes three phases:
Step 1: Download
- Downloads raw RDF from public source URL
- Handles compressed formats (.gz, .bz2, .zip)
- Extracts to usable RDF format
- Progress bars with download speed
Step 2: Convert
- Parses RDF file (Turtle, N-Triples, RDF/XML, etc.)
- Converts to HuggingFace dataset format
- Preserves all semantic information (lossless)
- Creates dataset card (README.md) with metadata
Step 3: Upload
- Uploads to HuggingFace Hub under CleverThis organization
- Creates/updates dataset repository
- Uploads dataset card with comprehensive documentation
- Makes dataset publicly accessible
DATASET CATEGORIES:
small < 1 GB (WordNet, Schema.org, Gene Ontology)
medium 1-10 GB (YAGO 4.5, DisGeNET, GeoNames)
large 10-100 GB (DBpedia, UniProt, DBLP)
xlarge > 100 GB (Wikidata, BabelNet, Bio2RDF)
DIRECTORY STRUCTURE:
dataset_processing/
├── downloads/ # Raw RDF downloads
│ ├── wordnet/
│ │ ├── english-wordnet-2024.ttl.gz
│ │ └── english-wordnet-2024.ttl
│ └── yago-4.5/
│ └── yago-4.5.0.2-tiny.ttl
│
└── hf_datasets/ # Converted HuggingFace datasets
├── wordnet/
│ ├── data/
│ ├── dataset_info.json
│ └── README.md
└── yago-4.5/
├── data/
├── dataset_info.json
└── README.md
REMOVE MODE (--rm):
When --rm flag is set:
1. Start: Removes all previous downloads from dataset_processing/downloads/
2. Per dataset: After successful processing, removes downloaded RDF files
3. Result: Only converted HuggingFace datasets remain on disk
Use cases:
- Processing large datasets with limited disk space
- Automated pipelines where raw downloads aren't needed
- Batch processing where intermediate files can be discarded
EXAMPLES:
# Test with small dataset
python scripts/upload_all_datasets.py --dataset wordnet
# Process all small datasets with cleanup
python scripts/upload_all_datasets.py --category small --rm
# Preview what would happen
python scripts/upload_all_datasets.py --category medium --dry-run
# Re-upload with updated metadata (skip download/convert)
python scripts/upload_all_datasets.py --dataset wordnet --skip-download --skip-convert
# Process specific datasets from different categories
python scripts/upload_all_datasets.py --dataset wordnet --dataset yago-4.5 --dataset dbpedia-core
AUTHENTICATION:
To upload to HuggingFace Hub, authentication is required:
# Option 1: Login via CLI (one-time)
huggingface-cli login
# Option 2: Set environment variable
export HF_TOKEN="your_token_here"
Get a token from: https://huggingface.co/settings/tokens
Required permissions: Write access to create/update datasets
REQUIREMENTS:
pip install rdflib datasets huggingface-hub rich httpx pyarrow
SYSTEM REQUIREMENTS:
Small datasets: 4 GB RAM, 5 GB disk, 5-15 minutes
Medium datasets: 16 GB RAM, 30 GB disk, 30-90 minutes
Large datasets: 32 GB RAM, 300 GB disk, 2-6 hours
XLarge datasets: 64 GB RAM, 1+ TB disk, 6-24+ hours
NOTES:
- Datasets marked as unavailable will be automatically skipped
- Processing summary provided at end with success/failure status
- Upload resumes if interrupted (HuggingFace Hub handles this)
- Large datasets may take hours; monitor progress with -v flag
RELATED:
- rdf_dataset_downloader.py: Download raw RDF datasets (Step 1)
- convert_rdf_to_hf_dataset.py: Convert RDF to HuggingFace (Step 2)
- dataset_registry.py: Registry of available datasets with metadata
For detailed documentation, see: docs/rdf_dataset_scripts.md
"""
from __future__ import annotations
import argparse
import bz2
import shutil
import subprocess
import sys
import threading
from concurrent.futures import ProcessPoolExecutor, ThreadPoolExecutor, as_completed
from datetime import datetime
from pathlib import Path
from typing import Dict, Tuple
from rich.console import Console
from rich.live import Live
from rich.progress import (
BarColumn,
MofNCompleteColumn,
Progress,
SpinnerColumn,
TextColumn,
TimeElapsedColumn,
)
from rich.table import Table
# Import our modules
from dataset_registry import DATASET_REGISTRY, DatasetInfo, get_dataset_config
console = Console()
def get_valid_hf_license(license_string: str | None) -> str:
"""Map dataset license to valid HuggingFace license identifier.
Args:
license_string: License string from dataset registry
Returns:
Valid HuggingFace license identifier
"""
if not license_string:
return "other"
# Get license mappings from config
config = get_dataset_config()
license_map = config.get("license_mappings", {})
# Convert keys to lowercase for comparison
license_map_lower = {k.lower(): v for k, v in license_map.items()}
normalized = license_string.lower().strip()
return license_map_lower.get(normalized, "other")
def get_rdf_format(dataset_info: DatasetInfo) -> str:
"""Get RDF format string for rdflib.
Args:
dataset_info: Dataset information
Returns:
Format string for rdflib
"""
# Get format mappings from config
config = get_dataset_config()
format_map = config.get("format_mappings", {})
return format_map.get(dataset_info.format.lower(), "turtle")
def get_rdf_format_documentation(dataset_info: DatasetInfo) -> str:
"""Get comprehensive RDF format documentation for README.
Args:
dataset_info: Dataset information
Returns:
Markdown documentation for RDF format
"""
return f"""
## Dataset Format: Lossless RDF Representation
This dataset uses a **standard lossless format** for representing RDF (Resource Description Framework)
data in HuggingFace Datasets. All semantic information from the original RDF knowledge graph is preserved,
enabling perfect round-trip conversion between RDF and HuggingFace formats.
### Schema
Each RDF triple is represented as a row with **6 fields**:
| Field | Type | Description | Example |
|-------|------|-------------|---------|
| `subject` | string | Subject of the triple (URI or blank node) | `"http://schema.org/Person"` |
| `predicate` | string | Predicate URI | `"http://www.w3.org/1999/02/22-rdf-syntax-ns#type"` |
| `object` | string | Object of the triple | `"John Doe"` or `"http://schema.org/Thing"` |
| `object_type` | string | Type of object: `"uri"`, `"literal"`, or `"blank_node"` | `"literal"` |
| `object_datatype` | string | XSD datatype URI (for typed literals) | `"http://www.w3.org/2001/XMLSchema#integer"` |
| `object_language` | string | Language tag (for language-tagged literals) | `"en"` |
### Example: RDF Triple Representation
**Original RDF (Turtle)**:
```turtle
<http://example.org/John> <http://schema.org/name> "John Doe"@en .
```
**HuggingFace Dataset Row**:
```python
{{
"subject": "http://example.org/John",
"predicate": "http://schema.org/name",
"object": "John Doe",
"object_type": "literal",
"object_datatype": None,
"object_language": "en"
}}
```
### Loading the Dataset
```python
from datasets import load_dataset
# Load the dataset
dataset = load_dataset("CleverThis/{dataset_info.id}")
# Access the data
data = dataset["data"]
# Iterate over triples
for row in data:
subject = row["subject"]
predicate = row["predicate"]
obj = row["object"]
obj_type = row["object_type"]
print(f"Triple: ({{subject}}, {{predicate}}, {{obj}})")
print(f" Object type: {{obj_type}}")
if row["object_language"]:
print(f" Language: {{row['object_language']}}")
if row["object_datatype"]:
print(f" Datatype: {{row['object_datatype']}}")
```
### Converting Back to RDF
The dataset can be converted back to any RDF format (Turtle, N-Triples, RDF/XML, etc.) with **zero information loss**:
```python
from datasets import load_dataset
from rdflib import Graph, URIRef, Literal, BNode
def convert_to_rdf(dataset_name, output_file="output.ttl", split="data"):
\"\"\"Convert HuggingFace dataset back to RDF Turtle format.\"\"\"
# Load dataset
dataset = load_dataset(dataset_name)
# Create RDF graph
graph = Graph()
# Convert each row to RDF triple
for row in dataset[split]:
# Subject
if row["subject"].startswith("_:"):
subject = BNode(row["subject"][2:])
else:
subject = URIRef(row["subject"])
# Predicate (always URI)
predicate = URIRef(row["predicate"])
# Object (depends on object_type)
if row["object_type"] == "uri":
obj = URIRef(row["object"])
elif row["object_type"] == "blank_node":
obj = BNode(row["object"][2:])
elif row["object_type"] == "literal":
if row["object_datatype"]:
obj = Literal(row["object"], datatype=URIRef(row["object_datatype"]))
elif row["object_language"]:
obj = Literal(row["object"], lang=row["object_language"])
else:
obj = Literal(row["object"])
graph.add((subject, predicate, obj))
# Serialize to Turtle (or any RDF format)
graph.serialize(output_file, format="turtle")
print(f"Exported {{len(graph)}} triples to {{output_file}}")
return graph
# Usage
graph = convert_to_rdf("CleverThis/{dataset_info.id}", "reconstructed.ttl")
```
### Information Preservation Guarantee
This format preserves **100% of RDF information**:
- ✅ **URIs**: Exact string representation preserved
- ✅ **Literals**: Full text content preserved
- ✅ **Datatypes**: XSD and custom datatypes preserved (e.g., `xsd:integer`, `xsd:dateTime`)
- ✅ **Language Tags**: BCP 47 language tags preserved (e.g., `@en`, `@fr`, `@ja`)
- ✅ **Blank Nodes**: Node structure preserved (identifiers may change but graph isomorphism maintained)
**Round-trip guarantee**: Original RDF → HuggingFace → Reconstructed RDF produces **semantically identical** graphs.
### Querying the Dataset
You can filter and query the dataset like any HuggingFace dataset:
```python
from datasets import load_dataset
dataset = load_dataset("CleverThis/{dataset_info.id}")
# Find all triples with English literals
english_literals = dataset["data"].filter(
lambda x: x["object_type"] == "literal" and x["object_language"] == "en"
)
print(f"Found {{len(english_literals)}} English literals")
# Find all rdf:type statements
type_statements = dataset["data"].filter(
lambda x: "rdf-syntax-ns#type" in x["predicate"]
)
print(f"Found {{len(type_statements)}} type statements")
# Convert to Pandas for analysis
import pandas as pd
df = dataset["data"].to_pandas()
# Analyze predicate distribution
print(df["predicate"].value_counts())
```
### Dataset Format
The dataset contains all triples in a single **data** split, suitable for machine learning tasks such as:
- Knowledge graph completion
- Link prediction
- Entity embedding
- Relation extraction
- Graph neural networks
### Format Specification
For complete technical documentation of the RDF-to-HuggingFace format, see:
📖 [RDF to HuggingFace Format Specification](https://github.com/CleverThis/cleverernie/blob/master/docs/rdf_huggingface_format_specification.md)
The specification includes:
- Detailed schema definition
- All RDF node type mappings
- Performance benchmarks
- Edge cases and limitations
- Complete code examples
### Conversion Metadata
- **Source Format**: {dataset_info.format}
- **Original Size**: {dataset_info.size_gb} GB
- **Conversion Tool**: [CleverErnie RDF Pipeline](https://github.com/CleverThis/cleverernie)
- **Format Version**: 1.0
- **Conversion Date**: {datetime.now().strftime("%Y-%m-%d")}
"""
def create_dataset_card(dataset_info: DatasetInfo) -> str:
"""Create README.md content for the dataset card.
Args:
dataset_info: Dataset information
Returns:
Markdown content for dataset card
"""
return f"""---
license: {get_valid_hf_license(dataset_info.license)}
task_categories:
- text-generation
- feature-extraction
language:
- en
tags:
- rdf
- knowledge-graph
- semantic-web
- triples
size_categories:
- {get_size_category(dataset_info)}
---
# {dataset_info.name}
## Dataset Description
{dataset_info.description}
**Original Source:** {dataset_info.url}
### Dataset Summary
This dataset contains RDF triples from {dataset_info.name} converted to HuggingFace dataset format
for easy use in machine learning pipelines.
- **Format:** Originally {dataset_info.format}, converted to HuggingFace Dataset
- **Size:** {dataset_info.size_gb} GB (extracted)
- **Entities:** {dataset_info.entities if dataset_info.entities else "N/A"}
- **Triples:** {dataset_info.triples if dataset_info.triples else "N/A"}
- **Original License:** {dataset_info.license if dataset_info.license else "See original source"}
### Recommended Use
{dataset_info.recommended_for}
{f"### Notes\\n\\n{dataset_info.notes}" if dataset_info.notes else ""}
{get_rdf_format_documentation(dataset_info)}
## Citation
If you use this dataset, please cite the original source:
**Original Dataset:** {dataset_info.name}
**URL:** {dataset_info.url}
**License:** {dataset_info.license if dataset_info.license else "See original source"}
## Dataset Preparation
This dataset was prepared using the CleverErnie GISM framework:
```bash
# Download original dataset
python scripts/rdf_dataset_downloader.py {dataset_info.id} -o datasets/
# Convert to HuggingFace format
python scripts/convert_rdf_to_hf_dataset.py \\
datasets/{dataset_info.id}/[file] \\
hf_datasets/{dataset_info.id} \\
--format {get_rdf_format(dataset_info)}
# Upload to HuggingFace Hub
python scripts/upload_all_datasets.py --dataset {dataset_info.id}
```
## Additional Information
### Original Source
{dataset_info.url}
### Conversion Details
- Converted using: [CleverErnie GISM](https://github.com/cleverthis/cleverernie)
- Conversion script: `scripts/convert_rdf_to_hf_dataset.py`
- Dataset format: Single 'data' split with all triples
### Maintenance
This dataset is maintained by the CleverThis organization.
"""
def get_size_category(dataset_info: DatasetInfo) -> str:
"""Get HuggingFace size category.
Args:
dataset_info: Dataset information
Returns:
Size category string
"""
if dataset_info.size_gb < 0.01:
return "n<1K"
elif dataset_info.size_gb < 1:
return "1K<n<10K"
elif dataset_info.size_gb < 10:
return "10K<n<100K"
elif dataset_info.size_gb < 100:
return "100K<n<1M"
elif dataset_info.size_gb < 1000:
return "1M<n<10M"
else:
return "10M<n<100M"
def process_dataset(
dataset_id: str,
base_dir: Path,
skip_download: bool = False,
skip_convert: bool = False,
skip_upload: bool = False,
dry_run: bool = False,
remove_after: bool = False,
) -> bool:
"""Process a single dataset: download, convert, upload.
Args:
dataset_id: Dataset identifier
base_dir: Base directory for operations
skip_download: Skip download step
skip_convert: Skip conversion step
skip_upload: Skip upload step
dry_run: Show what would be done without doing it
remove_after: Remove downloaded dataset after processing
Returns:
True if successful, False otherwise
"""
dataset_info = DATASET_REGISTRY.get(dataset_id)
if not dataset_info:
console.print(f"[red]Dataset '{dataset_id}' not found[/red]")
return False
# Check if dataset is available
if not dataset_info.available:
console.print(f"\n[bold yellow]⊘ Skipping '{dataset_info.name}'[/bold yellow]")
console.print("[yellow]Status:[/yellow] Dataset currently unavailable")
console.print(f"[yellow]Reason:[/yellow] {dataset_info.notes}")
console.print(f"[dim]URL: {dataset_info.url}[/dim]")
return None # Return None to indicate skipped (not failed)
console.print(f"\n[bold cyan]{'[DRY RUN] ' if dry_run else ''}Processing: {dataset_info.name}[/bold cyan]")
# Setup directories
download_dir = base_dir / "downloads" / dataset_id
hf_dataset_dir = base_dir / "hf_datasets" / dataset_id
download_dir.mkdir(parents=True, exist_ok=True)
hf_dataset_dir.mkdir(parents=True, exist_ok=True)
# Step 1: Download
if not skip_download:
console.print("\n[yellow]Step 1: Downloading raw RDF dataset...[/yellow]")
if dry_run:
console.print(
f"[dim]Would run: python scripts/rdf_dataset_downloader.py {dataset_id} -o {download_dir}[/dim]"
)
else:
result = subprocess.run(
[
sys.executable,
str(Path(__file__).parent / "rdf_dataset_downloader.py"),
dataset_id,
"-o",
str(download_dir),
],
capture_output=False,
)
if result.returncode != 0:
console.print(f"[red]Download failed for {dataset_id}[/red]")
return False
else:
console.print("[dim]Skipping download[/dim]")
# Find the downloaded file (only if we need it for conversion)
rdf_file = None
if not skip_convert:
# Special handling for FB15k-237 and similar TSV datasets
if dataset_info.format == "tsv" and dataset_id == "fb15k-237":
# For FB15k-237, look for train.txt in the Release directory
train_files = list(download_dir.rglob("train.txt"))
if train_files:
rdf_file = train_files[0]
console.print(f"[green]Found FB15k-237 train file: {rdf_file}[/green]")
else:
console.print(f"[red]No train.txt file found for FB15k-237 in {download_dir}[/red]")
return False
else:
# Original file finding logic for other datasets
rdf_files = (
list(download_dir.rglob("*.ttl"))
+ list(download_dir.rglob("*.nt"))
+ list(download_dir.rglob("*.rdf"))
+ list(download_dir.rglob("*.owl"))
+ list(download_dir.rglob("*.n3"))
+ list(download_dir.rglob("*.trig"))
+ list(download_dir.rglob("*.nq"))
+ list(download_dir.rglob("*.xml"))
# Also look for compressed versions
+ list(download_dir.rglob("*.ttl.bz2"))
+ list(download_dir.rglob("*.ttl.gz"))
+ list(download_dir.rglob("*.nt.bz2"))
+ list(download_dir.rglob("*.nt.gz"))
+ list(download_dir.rglob("*.rdf.bz2"))
+ list(download_dir.rglob("*.rdf.gz"))
)
# Special handling for geonames - it's XML format but stored as .txt
if dataset_id == "geonames" and dataset_info.format == "xml":
rdf_files += list(download_dir.rglob("*.txt"))
# For TSV datasets (other than FB15k-237), look for .txt or .tsv files
if dataset_info.format == "tsv":
rdf_files += list(download_dir.rglob("*.txt")) + list(download_dir.rglob("*.tsv"))
# For CSV datasets (like ConceptNet), look for .csv files
if dataset_info.format == "csv":
rdf_files += list(download_dir.rglob("*.csv")) + list(download_dir.rglob("*.csv.gz"))
# Check if there are unextracted compressed files
if not rdf_files:
# Check for ZIP files
zip_files = list(download_dir.rglob("*.zip"))
if zip_files:
console.print(f"[yellow]Found {len(zip_files)} ZIP file(s), attempting extraction...[/yellow]")
for zip_file in zip_files:
try:
import zipfile
# Check if it's a valid ZIP file
if zipfile.is_zipfile(zip_file):
extract_dir = zip_file.parent / zip_file.stem
extract_dir.mkdir(parents=True, exist_ok=True)
with zipfile.ZipFile(zip_file, "r") as zip_ref:
console.print(f"[cyan]Extracting {zip_file.name} to {extract_dir}...[/cyan]")
zip_ref.extractall(extract_dir)
console.print(
f"[green]Successfully extracted {len(zip_ref.namelist())} files[/green]"
)
else:
console.print(f"[red]Invalid or incomplete ZIP file: {zip_file}[/red]")
console.print(
f"[yellow]File size: {zip_file.stat().st_size / (1024**3):.2f} GB[/yellow]"
)
if (
dataset_info.compressed_size_gb
and zip_file.stat().st_size / (1024**3) < dataset_info.compressed_size_gb * 0.9
):
console.print(
f"[red]File appears incomplete. Expected ~{dataset_info.compressed_size_gb} GB[/red]"
)
console.print(
f"[yellow]Please re-download using: python scripts/rdf_dataset_downloader.py {dataset_id} --force[/yellow]"
)
return False
except Exception as e:
console.print(f"[red]Failed to extract {zip_file}: {e}[/red]")
if "end-of-stream marker" in str(e).lower() or "compressed file ended" in str(e).lower():
console.print(f"[red]The compressed file is incomplete or corrupted.[/red]")
console.print(
f"[yellow]File size: {zip_file.stat().st_size / (1024**3):.3f} GB[/yellow]"
)
if dataset_info.compressed_size_gb:
console.print(
f"[yellow]Expected size: ~{dataset_info.compressed_size_gb:.2f} GB[/yellow]"
)
console.print(
f"[yellow]Please re-download using: python scripts/rdf_dataset_downloader.py {dataset_id} --force[/yellow]"
)
return False
# Check for BZ2 files that haven't been extracted
bz2_files = list(download_dir.rglob("*.bz2"))
if bz2_files:
console.print(f"[yellow]Found {len(bz2_files)} BZ2 file(s), checking integrity...[/yellow]")
for bz2_file in bz2_files:
# First check file size if we have expected size
file_size_gb = bz2_file.stat().st_size / (1024**3)
if dataset_info.compressed_size_gb:
expected_gb = dataset_info.compressed_size_gb
if file_size_gb < expected_gb * 0.9:
console.print(f"[red]Incomplete BZ2 file: {bz2_file.name}[/red]")
console.print(
f"[yellow]File size: {file_size_gb:.3f} GB, Expected: ~{expected_gb:.2f} GB[/yellow]"
)
console.print(
f"[red]File is only {(file_size_gb / expected_gb * 100):.1f}% of expected size[/red]"
)
console.print(f"[yellow]Removing incomplete file and re-downloading...[/yellow]")
bz2_file.unlink()
# Trigger re-download by not skipping it
console.print(f"[cyan]Re-downloading {dataset_id} with proper size...[/cyan]")
result = subprocess.run(
[
sys.executable,
str(Path(__file__).parent / "rdf_dataset_downloader.py"),
dataset_id,
"-o",
str(download_dir.parent),
"--force",
],
capture_output=False,
)
if result.returncode != 0:
console.print(f"[red]Re-download failed for {dataset_id}[/red]")
return False
# After re-download, find the new file
bz2_files = list(download_dir.rglob("*.bz2"))
if not bz2_files:
console.print(f"[red]No BZ2 file found after re-download[/red]")
return False
bz2_file = bz2_files[0]
try:
# Test the bz2 file integrity
test_result = subprocess.run(
["bunzip2", "-t", str(bz2_file)], capture_output=True, text=True
)
if test_result.returncode != 0:
console.print(f"[red]Corrupted BZ2 file: {bz2_file.name}[/red]")
console.print(
f"[yellow]File size: {bz2_file.stat().st_size / (1024**2):.2f} MB[/yellow]"
)
console.print(f"[yellow]Removing corrupted file and attempting re-download...[/yellow]")
bz2_file.unlink()
# Trigger re-download
result = subprocess.run(
[
sys.executable,
str(Path(__file__).parent / "rdf_dataset_downloader.py"),
dataset_id,
"-o",
str(download_dir.parent),
"--force",
],
capture_output=False,
)
if result.returncode != 0:
console.print(f"[red]Re-download failed for {dataset_id}[/red]")
return False
# After re-download, find the new file
bz2_files = list(download_dir.rglob("*.bz2"))
if not bz2_files:
console.print(f"[red]No BZ2 file found after re-download[/red]")
return False
bz2_file = bz2_files[0]
# Test again
test_result = subprocess.run(
["bunzip2", "-t", str(bz2_file)], capture_output=True, text=True
)
if test_result.returncode != 0:
console.print(
f"[red]File still corrupted after re-download. Manual intervention required.[/red]"
)
return False
# Try to extract the bz2 file
console.print(f"[cyan]Extracting {bz2_file.name}...[/cyan]")
extracted_file = bz2_file.with_suffix("")
with bz2.open(bz2_file, "rb") as f_in:
with extracted_file.open("wb") as f_out:
shutil.copyfileobj(f_in, f_out, length=1024 * 1024)
console.print(f"[green]Successfully extracted to {extracted_file.name}[/green]")
# Re-scan for RDF files after extraction
rdf_files = (
list(download_dir.rglob("*.ttl"))
+ list(download_dir.rglob("*.nt"))
+ list(download_dir.rglob("*.rdf"))
+ list(download_dir.rglob("*.owl"))
+ list(download_dir.rglob("*.xml"))
)
except Exception as e:
console.print(f"[red]Failed to process {bz2_file}: {e}[/red]")
if "end-of-stream marker" in str(e).lower() or "compressed file ended" in str(e).lower():
console.print(f"[yellow]The file is incomplete. Attempting to re-download...[/yellow]")
bz2_file.unlink()
result = subprocess.run(
[
sys.executable,
str(Path(__file__).parent / "rdf_dataset_downloader.py"),
dataset_id,
"-o",
str(download_dir.parent),
"--force",
],
capture_output=False,
)
if result.returncode != 0:
console.print(f"[red]Re-download failed for {dataset_id}[/red]")
return False
# Try again with the new file
bz2_files = list(download_dir.rglob("*.bz2"))
if bz2_files:
bz2_file = bz2_files[0]
try:
console.print(f"[cyan]Extracting re-downloaded {bz2_file.name}...[/cyan]")
extracted_file = bz2_file.with_suffix("")
with bz2.open(bz2_file, "rb") as f_in:
with extracted_file.open("wb") as f_out:
shutil.copyfileobj(f_in, f_out, length=1024 * 1024)
console.print(f"[green]Successfully extracted to {extracted_file.name}[/green]")
# Re-scan for RDF files after extraction
rdf_files = (
list(download_dir.rglob("*.ttl"))
+ list(download_dir.rglob("*.nt"))
+ list(download_dir.rglob("*.rdf"))
+ list(download_dir.rglob("*.owl"))
+ list(download_dir.rglob("*.xml"))
)
except Exception as e2:
console.print(f"[red]Failed again after re-download: {e2}[/red]")
return False
else:
console.print(f"[red]No BZ2 file found after re-download[/red]")
return False
else:
return False
# Re-scan for RDF files after extraction
rdf_files = (
list(download_dir.rglob("*.ttl"))
+ list(download_dir.rglob("*.nt"))
+ list(download_dir.rglob("*.rdf"))
+ list(download_dir.rglob("*.owl"))
+ list(download_dir.rglob("*.n3"))
+ list(download_dir.rglob("*.trig"))
+ list(download_dir.rglob("*.nq"))
+ list(download_dir.rglob("*.xml"))
# Also look for compressed versions
+ list(download_dir.rglob("*.ttl.bz2"))
+ list(download_dir.rglob("*.ttl.gz"))
+ list(download_dir.rglob("*.nt.bz2"))
+ list(download_dir.rglob("*.nt.gz"))
+ list(download_dir.rglob("*.rdf.bz2"))
+ list(download_dir.rglob("*.rdf.gz"))
)
if not rdf_files:
console.print(f"[red]No RDF files found in {download_dir}[/red]")
return False
# Filter out README files if there are other options
non_readme_files = [f for f in rdf_files if "readme" not in f.name.lower()]
if non_readme_files:
rdf_file = non_readme_files[0]
else:
rdf_file = rdf_files[0]
# Check if the selected file is a compressed file and validate its size
if str(rdf_file).endswith(".bz2"):
file_size_gb = rdf_file.stat().st_size / (1024**3)
if dataset_info.compressed_size_gb:
expected_gb = dataset_info.compressed_size_gb
if file_size_gb < expected_gb * 0.9:
console.print(f"[red]Incomplete BZ2 file: {rdf_file.name}[/red]")
console.print(
f"[yellow]File size: {file_size_gb:.3f} GB, Expected: ~{expected_gb:.2f} GB[/yellow]"
)
console.print(
f"[red]File is only {(file_size_gb / expected_gb * 100):.1f}% of expected size[/red]"
)
console.print(f"[yellow]Removing incomplete file and re-downloading...[/yellow]")
rdf_file.unlink()
if skip_download:
console.print(f"[red]Cannot re-download because --skip-download is set[/red]")
console.print(
f"[yellow]Please run: python scripts/rdf_dataset_downloader.py {dataset_id} --force[/yellow]"
)
return False
# Trigger re-download
console.print(f"[cyan]Re-downloading {dataset_id} with proper size...[/cyan]")
result = subprocess.run(
[
sys.executable,
str(Path(__file__).parent / "rdf_dataset_downloader.py"),
dataset_id,
"-o",
str(download_dir.parent),
"--force",
],
capture_output=False,
)
if result.returncode != 0:
console.print(f"[red]Re-download failed for {dataset_id}[/red]")
return False
# After re-download, find the new file again
rdf_files = (
list(download_dir.rglob("*.ttl"))
+ list(download_dir.rglob("*.nt"))
+ list(download_dir.rglob("*.rdf"))
+ list(download_dir.rglob("*.owl"))
+ list(download_dir.rglob("*.n3"))
+ list(download_dir.rglob("*.trig"))
+ list(download_dir.rglob("*.nq"))
+ list(download_dir.rglob("*.xml"))
+ list(download_dir.rglob("*.ttl.bz2"))
+ list(download_dir.rglob("*.ttl.gz"))
+ list(download_dir.rglob("*.nt.bz2"))
+ list(download_dir.rglob("*.nt.gz"))
+ list(download_dir.rglob("*.rdf.bz2"))
+ list(download_dir.rglob("*.rdf.gz"))
)
if not rdf_files:
console.print(f"[red]No RDF files found after re-download[/red]")
return False
# Filter out README files again
non_readme_files = [f for f in rdf_files if "readme" not in f.name.lower()]
if non_readme_files:
rdf_file = non_readme_files[0]
else:
rdf_file = rdf_files[0]
# Also test BZ2 file integrity
console.print(f"[cyan]Testing BZ2 file integrity...[/cyan]")
test_result = subprocess.run(["bunzip2", "-t", str(rdf_file)], capture_output=True, text=True)
if test_result.returncode != 0:
console.print(f"[red]Corrupted BZ2 file: {rdf_file.name}[/red]")
console.print(f"[red]Error: {test_result.stderr}[/red]")
if skip_download:
console.print(f"[red]Cannot re-download because --skip-download is set[/red]")
console.print(
f"[yellow]Please run: python scripts/rdf_dataset_downloader.py {dataset_id} --force[/yellow]"
)
return False
console.print(f"[yellow]Removing corrupted file and re-downloading...[/yellow]")
rdf_file.unlink()
# Re-download
console.print(f"[cyan]Re-downloading {dataset_id}...[/cyan]")
result = subprocess.run(
[
sys.executable,
str(Path(__file__).parent / "rdf_dataset_downloader.py"),
dataset_id,
"-o",
str(download_dir.parent),
"--force",
],
capture_output=False,
)
if result.returncode != 0:
console.print(f"[red]Re-download failed for {dataset_id}[/red]")
return False
# Find the file again
rdf_files = (
list(download_dir.rglob("*.ttl.bz2"))
+ list(download_dir.rglob("*.nt.bz2"))
+ list(download_dir.rglob("*.rdf.bz2"))
)
if not rdf_files:
console.print(f"[red]No BZ2 files found after re-download[/red]")
return False
rdf_file = rdf_files[0]
# Test again
test_result = subprocess.run(["bunzip2", "-t", str(rdf_file)], capture_output=True, text=True)
if test_result.returncode != 0:
console.print(
f"[red]File still corrupted after re-download. Manual intervention required.[/red]"
)
return False
console.print(f"[green]BZ2 file integrity check passed[/green]")
console.print(f"[green]Found RDF file: {rdf_file}[/green]")
# Step 2: Convert to HuggingFace format
if not skip_convert:
console.print("\n[yellow]Step 2: Converting to HuggingFace format...[/yellow]")
# Special handling for FB15k-237
if dataset_id == "fb15k-237":
# Use specialized FB15k-237 converter that handles all splits
fb15k_script = Path(__file__).parent / "convert_fb15k237_to_hf.py"
fb15k_dir = download_dir / "fb15k-237" / "FB15K-237.2"
if dry_run:
console.print(
f"[dim]Would run: python scripts/convert_fb15k237_to_hf.py {fb15k_dir} {hf_dataset_dir}[/dim]"
)
else:
result = subprocess.run(
[
sys.executable,
str(fb15k_script),
str(fb15k_dir),
str(hf_dataset_dir),
],
capture_output=False,
)
if result.returncode != 0:
console.print(f"[red]Conversion failed for {dataset_id}[/red]")
return False
# Create dataset card
readme_path = hf_dataset_dir / "README.md"
readme_path.write_text(create_dataset_card(dataset_info))
console.print(f"[green]Created dataset card: {readme_path}[/green]")
# Special handling for NELL-995
elif dataset_id == "nell-995":
# Use specialized NELL-995 converter that handles raw.kb format
nell_script = Path(__file__).parent / "convert_nell995_to_hf.py"
if dry_run:
console.print(
f"[dim]Would run: python scripts/convert_nell995_to_hf.py {download_dir} {hf_dataset_dir}[/dim]"
)
else:
result = subprocess.run(
[
sys.executable,
str(nell_script),
str(download_dir),
str(hf_dataset_dir),
],
capture_output=False,
)
if result.returncode != 0:
console.print(f"[red]Conversion failed for {dataset_id}[/red]")
return False
# Create dataset card
readme_path = hf_dataset_dir / "README.md"
readme_path.write_text(create_dataset_card(dataset_info))
console.print(f"[green]Created dataset card: {readme_path}[/green]")
# Special handling for ConceptNet
elif dataset_id == "conceptnet":
# Use specialized ConceptNet converter that handles CSV format
conceptnet_script = Path(__file__).parent / "convert_conceptnet_to_hf.py"
# Find the downloaded CSV file
csv_files = list(download_dir.rglob("*.csv")) + list(download_dir.rglob("*.csv.gz"))
if not csv_files:
console.print(f"[red]No CSV files found for ConceptNet in {download_dir}[/red]")
return False
csv_file = csv_files[0]
console.print(f"[green]Found ConceptNet CSV file: {csv_file}[/green]")
if dry_run:
console.print(
f"[dim]Would run: python scripts/convert_conceptnet_to_hf.py {csv_file} {hf_dataset_dir}[/dim]"
)
else:
result = subprocess.run(
[
sys.executable,
str(conceptnet_script),
str(csv_file),
str(hf_dataset_dir),
],
capture_output=False,
)
if result.returncode != 0:
console.print(f"[red]Conversion failed for {dataset_id}[/red]")
return False
# Create dataset card
readme_path = hf_dataset_dir / "README.md"
readme_path.write_text(create_dataset_card(dataset_info))
console.print(f"[green]Created dataset card: {readme_path}[/green]")
else:
# Original conversion logic for other datasets
# Determine if we should use streaming converter based on dataset size
use_streaming = False
streaming_script = Path(__file__).parent / "convert_rdf_to_hf_dataset_streaming_parallel.py"
# Use streaming for medium and larger datasets, or specifically for geonames
if dataset_info.category in ["medium", "large", "xlarge"] or "geonames" in dataset_id.lower():
use_streaming = streaming_script.exists()
if use_streaming:
console.print(f"[dim]Using parallel streaming converter for {dataset_info.category} dataset[/dim]")
if dry_run:
script_name = (
"convert_rdf_to_hf_dataset_streaming_parallel.py"
if use_streaming
else "convert_rdf_to_hf_dataset.py"
)
console.print(f"[dim]Would run: python scripts/{script_name} {rdf_file} {hf_dataset_dir}[/dim]")
else:
# Choose appropriate converter script
if use_streaming:
# Use streaming converter with appropriate chunk size based on dataset size
# Larger chunk sizes are more memory efficient but still reasonable
chunk_size = 50000 # Default for streaming
if dataset_info.size_gb > 10:
chunk_size = 100000 # Larger chunks for very large datasets
elif dataset_info.size_gb < 1:
chunk_size = 10000 # Smaller chunks for small datasets
result = subprocess.run(
[
sys.executable,
str(streaming_script),
str(rdf_file),
str(hf_dataset_dir),
"--format",
get_rdf_format(dataset_info),
"--chunk-size",
str(chunk_size),
"--description",
dataset_info.description,
"--homepage",
dataset_info.url,
"--license",
dataset_info.license if dataset_info.license else "Unknown",
],
capture_output=False,
)
else:
# Use original converter for small datasets
result = subprocess.run(
[
sys.executable,
str(Path(__file__).parent / "convert_rdf_to_hf_dataset.py"),
str(rdf_file),
str(hf_dataset_dir),
"--format",
get_rdf_format(dataset_info),
"--description",
dataset_info.description,
"--homepage",
dataset_info.url,
"--license",
dataset_info.license if dataset_info.license else "Unknown",
],
capture_output=False,
)
if result.returncode != 0:
console.print(f"[red]Conversion failed for {dataset_id}[/red]")
# If compressed file failed, suggest decompressing first
if rdf_file.suffix == ".gz":
console.print(f"[yellow]The compressed file may be corrupted or incomplete.[/yellow]")
console.print(f"[yellow]Attempting to decompress and use partial data...[/yellow]")
# Try to decompress what we can
decompressed_file = rdf_file.with_suffix("")
decompress_result = subprocess.run(
["gunzip", "-c", str(rdf_file)], capture_output=True, text=False
)
if decompress_result.stdout:
# Write whatever we could decompress
decompressed_file.write_bytes(decompress_result.stdout)
console.print(
f"[green]Decompressed {len(decompress_result.stdout) / (1024**2):.2f} MB of data[/green]"
)
# Try conversion again with decompressed file
console.print(f"[yellow]Retrying conversion with decompressed file...[/yellow]")
retry_result = subprocess.run(
[
sys.executable,
str(Path(__file__).parent / "convert_rdf_to_hf_dataset.py"),
str(decompressed_file),
str(hf_dataset_dir),
"--format",
get_rdf_format(dataset_info),
"--description",
dataset_info.description,
"--homepage",
dataset_info.url,
"--license",
dataset_info.license if dataset_info.license else "Unknown",
],
capture_output=False,
)
if retry_result.returncode == 0:
console.print(f"[green]Successfully converted partial dataset[/green]")
# Create dataset card
readme_path = hf_dataset_dir / "README.md"
readme_path.write_text(create_dataset_card(dataset_info))
console.print(f"[green]Created dataset card: {readme_path}[/green]")
return True
else:
console.print(f"[red]Retry also failed[/red]")
return False
else:
console.print(f"[red]Could not decompress any data from the file[/red]")
return False
return False
# Create dataset card
readme_path = hf_dataset_dir / "README.md"
readme_path.write_text(create_dataset_card(dataset_info))
console.print(f"[green]Created dataset card: {readme_path}[/green]")
else:
console.print("[dim]Skipping conversion[/dim]")
# Step 3: Upload to HuggingFace Hub
if not skip_upload:
console.print("\n[yellow]Step 3: Uploading to HuggingFace Hub...[/yellow]")
# Get repository ID from dataset info or config
if hasattr(dataset_info, 'huggingface_repo') and dataset_info.huggingface_repo:
repo_id = dataset_info.huggingface_repo
else:
# Fallback to default pattern from config
config = get_dataset_config()
org = config.get("organization", "CleverThis")
repo_id = f"{org}/{dataset_id}"
if dry_run:
console.print(f"[dim]Would upload to: {repo_id}[/dim]")
else:
try:
from datasets import load_from_disk
from huggingface_hub import HfApi
# Load dataset
dataset = load_from_disk(str(hf_dataset_dir))
# Upload
console.print(f"[cyan]Uploading to {repo_id}...[/cyan]")
dataset.push_to_hub(repo_id, private=False)
# Upload README separately if it exists
readme_path = hf_dataset_dir / "README.md"
if readme_path.exists():
api = HfApi()
api.upload_file(
path_or_fileobj=str(readme_path),
path_in_repo="README.md",
repo_id=repo_id,
repo_type="dataset",
)
else:
# Create README if it doesn't exist (e.g., when skipping conversion)
console.print("[yellow]README.md not found, creating it now...[/yellow]")
readme_path.write_text(create_dataset_card(dataset_info))
api = HfApi()
api.upload_file(
path_or_fileobj=str(readme_path),
path_in_repo="README.md",
repo_id=repo_id,
repo_type="dataset",
)
console.print(f"[bold green]✓ Successfully uploaded to {repo_id}[/bold green]")
except Exception as e:
console.print(f"[red]Upload failed: {e}[/red]")
return False
else:
console.print("[dim]Skipping upload[/dim]")
# Step 4: Cleanup if --rm flag is set
if remove_after and not dry_run:
console.print(f"\n[yellow]--rm: Removing downloaded dataset for {dataset_id}[/yellow]")
if download_dir.exists():
shutil.rmtree(download_dir)
console.print(f"[green]✓ Removed {download_dir}[/green]")
return True
def main():
"""Main entry point."""
parser = argparse.ArgumentParser(description="Download, convert, and upload RDF datasets to HuggingFace")
parser.add_argument(
"--dataset",
"-d",
type=str,
action="append",
help="Specific dataset(s) to process (can be specified multiple times)",
)
parser.add_argument(
"--category",
"-c",
choices=["small", "medium", "large", "xlarge", "all"],
default="all",
help="Process datasets by category",
)
parser.add_argument("--base-dir", type=Path, default=Path("./dataset_processing"), help="Base directory for work")
parser.add_argument("--skip-download", action="store_true", help="Skip download step")
parser.add_argument("--skip-convert", action="store_true", help="Skip conversion step")
parser.add_argument("--skip-upload", action="store_true", help="Skip upload step")
parser.add_argument("--dry-run", action="store_true", help="Show what would be done without doing it")
parser.add_argument("--list", action="store_true", help="List datasets that would be processed")
parser.add_argument(
"--rm",
action="store_true",
help="Remove mode: wipe all downloads at start, delete each dataset after processing",
)
parser.add_argument(
"--parallel",
"-p",
type=int,
default=1,
help="Number of datasets to process in parallel (default: 1, sequential)",
)
args = parser.parse_args()
# Determine which datasets to process
if args.dataset:
datasets_to_process = args.dataset
elif args.category == "all":
datasets_to_process = list(DATASET_REGISTRY.keys())
else:
datasets_to_process = [d_id for d_id, d_info in DATASET_REGISTRY.items() if d_info.category == args.category]
# List mode
if args.list:
console.print("\n[bold cyan]Datasets to Process:[/bold cyan]\n")
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")
table.add_column("Available", style="magenta")
for dataset_id in datasets_to_process:
info = DATASET_REGISTRY[dataset_id]
available = "[green]✓[/green]" if info.available else "[red]✗[/red]"
table.add_row(dataset_id, info.name, info.category, f"{info.size_gb} GB", available)
console.print(table)
available_count = sum(1 for d_id in datasets_to_process if DATASET_REGISTRY[d_id].available)
console.print(f"\n[bold]Total: {len(datasets_to_process)} datasets ({available_count} available)[/bold]")
return
# Process datasets
console.print(f"\n[bold]Processing {len(datasets_to_process)} dataset(s)[/bold]")
console.print(f"Base directory: {args.base_dir}")
console.print(f"Dry run: {args.dry_run}")
console.print(f"Parallel workers: {args.parallel}\n")
# If --rm flag is set, wipe all previously downloaded datasets at start
if args.rm and not args.dry_run:
downloads_dir = args.base_dir / "downloads"
if downloads_dir.exists():
console.print(f"[yellow]--rm: Removing all previous downloads from {downloads_dir}[/yellow]")
shutil.rmtree(downloads_dir)
console.print("[green]✓ Previous downloads removed[/green]\n")
results = {}
# Check if parallel processing is requested
if args.parallel > 1:
console.print(f"[cyan]Using parallel processing with {args.parallel} workers[/cyan]\n")
# Create progress tracking for parallel execution
overall_progress = Progress(
SpinnerColumn(),
TextColumn("[bold blue]Overall Progress"),
BarColumn(),
MofNCompleteColumn(),
TimeElapsedColumn(),
)
dataset_progress = Progress(
TextColumn("[bold cyan]{task.fields[dataset_id]:30}"),
SpinnerColumn(),
TextColumn("{task.description}"),
BarColumn(),
TextColumn("[progress.percentage]{task.percentage:>3.0f}%"),
TimeElapsedColumn(),
expand=True,
)
# Create progress table
def generate_progress_table() -> Table:
"""Generate a table showing all progress bars."""
table = Table(title="Dataset Processing Progress", expand=True)
table.add_column("Progress", ratio=1)
# Add overall progress
table.add_row(overall_progress)
table.add_row("") # Spacer
# Add individual dataset progress
table.add_row(dataset_progress)
return table
# Process datasets with live progress display
with Live(generate_progress_table(), console=console, refresh_per_second=4) as live:
# Create a lock for thread-safe progress updates
progress_lock = threading.Lock()
# Track tasks for each dataset
task_ids = {}
for dataset_id in datasets_to_process:
task_id = dataset_progress.add_task("[dim]Waiting...[/dim]", dataset_id=dataset_id, total=100)
task_ids[dataset_id] = task_id
# Create overall progress task
overall_task = overall_progress.add_task("Processing datasets", total=len(datasets_to_process))
# Function to monitor subprocess and update progress
def monitor_subprocess(proc, task_id, stage_name, stage_color, start_pct, end_pct):
"""Monitor subprocess output and update progress in real-time."""
import re
import time
stage_range = end_pct - start_pct
last_progress = 0
last_update_time = time.time()
# For initial stage update
with progress_lock:
dataset_progress.update(
task_id,
description=f"{stage_color}{stage_name}[/{stage_color.strip('[]')}] (0%)",
completed=start_pct,
)
# Time-based progress estimation as fallback
start_time = time.time()
estimated_duration = 30 # Default estimate in seconds
while True:
line = proc.stdout.readline()
if not line:
break
line_str = line.strip()
# Look for PROGRESS markers first
progress_match = re.search(r"PROGRESS:\s*(\d+)", line_str)
if progress_match:
progress_pct = float(progress_match.group(1))
# Calculate the absolute progress across all stages
overall_pct = start_pct + (progress_pct / 100) * stage_range
with progress_lock:
dataset_progress.update(
task_id,
description=f"{stage_color}{stage_name}[/{stage_color.strip('[]')}] ({progress_pct:.0f}%)",
completed=overall_pct,
)
last_progress = progress_pct
last_update_time = time.time()
continue
# Also try other patterns for incremental updates
patterns = [
(r"\[(\d+)%\]", lambda m: float(m.group(1))), # [45%]
(r"completed.*?(\d+)%", lambda m: float(m.group(1))), # completed: 45%
(
r"(\d+\.?\d*)\s*/\s*(\d+\.?\d*)\s*MB",
lambda m: (float(m.group(1)) / float(m.group(2))) * 100,
), # 45.5/100 MB
(
r"(\d+)\s*/\s*(\d+)\s+triples",
lambda m: (float(m.group(1)) / float(m.group(2))) * 100,
), # 45/100 triples
(r"Downloading", lambda m: 5), # Starting download
(r"Extracting", lambda m: 85), # Extracting phase
(r"Phase 1", lambda m: 10), # Phase 1 marker
(r"Phase 2", lambda m: 40), # Phase 2 marker
(r"Converting", lambda m: 70), # Converting phase
(r"✓.*Successfully", lambda m: 95), # Success marker
]
for pattern, extractor in patterns:
match = re.search(pattern, line_str, re.IGNORECASE)
if match:
try:
progress_pct = extractor(match)
if progress_pct > last_progress:
last_progress = progress_pct
overall_pct = start_pct + (progress_pct / 100) * stage_range
with progress_lock:
dataset_progress.update(
task_id,
description=f"{stage_color}{stage_name}[/{stage_color.strip('[]')}] ({progress_pct:.0f}%)",
completed=overall_pct,
)
last_update_time = time.time()
break
except:
pass
# If no progress for 2 seconds, update based on elapsed time
current_time = time.time()
if current_time - last_update_time > 2:
elapsed = current_time - start_time
time_progress = min(90, (elapsed / estimated_duration) * 90) # Cap at 90%
if time_progress > last_progress + 5:
last_progress = time_progress
overall_pct = start_pct + (time_progress / 100) * stage_range
with progress_lock:
dataset_progress.update(
task_id,
description=f"{stage_color}{stage_name}[/{stage_color.strip('[]')}] (~{time_progress:.0f}%)",
completed=overall_pct,
)
last_update_time = current_time
proc.wait()
# Ensure we reach the end percentage for this stage if successful
if proc.returncode == 0:
with progress_lock:
dataset_progress.update(
task_id,
description=f"{stage_color}{stage_name}[/{stage_color.strip('[]')}] (100%)",
completed=end_pct,
)
return proc.returncode == 0
# Function to process dataset and update progress
def process_with_progress(dataset_id: str) -> Tuple[str, bool]:
"""Process a dataset and update progress bar with real-time updates."""
task_id = task_ids[dataset_id]
try:
# Get dataset info
dataset_info = DATASET_REGISTRY.get(dataset_id)
if not dataset_info:
with progress_lock:
dataset_progress.update(task_id, description="[red]✗ Unknown dataset[/red]", completed=100)
return (dataset_id, False)
if not dataset_info.available:
with progress_lock:
dataset_progress.update(
task_id, description="[yellow]⊘ Unavailable[/yellow]", completed=100
)
return (dataset_id, None)
# Setup directories
download_dir = args.base_dir / "downloads" / dataset_id
hf_dataset_dir = args.base_dir / "hf_datasets" / dataset_id
# Calculate progress increments
total_steps = 3 - (args.skip_download + args.skip_convert + args.skip_upload)
if total_steps == 0:
total_steps = 1
step_size = 100 / total_steps
current_step = 0
# Step 1: Download with real-time progress
if not args.skip_download:
start_pct = current_step * step_size
end_pct = (current_step + 1) * step_size
if not args.dry_run:
# Check if wrapper script exists for progress monitoring
download_wrapper = Path(__file__).parent / "download_with_progress.py"
if download_wrapper.exists():
download_script = download_wrapper
else:
download_script = Path(__file__).parent / "rdf_dataset_downloader.py"
# Use Popen for real-time output monitoring
proc = subprocess.Popen(
[
sys.executable,
"-u", # Unbuffered output
str(download_script),
dataset_id,
"-o",
str(download_dir),
],
stdout=subprocess.PIPE,
stderr=subprocess.STDOUT,
text=True,
bufsize=1, # Line buffering
)
# Monitor and update progress
success = monitor_subprocess(
proc, task_id, "📥 Downloading", "[yellow]", start_pct, end_pct
)
if not success:
with progress_lock:
dataset_progress.update(
task_id, description="[red]✗ Download failed[/red]", completed=100
)
return (dataset_id, False)
current_step += 1
# Step 2: Convert with real-time progress
if not args.skip_convert:
start_pct = current_step * step_size
end_pct = (current_step + 1) * step_size
if not args.dry_run:
# Special handling for FB15k-237
if dataset_id == "fb15k-237":
# Use specialized FB15k-237 converter
fb_script = Path(__file__).parent / "convert_fb15k237_to_hf.py"
fb15k_dir = download_dir / "fb15k-237" / "FB15K-237.2"
proc = subprocess.Popen(
[
sys.executable,
"-u",
str(fb_script),
str(fb15k_dir),
str(hf_dataset_dir),
],
stdout=subprocess.PIPE,
stderr=subprocess.STDOUT,
text=True,
bufsize=1,
)
success = monitor_subprocess(
proc, task_id, "🔄 Converting FB15k", "[cyan]", start_pct, end_pct
)
if not success:
with progress_lock:
dataset_progress.update(
task_id, description="[red]✗ Conversion failed[/red]", completed=100
)
return (dataset_id, False)
# Create dataset card
readme_path = hf_dataset_dir / "README.md"
readme_path.write_text(create_dataset_card(dataset_info))
# Special handling for NELL-995
elif dataset_id == "nell-995":
# Use specialized NELL-995 converter
nell_script = Path(__file__).parent / "convert_nell995_to_hf.py"
proc = subprocess.Popen(
[
sys.executable,
"-u",
str(nell_script),
str(download_dir),
str(hf_dataset_dir),
],
stdout=subprocess.PIPE,
stderr=subprocess.STDOUT,
text=True,
bufsize=1,
)
success = monitor_subprocess(
proc, task_id, "🔄 Converting NELL", "[cyan]", start_pct, end_pct
)
if not success:
with progress_lock:
dataset_progress.update(
task_id, description="[red]✗ Conversion failed[/red]", completed=100
)
return (dataset_id, False)
# Create dataset card
readme_path = hf_dataset_dir / "README.md"
readme_path.write_text(create_dataset_card(dataset_info))
# Special handling for ConceptNet
elif dataset_id == "conceptnet":
# Use specialized ConceptNet converter
conceptnet_script = Path(__file__).parent / "convert_conceptnet_to_hf.py"
# Find the downloaded CSV file
csv_files = list(download_dir.rglob("*.csv")) + list(download_dir.rglob("*.csv.gz"))
if not csv_files:
with progress_lock:
dataset_progress.update(
task_id, description="[red]✗ No CSV files found[/red]", completed=100
)
return (dataset_id, False)
csv_file = csv_files[0]
proc = subprocess.Popen(
[
sys.executable,
"-u",
str(conceptnet_script),
str(csv_file),
str(hf_dataset_dir),
],
stdout=subprocess.PIPE,
stderr=subprocess.STDOUT,
text=True,
bufsize=1,
)
success = monitor_subprocess(
proc, task_id, "🔄 Converting ConceptNet", "[cyan]", start_pct, end_pct
)
if not success:
with progress_lock:
dataset_progress.update(
task_id, description="[red]✗ Conversion failed[/red]", completed=100
)
return (dataset_id, False)
# Create dataset card
readme_path = hf_dataset_dir / "README.md"
readme_path.write_text(create_dataset_card(dataset_info))
# Handle regular RDF datasets
else:
# Find the downloaded file
rdf_files = (
list(download_dir.rglob("*.ttl"))
+ list(download_dir.rglob("*.nt"))
+ list(download_dir.rglob("*.rdf"))
+ list(download_dir.rglob("*.owl"))
+ list(download_dir.rglob("*.xml"))
)
# Special handling for geonames - it's XML format but stored as .txt
if dataset_id == "geonames" and dataset_info.format == "xml":
rdf_files += list(download_dir.rglob("*.txt"))
if dataset_info.format == "tsv":
rdf_files += list(download_dir.rglob("*.txt")) + list(download_dir.rglob("*.tsv"))
# Also check for CSV files for special datasets
if dataset_info.format == "csv":
rdf_files += list(download_dir.rglob("*.csv"))
if not rdf_files:
with progress_lock:
dataset_progress.update(
task_id, description="[red]✗ No RDF files found[/red]", completed=100
)
return (dataset_id, False)
# Filter out README files if there are other options
non_readme_files = [f for f in rdf_files if "readme" not in f.name.lower()]
if non_readme_files:
rdf_file = non_readme_files[0]
else:
rdf_file = rdf_files[0]
# Determine converter with progress wrapper
use_streaming = (
dataset_info.category in ["medium", "large", "xlarge"]
or "geonames" in dataset_id.lower()
)
# Check for wrapper scripts that provide progress
convert_wrapper = Path(__file__).parent / "convert_rdf_progress.py"
streaming_wrapper = Path(__file__).parent / "convert_rdf_streaming_progress.py"
if use_streaming:
if streaming_wrapper.exists():
converter_script = streaming_wrapper
else:
converter_script = (
Path(__file__).parent / "convert_rdf_to_hf_dataset_streaming_parallel.py"
)
else:
if convert_wrapper.exists():
converter_script = convert_wrapper
else:
converter_script = Path(__file__).parent / "convert_rdf_to_hf_dataset.py"
# Build command
cmd = [
sys.executable,
"-u",
str(converter_script),
str(rdf_file),
str(hf_dataset_dir),
"--format",
get_rdf_format(dataset_info),
]
# Add additional arguments for streaming converter
if use_streaming and "streaming" in converter_script.name:
# Use appropriate chunk size based on dataset size
chunk_size = 50000 # Default for streaming
if dataset_info.size_gb > 10:
chunk_size = 100000 # Larger chunks for very large datasets
elif dataset_info.size_gb < 1:
chunk_size = 10000 # Smaller chunks for small datasets
cmd.extend(["--chunk-size", str(chunk_size)])
# Add metadata arguments
cmd.extend(
[
"--description",
dataset_info.description,
"--homepage",
dataset_info.url,
"--license",
dataset_info.license if dataset_info.license else "Unknown",
]
)
# Run conversion with monitoring
proc = subprocess.Popen(
cmd, stdout=subprocess.PIPE, stderr=subprocess.STDOUT, text=True, bufsize=1
)
success = monitor_subprocess(
proc, task_id, "🔄 Converting", "[cyan]", start_pct, end_pct
)
if not success:
with progress_lock:
dataset_progress.update(
task_id, description="[red]✗ Conversion failed[/red]", completed=100
)
return (dataset_id, False)
# Create dataset card
readme_path = hf_dataset_dir / "README.md"
readme_path.write_text(create_dataset_card(dataset_info))
current_step += 1
# Step 3: Upload with progress tracking
if not args.skip_upload:
start_pct = current_step * step_size
end_pct = (current_step + 1) * step_size
if not args.dry_run:
try:
from datasets import load_from_disk
from huggingface_hub import HfApi
import time
# Update start of upload
with progress_lock:
dataset_progress.update(
task_id,
description="[green]📤 Loading dataset...[/green] (0%)",
completed=start_pct,
)
# Load dataset
dataset = load_from_disk(str(hf_dataset_dir))
# Update after loading
with progress_lock:
dataset_progress.update(
task_id,
description="[green]📤 Uploading to Hub...[/green] (20%)",
completed=start_pct + 0.2 * (end_pct - start_pct),
)
# Upload dataset
repo_id = f"CleverThis/{dataset_id}"
dataset.push_to_hub(repo_id, private=False)
# Update after main upload
with progress_lock:
dataset_progress.update(
task_id,
description="[green]📤 Uploading README...[/green] (80%)",
completed=start_pct + 0.8 * (end_pct - start_pct),
)
# Upload README separately if it exists
readme_path = hf_dataset_dir / "README.md"
if not readme_path.exists():
# Create README if it doesn't exist
readme_path.write_text(create_dataset_card(dataset_info))
api = HfApi()
api.upload_file(
path_or_fileobj=str(readme_path),
path_in_repo="README.md",
repo_id=repo_id,
repo_type="dataset",
)
# Mark complete
with progress_lock:
dataset_progress.update(
task_id, description="[green]📤 Uploaded[/green] (100%)", completed=end_pct
)
except Exception as e:
with progress_lock:
dataset_progress.update(
task_id, description=f"[red]✗ Upload failed: {str(e)[:30]}[/red]", completed=100
)
return (dataset_id, False)
current_step += 1
# Cleanup if requested
if args.rm and not args.dry_run:
if download_dir.exists():
shutil.rmtree(download_dir)
# Mark as complete
with progress_lock:
dataset_progress.update(task_id, description="[green]✓ Complete[/green]", completed=100)
return (dataset_id, True)
except Exception as e:
with progress_lock:
dataset_progress.update(
task_id, description=f"[red]✗ Error: {str(e)[:30]}[/red]", completed=100
)
return (dataset_id, False)
# Use ThreadPoolExecutor for I/O-bound download/upload operations
with ThreadPoolExecutor(max_workers=args.parallel) as executor:
# Submit all tasks
futures = {
executor.submit(process_with_progress, dataset_id): dataset_id for dataset_id in datasets_to_process
}
# Process completed tasks
completed = 0
for future in as_completed(futures):
dataset_id = futures[future]
try:
dataset_id, success = future.result()
results[dataset_id] = success
except Exception as e:
console.print(f"[red]Error processing {dataset_id}: {e}[/red]")
results[dataset_id] = False
completed += 1
with progress_lock:
overall_progress.update(overall_task, completed=completed)
# Update live display
live.update(generate_progress_table())
else:
# Sequential processing (original behavior)
for dataset_id in datasets_to_process:
success = process_dataset(
dataset_id=dataset_id,
base_dir=args.base_dir,
skip_download=args.skip_download,
skip_convert=args.skip_convert,
skip_upload=args.skip_upload,
dry_run=args.dry_run,
remove_after=args.rm,
)
results[dataset_id] = success
# Summary
console.print("\n[bold cyan]Processing Summary:[/bold cyan]\n")
table = Table(show_header=True)
table.add_column("Dataset", style="cyan")
table.add_column("Status", style="green")
for dataset_id, success in results.items():
if success is None:
status = "[yellow]⊘ Skipped[/yellow]"
elif success:
status = "[green]✓ Success[/green]"
else:
status = "[red]✗ Failed[/red]"
table.add_row(dataset_id, status)
console.print(table)
successful = sum(1 for v in results.values() if v is True)
skipped = sum(1 for v in results.values() if v is None)
failed = sum(1 for v in results.values() if v is False)
result_msg = f"Results: {successful}/{len(results)} successful"
if skipped:
result_msg += f", {skipped} skipped"
if failed:
result_msg += f", {failed} failed"
console.print(f"\n[bold]{result_msg}[/bold]")
if __name__ == "__main__":
main()
Reference in New Issue View Git Blame Copy Permalink