dataset-uploader/scripts/convert_conceptnet_to_hf.py

#!/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 gzip
import sys
import tempfile
from pathlib import Path

from datasets import Dataset, DatasetDict
from rich.console import Console
from rich.progress import (
    BarColumn,
    Progress,
    SpinnerColumn,
    TextColumn,
    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 | None]]:
    """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 = []

    try:
        with open(file_path, 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
                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

                # 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)
    except FileNotFoundError:
        console = Console()
        console.print(f"[red]File not found: {file_path}[/red]")
        console.print(
            "[yellow]Make sure the ConceptNet 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


# pylint: disable=too-many-locals,too-many-statements
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]")
        console.print("[yellow]Expected ConceptNet CSV file (tab-separated)[/yellow]")
        return

    # Check if file is readable and not empty
    try:
        file_size = input_file.stat().st_size
        if file_size == 0:
            console.print(f"[red]Error: Input file is empty: {input_file}[/red]")
            return
    except OSError as e:
        console.print(f"[red]Error accessing input file: {e}[/red]")
        return

    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)
        try:
            with open(input_file, encoding="utf-8") as f:
                total_lines = sum(1 for _ in f)
            progress.update(
                task, description=f"[green]✓ Found {total_lines:,} assertions[/green]"
            )
        except OSError as e:
            console.print(f"[red]Error counting lines: {e}[/red]")
            return

        # 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.0

        with open(input_file, 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:
                    relation = row[1]
                    start_node = row[2]
                    end_node = row[3]

                    # 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):
                    # 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()

        if not all_triples:
            console.print(
                "[yellow]Warning: No valid triples found in input file[/yellow]"
            )
            console.print("[yellow]Check file format and content[/yellow]")
        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()

    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

    print("\nPROGRESS: 90", flush=True)
    sys.stdout.flush()

    # Add metadata
    # pylint: disable=protected-access
    # Add metadata
    metadata_dict = {
        "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",
    }

    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
    console.print("[yellow]Saving dataset...[/yellow]")
    print("\nPROGRESS: 95", flush=True)
    sys.stdout.flush()

    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

    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: dict[str, int] = {}
    for triple in all_triples[:10000]:  # Sample first 10k for stats
        rel = triple["predicate"]
        assert rel is not None
        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() -> None:
    """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":
        console = Console()
        console.print("[yellow]Decompressing .gz file...[/yellow]")

        try:
            # 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:
                    decompressed_data = gz.read()
                    if isinstance(decompressed_data, str):
                        decompressed_data = decompressed_data.encode("utf-8")
                    tmp.write(decompressed_data)
                temp_file = Path(tmp.name)
        except gzip.BadGzipFile as e:
            console.print(f"[red]Invalid or corrupted gzip file: {e}[/red]")
            console.print(f"[yellow]File: {args.input}[/yellow]")
            console.print("[yellow]Please re-download the dataset[/yellow]")
            return
        except OSError as e:
            console.print(f"[red]Error decompressing file: {e}[/red]")
            console.print(f"[yellow]File: {args.input}[/yellow]")
            return

        try:
            convert_conceptnet(temp_file, args.output)
        finally:
            # Clean up temp file
            try:
                temp_file.unlink()
            except OSError:
                console.print(
                    f"[yellow]Warning: Could not remove temporary file: "
                    f"{temp_file}[/yellow]"
                )
    else:
        convert_conceptnet(args.input, args.output)


if __name__ == "__main__":
    main()