241 lines
7.7 KiB
Python
241 lines
7.7 KiB
Python
#!/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 (
|
|
BarColumn,
|
|
Progress,
|
|
SpinnerColumn,
|
|
TextColumn,
|
|
TimeElapsedColumn,
|
|
)
|
|
|
|
|
|
def parse_nell995_kb_file(file_path: Path) -> list[dict[str, str | None]]:
|
|
"""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 = []
|
|
|
|
try:
|
|
with open(file_path, 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
|
|
"object_datatype": None,
|
|
"object_language": None,
|
|
}
|
|
triples.append(triple)
|
|
except FileNotFoundError:
|
|
console = Console()
|
|
console.print(f"[red]File not found: {file_path}[/red]")
|
|
console.print(
|
|
"[yellow]Make sure the NELL-995 dataset has been downloaded "
|
|
"correctly[/yellow]"
|
|
)
|
|
raise
|
|
except OSError as e:
|
|
console = Console()
|
|
console.print(f"[red]Error reading file {file_path}: {e}[/red]")
|
|
raise
|
|
|
|
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}")
|
|
# List what files are actually there
|
|
try:
|
|
all_files = list(input_dir.rglob("*"))
|
|
if all_files:
|
|
console.print(
|
|
f"[yellow]Files found: {[f.name for f in all_files[:10]]}[/yellow]"
|
|
)
|
|
if len(all_files) > 10:
|
|
console.print(
|
|
f"[yellow]... and {len(all_files) - 10} more files[/yellow]"
|
|
)
|
|
else:
|
|
console.print(f"[yellow]Directory is empty: {input_dir}[/yellow]")
|
|
except OSError:
|
|
console.print(f"[yellow]Cannot access directory: {input_dir}[/yellow]")
|
|
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)
|
|
if not triples:
|
|
console.print(
|
|
f"[yellow]Warning: No valid triples found in {kb_file}[/yellow]"
|
|
)
|
|
progress.update(
|
|
task, description=f"[green]✓ Parsed {len(triples):,} triples[/green]"
|
|
)
|
|
|
|
# Create dataset dictionary with single "data" split
|
|
try:
|
|
dataset_dict = DatasetDict({"data": Dataset.from_list(triples)})
|
|
except Exception as e:
|
|
console.print(f"[red]Error creating dataset: {e}[/red]")
|
|
console.print(f"[yellow]Total triples processed: {len(triples)}[/yellow]")
|
|
raise
|
|
|
|
# Add metadata
|
|
metadata_dict = {
|
|
"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",
|
|
}
|
|
for split in dataset_dict:
|
|
dataset_dict[split].info.description = metadata_dict.get("description", "")
|
|
dataset_dict[split].info.citation = metadata_dict.get("citation", "")
|
|
dataset_dict[split].info.homepage = metadata_dict.get("homepage", "")
|
|
dataset_dict[split].info.license = metadata_dict.get("license", "")
|
|
|
|
# Save dataset
|
|
try:
|
|
output_dir.mkdir(parents=True, exist_ok=True)
|
|
except OSError as e:
|
|
console.print(f"[red]Error creating output directory: {e}[/red]")
|
|
console.print(f"[yellow]Output path: {output_dir}[/yellow]")
|
|
console.print("[yellow]Check permissions and disk space[/yellow]")
|
|
raise
|
|
|
|
try:
|
|
dataset_dict.save_to_disk(str(output_dir))
|
|
except Exception as e:
|
|
console.print(f"[red]Error saving dataset: {e}[/red]")
|
|
console.print(f"[yellow]Output directory: {output_dir}[/yellow]")
|
|
console.print("[yellow]Check disk space and permissions[/yellow]")
|
|
raise
|
|
|
|
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() -> None:
|
|
"""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()
|