300 lines
10 KiB
Python
Executable File
300 lines
10 KiB
Python
Executable File
#!/usr/bin/env python3
|
|
"""Convert ConceptNet CSV to HuggingFace format.
|
|
|
|
This script handles the ConceptNet 5.7 dataset which is in CSV format with
|
|
tab-separated values and specific columns for assertions.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import csv
|
|
import json
|
|
import sys
|
|
from pathlib import Path
|
|
|
|
from datasets import Dataset, DatasetDict
|
|
from rich.console import Console
|
|
from rich.progress import Progress, SpinnerColumn, TextColumn, BarColumn, TimeElapsedColumn
|
|
|
|
|
|
def convert_conceptnet_uri_to_full(uri: str) -> str:
|
|
"""Convert ConceptNet URI format to full URI.
|
|
|
|
Args:
|
|
uri: ConceptNet URI like /c/en/word or /r/RelatedTo
|
|
|
|
Returns:
|
|
Full URI like http://conceptnet.io/c/en/word
|
|
"""
|
|
if uri.startswith("/"):
|
|
return f"http://conceptnet.io{uri}"
|
|
return uri
|
|
|
|
|
|
def parse_conceptnet_csv(file_path: Path, max_rows: int | None = None) -> list[dict[str, str]]:
|
|
"""Parse ConceptNet CSV file.
|
|
|
|
ConceptNet CSV format (tab-separated):
|
|
- assertion URI
|
|
- relation
|
|
- start node (subject)
|
|
- end node (object)
|
|
- edge info JSON (contains weight, sources, etc.)
|
|
|
|
Args:
|
|
file_path: Path to CSV file
|
|
max_rows: Maximum number of rows to parse (for testing)
|
|
|
|
Returns:
|
|
List of triple dictionaries
|
|
"""
|
|
triples = []
|
|
|
|
with open(file_path, "r", encoding="utf-8") as f:
|
|
reader = csv.reader(f, delimiter="\t")
|
|
|
|
for i, row in enumerate(reader):
|
|
if max_rows and i >= max_rows:
|
|
break
|
|
|
|
# Skip malformed rows
|
|
if len(row) < 5:
|
|
continue
|
|
|
|
# Extract fields
|
|
assertion_uri = row[0] # e.g., /a/[/r/Antonym/,/c/en/able/,/c/en/unable/]
|
|
relation = row[1] # e.g., /r/Antonym
|
|
start_node = row[2] # e.g., /c/en/able
|
|
end_node = row[3] # e.g., /c/en/unable
|
|
edge_info = row[4] # JSON with weight, sources, etc.
|
|
|
|
# Parse edge info to get weight (confidence score)
|
|
try:
|
|
info = json.loads(edge_info)
|
|
weight = info.get("weight", 1.0)
|
|
except (json.JSONDecodeError, KeyError):
|
|
weight = 1.0
|
|
|
|
# Convert to RDF-style triple
|
|
triple = {
|
|
"subject": convert_conceptnet_uri_to_full(start_node),
|
|
"predicate": convert_conceptnet_uri_to_full(relation),
|
|
"object": convert_conceptnet_uri_to_full(end_node),
|
|
"object_type": "uri", # ConceptNet edges link concepts (URIs)
|
|
"object_datatype": None,
|
|
"object_language": None,
|
|
# Store weight as metadata - could be used for filtering
|
|
# Note: We're storing this as string to maintain compatibility
|
|
# with the standard RDF format, but it could be parsed as float
|
|
# when needed for ML tasks
|
|
}
|
|
|
|
# For high-confidence assertions only (optional filter)
|
|
# Typical ConceptNet weights range from 0.5 to 20+
|
|
# if weight >= 1.0: # Uncomment to filter low-confidence edges
|
|
triples.append(triple)
|
|
|
|
return triples
|
|
|
|
|
|
def convert_conceptnet(input_file: Path, output_dir: Path) -> None:
|
|
"""Convert ConceptNet CSV to HuggingFace format.
|
|
|
|
Args:
|
|
input_file: Path to ConceptNet CSV file
|
|
output_dir: Output directory for HuggingFace dataset
|
|
"""
|
|
console = Console()
|
|
|
|
console.print("\n[bold cyan]Converting ConceptNet to HuggingFace Dataset[/bold cyan]")
|
|
console.print(f"Input: {input_file}")
|
|
console.print(f"Output: {output_dir}\n")
|
|
|
|
# Emit initial progress
|
|
print("\nPROGRESS: 5", flush=True)
|
|
sys.stdout.flush()
|
|
|
|
if not input_file.exists():
|
|
console.print(f"[red]Error: Input file not found: {input_file}[/red]")
|
|
return
|
|
|
|
# Get file size for progress estimation
|
|
file_size = input_file.stat().st_size
|
|
|
|
all_triples = []
|
|
|
|
with Progress(
|
|
SpinnerColumn(),
|
|
TextColumn("[progress.description]{task.description}"),
|
|
BarColumn(),
|
|
TimeElapsedColumn(),
|
|
console=console,
|
|
transient=False,
|
|
) as progress:
|
|
# Count lines first for accurate progress
|
|
task = progress.add_task("[yellow]Counting assertions...[/yellow]", total=None)
|
|
with open(input_file, "r", encoding="utf-8") as f:
|
|
total_lines = sum(1 for _ in f)
|
|
progress.update(task, description=f"[green]✓ Found {total_lines:,} assertions[/green]")
|
|
|
|
# Emit progress after counting (10%)
|
|
print("\nPROGRESS: 10", flush=True)
|
|
sys.stdout.flush()
|
|
|
|
# Parse CSV file
|
|
task = progress.add_task("[yellow]Parsing ConceptNet assertions...[/yellow]", total=total_lines)
|
|
|
|
# Track progress emission
|
|
last_progress_pct = 10
|
|
|
|
with open(input_file, "r", encoding="utf-8") as f:
|
|
reader = csv.reader(f, delimiter="\t")
|
|
|
|
for i, row in enumerate(reader):
|
|
# Skip malformed rows
|
|
if len(row) < 5:
|
|
progress.advance(task)
|
|
continue
|
|
|
|
# Extract fields
|
|
try:
|
|
assertion_uri = row[0]
|
|
relation = row[1]
|
|
start_node = row[2]
|
|
end_node = row[3]
|
|
edge_info = row[4]
|
|
|
|
# Parse edge info for weight
|
|
try:
|
|
info = json.loads(edge_info)
|
|
weight = info.get("weight", 1.0)
|
|
except (json.JSONDecodeError, KeyError):
|
|
weight = 1.0
|
|
|
|
# Create triple
|
|
triple = {
|
|
"subject": convert_conceptnet_uri_to_full(start_node),
|
|
"predicate": convert_conceptnet_uri_to_full(relation),
|
|
"object": convert_conceptnet_uri_to_full(end_node),
|
|
"object_type": "uri",
|
|
"object_datatype": None,
|
|
"object_language": None,
|
|
}
|
|
|
|
all_triples.append(triple)
|
|
|
|
except (IndexError, ValueError) as e:
|
|
# Skip problematic rows
|
|
pass
|
|
|
|
# Update progress and emit PROGRESS markers
|
|
if i % 10000 == 0:
|
|
progress.update(task, completed=i)
|
|
|
|
# Calculate overall progress (10-80% for parsing phase)
|
|
parse_pct = (i / total_lines) * 100 if total_lines > 0 else 0
|
|
overall_progress = 10 + (parse_pct / 100) * 70
|
|
|
|
# Emit progress every 5% or so
|
|
if overall_progress >= last_progress_pct + 5:
|
|
print(f"\nPROGRESS: {overall_progress:.0f}", flush=True)
|
|
sys.stdout.flush()
|
|
last_progress_pct = overall_progress
|
|
|
|
progress.update(task, completed=total_lines)
|
|
|
|
# Final parsing progress
|
|
print("\nPROGRESS: 80", flush=True)
|
|
sys.stdout.flush()
|
|
|
|
progress.update(task, description=f"[green]✓ Parsed {len(all_triples):,} valid triples[/green]")
|
|
|
|
# Create dataset dictionary with single "data" split
|
|
console.print("[yellow]Creating HuggingFace dataset...[/yellow]")
|
|
print("\nPROGRESS: 85", flush=True)
|
|
sys.stdout.flush()
|
|
|
|
dataset_dict = DatasetDict({"data": Dataset.from_list(all_triples)})
|
|
|
|
print("\nPROGRESS: 90", flush=True)
|
|
sys.stdout.flush()
|
|
|
|
# Add metadata
|
|
dataset_dict._info = {
|
|
"description": "ConceptNet 5.7 - A multilingual knowledge graph of common sense knowledge",
|
|
"citation": """@inproceedings{speer2017conceptnet,
|
|
title={ConceptNet 5.5: An Open Multilingual Graph of General Knowledge},
|
|
author={Speer, Robyn and Chin, Joshua and Havasi, Catherine},
|
|
booktitle={Proceedings of the Thirty-First AAAI Conference on Artificial Intelligence},
|
|
pages={4444--4451},
|
|
year={2017}
|
|
}""",
|
|
"homepage": "http://conceptnet.io",
|
|
"license": "CC BY-SA 4.0",
|
|
}
|
|
|
|
# Save dataset
|
|
console.print("[yellow]Saving dataset...[/yellow]")
|
|
print("\nPROGRESS: 95", flush=True)
|
|
sys.stdout.flush()
|
|
|
|
output_dir.mkdir(parents=True, exist_ok=True)
|
|
dataset_dict.save_to_disk(str(output_dir))
|
|
|
|
print("\nPROGRESS: 100", flush=True)
|
|
sys.stdout.flush()
|
|
|
|
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):,}")
|
|
|
|
# Sample some relations to show variety
|
|
relations = {}
|
|
for triple in all_triples[:10000]: # Sample first 10k for stats
|
|
rel = triple["predicate"]
|
|
relations[rel] = relations.get(rel, 0) + 1
|
|
|
|
console.print(f" • Sample relation types: {len(relations)}")
|
|
console.print("\n[bold]Top 5 relations in sample:[/bold]")
|
|
for rel, count in sorted(relations.items(), key=lambda x: x[1], reverse=True)[:5]:
|
|
rel_name = rel.split("/")[-1] if "/" in rel else rel
|
|
console.print(f" • {rel_name}: {count:,}")
|
|
|
|
|
|
def main():
|
|
"""Main entry point."""
|
|
parser = argparse.ArgumentParser(description="Convert ConceptNet CSV to HuggingFace format")
|
|
parser.add_argument("input", type=Path, help="Input ConceptNet CSV file")
|
|
parser.add_argument("output", type=Path, help="Output directory for HuggingFace dataset")
|
|
|
|
args = parser.parse_args()
|
|
|
|
# Handle both compressed and uncompressed files
|
|
if args.input.suffix == ".gz":
|
|
import gzip
|
|
import tempfile
|
|
|
|
console = Console()
|
|
console.print("[yellow]Decompressing .gz file...[/yellow]")
|
|
|
|
# Create temp file for decompressed content
|
|
with tempfile.NamedTemporaryFile(mode="wb", suffix=".csv", delete=False) as tmp:
|
|
with gzip.open(args.input, "rb") as gz:
|
|
tmp.write(gz.read())
|
|
temp_file = Path(tmp.name)
|
|
|
|
try:
|
|
convert_conceptnet(temp_file, args.output)
|
|
finally:
|
|
# Clean up temp file
|
|
temp_file.unlink()
|
|
else:
|
|
convert_conceptnet(args.input, args.output)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|