dataset-uploader/scripts/convert_nell995_to_hf.py

#!/usr/bin/env python3
"""Convert NELL-995 dataset to HuggingFace format.

This script specifically handles the NELL-995 dataset which has a raw.kb file
with triples in tab-separated format.
"""

from __future__ import annotations

import argparse
from pathlib import Path

from datasets import Dataset, DatasetDict
from rich.console import Console
from rich.progress import Progress, SpinnerColumn, TextColumn, BarColumn, TimeElapsedColumn


def parse_nell995_kb_file(file_path: Path) -> list[dict[str, str]]:
    """Parse NELL-995 raw.kb file.

    The raw.kb file format:
    - Tab-separated values: subject\trelation\tobject
    - All values are entity/relation names (not URIs)

    Args:
        file_path: Path to raw.kb file

    Returns:
        List of triple dictionaries
    """
    triples = []

    with open(file_path, "r", encoding="utf-8") as f:
        for line in f:
            line = line.strip()
            if not line:
                continue

            parts = line.split("\t")
            if len(parts) != 3:
                continue

            subject, predicate, obj = parts
            # Convert to NELL namespace URIs
            triple = {
                "subject": f"http://nell.ml.cmu.edu/entity/{subject.strip()}",
                "predicate": f"http://nell.ml.cmu.edu/relation/{predicate.strip()}",
                "object": f"http://nell.ml.cmu.edu/entity/{obj.strip()}",
                "object_type": "uri",  # NELL-995 only has entity-to-entity relations
                "object_datatype": None,
                "object_language": None,
            }
            triples.append(triple)

    return triples


def convert_nell995(input_dir: Path, output_dir: Path) -> None:
    """Convert NELL-995 dataset to HuggingFace format.

    Args:
        input_dir: Directory containing NELL-995 files
        output_dir: Output directory for HuggingFace dataset
    """
    console = Console()

    console.print("\n[bold cyan]Converting NELL-995 to HuggingFace Dataset[/bold cyan]")
    console.print(f"Input: {input_dir}")
    console.print(f"Output: {output_dir}\n")

    # Look for raw.kb file
    kb_file = None

    # Check various possible locations
    possible_paths = [
        input_dir / "raw.kb",
        input_dir / "NELL-995" / "raw.kb",
        input_dir / "KB-Reasoning-Data-master" / "NELL-995" / "raw.kb",
        input_dir / "nell-995" / "master" / "KB-Reasoning-Data-master" / "NELL-995" / "raw.kb",
    ]

    for path in possible_paths:
        if path.exists():
            kb_file = path
            break

    if not kb_file:
        console.print(f"[red]Error: raw.kb not found in {input_dir}[/red]")
        console.print("[yellow]Looking for: raw.kb in NELL-995 directory structure[/yellow]")
        console.print("[yellow]Tried paths:[/yellow]")
        for path in possible_paths:
            console.print(f"  - {path}")
        return

    console.print(f"[green]Found raw.kb at: {kb_file}[/green]\n")

    with Progress(
        SpinnerColumn(),
        TextColumn("[progress.description]{task.description}"),
        BarColumn(),
        TimeElapsedColumn(),
        console=console,
        transient=False,
    ) as progress:
        # Parse KB file
        task = progress.add_task("[yellow]Parsing raw.kb...[/yellow]", total=None)
        triples = parse_nell995_kb_file(kb_file)
        progress.update(task, description=f"[green]✓ Parsed {len(triples):,} triples[/green]")

    # Create dataset dictionary with single "data" split
    dataset_dict = DatasetDict({"data": Dataset.from_list(triples)})

    # Add metadata
    dataset_dict._info = {
        "description": "NELL-995 Knowledge Graph for Link Prediction",
        "citation": """@inproceedings{zhang2018variational,
  title={Variational reasoning for question answering with knowledge graph},
  author={Zhang, Yuyu and Dai, Hanjun and Kozareva, Zornitsa and Smola, Alexander J and Song, Le},
  booktitle={Proceedings of the AAAI Conference on Artificial Intelligence},
  volume={32},
  number={1},
  year={2018}
}""",
        "homepage": "https://github.com/wenhuchen/KB-Reasoning-Data",
        "license": "CC BY-NC",
    }

    # Save dataset
    output_dir.mkdir(parents=True, exist_ok=True)
    dataset_dict.save_to_disk(str(output_dir))

    console.print(f"\n[bold green]✓ Dataset saved to {output_dir}[/bold green]")

    # Print statistics
    console.print("\n[bold]Dataset Statistics:[/bold]")
    console.print(f"  • Total triples: {len(triples):,}")

    # Count unique entities and relations
    entities = set()
    relations = set()
    for triple in triples:
        entities.add(triple["subject"])
        entities.add(triple["object"])
        relations.add(triple["predicate"])

    console.print(f"  • Unique entities: {len(entities):,}")
    console.print(f"  • Unique relations: {len(relations):,}")


def main():
    """Main entry point."""
    parser = argparse.ArgumentParser(description="Convert NELL-995 to HuggingFace format")
    parser.add_argument("input", type=Path, help="Input directory containing NELL-995 files")
    parser.add_argument("output", type=Path, help="Output directory for HuggingFace dataset")

    args = parser.parse_args()

    convert_nell995(args.input, args.output)


if __name__ == "__main__":
    main()