243 lines
7.9 KiB
Python
Executable File
243 lines
7.9 KiB
Python
Executable File
#!/usr/bin/env python3
|
|
"""Convert FB15k-237 dataset to HuggingFace format with proper splits.
|
|
|
|
This script specifically handles the FB15k-237 dataset which has train/valid/test
|
|
splits in separate TSV files.
|
|
"""
|
|
|
|
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_fb15k237_file(file_path: Path) -> list[dict[str, str | None]]:
|
|
"""Parse FB15k-237 TSV file.
|
|
|
|
Args:
|
|
file_path: Path to TSV file (train.txt, valid.txt, or test.txt)
|
|
|
|
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 Freebase IDs to URIs
|
|
triple = {
|
|
"subject": f"http://freebase.com/{subject.strip()}",
|
|
"predicate": f"http://freebase.com/{predicate.strip()}",
|
|
"object": f"http://freebase.com/{obj.strip()}",
|
|
"object_type": "uri", # FB15k-237 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 FB15k-237 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_fb15k237(input_dir: Path, output_dir: Path) -> None:
|
|
"""Convert FB15k-237 dataset to HuggingFace format.
|
|
|
|
Args:
|
|
input_dir: Directory containing train.txt, valid.txt, test.txt
|
|
output_dir: Output directory for HuggingFace dataset
|
|
"""
|
|
console = Console()
|
|
|
|
console.print(
|
|
"\n[bold cyan]Converting FB15k-237 to HuggingFace Dataset[/bold cyan]"
|
|
)
|
|
console.print(f"Input: {input_dir}")
|
|
console.print(f"Output: {output_dir}\n")
|
|
|
|
# Find the Release directory if it exists
|
|
if (input_dir / "Release").exists():
|
|
input_dir = input_dir / "Release"
|
|
|
|
# Check for required files
|
|
train_file = input_dir / "train.txt"
|
|
valid_file = input_dir / "valid.txt"
|
|
test_file = input_dir / "test.txt"
|
|
|
|
if not train_file.exists():
|
|
console.print(f"[red]Error: train.txt not found in {input_dir}[/red]")
|
|
console.print(
|
|
"[yellow]Expected FB15k-237 structure with train.txt, valid.txt, "
|
|
"test.txt[/yellow]"
|
|
)
|
|
# List what files are actually there
|
|
try:
|
|
all_files = list(input_dir.glob("*"))
|
|
if all_files:
|
|
console.print(
|
|
f"[yellow]Files found: {[f.name for f in all_files[:10]]}[/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
|
|
|
|
all_triples = []
|
|
|
|
with Progress(
|
|
SpinnerColumn(),
|
|
TextColumn("[progress.description]{task.description}"),
|
|
BarColumn(),
|
|
TimeElapsedColumn(),
|
|
console=console,
|
|
transient=False,
|
|
) as progress:
|
|
# Parse train split
|
|
if train_file.exists():
|
|
task = progress.add_task(
|
|
"[yellow]Parsing train.txt...[/yellow]", total=None
|
|
)
|
|
train_triples = parse_fb15k237_file(train_file)
|
|
if not train_triples:
|
|
console.print(
|
|
f"[yellow]Warning: No valid triples found in {train_file}[/yellow]"
|
|
)
|
|
progress.update(
|
|
task,
|
|
description=(
|
|
f"[green]✓ Parsed {len(train_triples):,} training triples[/green]"
|
|
),
|
|
)
|
|
all_triples.extend(train_triples)
|
|
|
|
# Parse validation split
|
|
if valid_file.exists():
|
|
task = progress.add_task(
|
|
"[yellow]Parsing valid.txt...[/yellow]", total=None
|
|
)
|
|
valid_triples = parse_fb15k237_file(valid_file)
|
|
progress.update(
|
|
task,
|
|
description=(
|
|
f"[green]✓ Parsed {len(valid_triples):,} validation triples[/green]"
|
|
),
|
|
)
|
|
all_triples.extend(valid_triples)
|
|
|
|
# Parse test split
|
|
if test_file.exists():
|
|
task = progress.add_task("[yellow]Parsing test.txt...[/yellow]", total=None)
|
|
test_triples = parse_fb15k237_file(test_file)
|
|
progress.update(
|
|
task,
|
|
description=(
|
|
f"[green]✓ Parsed {len(test_triples):,} test triples[/green]"
|
|
),
|
|
)
|
|
all_triples.extend(test_triples)
|
|
|
|
# Create dataset dictionary with single "data" split
|
|
try:
|
|
dataset_dict = DatasetDict({"data": Dataset.from_list(all_triples)})
|
|
except Exception as e:
|
|
console.print(f"[red]Error creating dataset: {e}[/red]")
|
|
console.print(f"[yellow]Total triples processed: {len(all_triples)}[/yellow]")
|
|
raise
|
|
|
|
# Add metadata
|
|
metadata_dict = {
|
|
"description": "FB15k-237 Knowledge Base Completion Dataset",
|
|
"citation": """@inproceedings{toutanova2015observed,
|
|
title={Observed versus latent features for knowledge base and text inference},
|
|
author={Toutanova, Kristina and Chen, Danqi},
|
|
booktitle={Proceedings of the 3rd workshop on continuous vector space """
|
|
"""models and their compositionality},
|
|
pages={57--66},
|
|
year={2015}
|
|
}""",
|
|
"homepage": "https://www.microsoft.com/en-us/download/details.aspx?id=52312",
|
|
"license": "CC BY 2.5",
|
|
}
|
|
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(all_triples):,} (combined from train/valid/test)"
|
|
)
|
|
|
|
|
|
def main() -> None:
|
|
"""Main entry point."""
|
|
parser = argparse.ArgumentParser(
|
|
description="Convert FB15k-237 to HuggingFace format"
|
|
)
|
|
parser.add_argument(
|
|
"input", type=Path, help="Input directory containing FB15k-237 files"
|
|
)
|
|
parser.add_argument(
|
|
"output", type=Path, help="Output directory for HuggingFace dataset"
|
|
)
|
|
|
|
args = parser.parse_args()
|
|
|
|
convert_fb15k237(args.input, args.output)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|