From 1bb26305ce0d55cf1880b547d4af8217cc51b9dc Mon Sep 17 00:00:00 2001 From: Aditya Chhabra Date: Fri, 26 Dec 2025 18:21:02 +0530 Subject: [PATCH 1/8] fix: fix path too long issue; fix pickling issue in unified converter script --- scripts/convert_rdf_to_hf_dataset_unified.py | 411 ++++++++++++------- scripts/upload_all_datasets.py | 5 + 2 files changed, 274 insertions(+), 142 deletions(-) diff --git a/scripts/convert_rdf_to_hf_dataset_unified.py b/scripts/convert_rdf_to_hf_dataset_unified.py index aed5e89..934d0de 100644 --- a/scripts/convert_rdf_to_hf_dataset_unified.py +++ b/scripts/convert_rdf_to_hf_dataset_unified.py @@ -374,6 +374,216 @@ def batch_file_lines(file_path: Path, batch_size: int = 1000, format: str = "geo finally: file_obj.close() + +def stream_geonames_parallel( + file_path: Path, chunk_size: int = 10000, num_workers: int | None = None +) -> Iterator[list[RDFTriple]]: + """Stream GeoNames RDF file with parallel processing. + + Module-level function for use in picklable generators. + """ + if num_workers is None: + num_workers = max(1, mp.cpu_count() - 1) + + batches = batch_file_lines(file_path, batch_size=chunk_size, format="geonames") + + try: + with mp.Pool(processes=num_workers) as pool: + for triples in pool.imap_unordered(process_geonames_lines, batches, chunksize=1): + if triples: + yield triples + except Exception as e: + logger.error(f"Error processing GeoNames file: {e}") + raise + + +def stream_ntriples_parallel( + file_path: Path, chunk_size: int = 10000, num_workers: int | None = None +) -> Iterator[list[RDFTriple]]: + """Stream N-Triples file with parallel processing. + + Module-level function for use in picklable generators. + """ + if num_workers is None: + num_workers = max(1, mp.cpu_count() - 1) + + batches = batch_file_lines(file_path, batch_size=chunk_size, format="ntriples") + + try: + with mp.Pool(processes=num_workers) as pool: + for triples in pool.imap_unordered(process_ntriples_lines, batches, chunksize=1): + if triples: + yield triples + except Exception as e: + logger.error(f"Error processing N-Triples file: {e}") + raise + + +def stream_turtle_chunks(file_path: Path, chunk_size: int = 10000) -> Iterator[list[RDFTriple]]: + """Stream Turtle file in chunks using incremental parsing. + + Module-level function for use in picklable generators. + """ + is_gzipped = file_path.suffix == ".gz" + is_bz2 = file_path.suffix == ".bz2" or str(file_path).endswith(".ttl.bz2") + + if is_bz2: + import bz2 + file_obj = bz2.open(file_path, "rt", encoding="utf-8", errors="ignore") + elif is_gzipped: + import gzip + file_obj = gzip.open(file_path, "rt", encoding="utf-8", errors="ignore") + else: + file_obj = open(file_path, encoding="utf-8", errors="ignore") + + try: + current_chunk = [] + triple_count = 0 + prefix_lines = [] + in_prefixes = True + + for line in file_obj: + stripped = line.strip() + if not stripped or stripped.startswith("#"): + continue + + if in_prefixes and (stripped.startswith("@prefix") or stripped.startswith("@base")): + prefix_lines.append(line) + continue + elif in_prefixes: + in_prefixes = False + + if stripped.endswith("."): + triple_count += 1 + + if triple_count >= chunk_size: + chunk_text = "".join(prefix_lines) + "\n" + "".join(current_chunk) + "\n" + line + + try: + graph = Graph() + graph.parse(data=chunk_text, format="turtle") + + triples = [] + for s, p, o in graph: + if isinstance(o, Literal): + object_type = "literal" + object_datatype = str(o.datatype) if o.datatype else None + object_language = o.language if o.language else None + elif isinstance(o, URIRef): + object_type = "uri" + object_datatype = None + object_language = None + else: + object_type = "blank_node" + object_datatype = None + object_language = None + + triple = { + "subject": str(s), + "predicate": str(p), + "object": str(o), + "object_type": object_type, + "object_datatype": object_datatype, + "object_language": object_language, + } + triples.append(triple) + + if triples: + yield triples + + except Exception as e: + logger.error(f"Error parsing Turtle chunk: {e}") + + current_chunk = [] + triple_count = 0 + else: + current_chunk.append(line) + else: + current_chunk.append(line) + + # Process remaining lines + if current_chunk: + chunk_text = "".join(prefix_lines) + "\n" + "".join(current_chunk) + try: + graph = Graph() + graph.parse(data=chunk_text, format="turtle") + + triples = [] + for s, p, o in graph: + if isinstance(o, Literal): + object_type = "literal" + object_datatype = str(o.datatype) if o.datatype else None + object_language = o.language if o.language else None + elif isinstance(o, URIRef): + object_type = "uri" + object_datatype = None + object_language = None + else: + object_type = "blank_node" + object_datatype = None + object_language = None + + triple = { + "subject": str(s), + "predicate": str(p), + "object": str(o), + "object_type": object_type, + "object_datatype": object_datatype, + "object_language": object_language, + } + triples.append(triple) + + if triples: + yield triples + + except Exception as e: + logger.error(f"Error parsing final Turtle chunk: {e}") + + finally: + file_obj.close() + + +def stream_generic_rdf(file_path: Path, rdf_format: str, chunk_size: int = 10000) -> Iterator[list[RDFTriple]]: + """Stream generic RDF file using standard parsing. + + Module-level function for use in picklable generators. + """ + graph = Graph() + graph.parse(str(file_path), format=rdf_format) + + current_chunk = [] + for s, p, o in graph: + if isinstance(o, Literal): + object_type = "literal" + object_datatype = str(o.datatype) if o.datatype else None + object_language = o.language if o.language else None + elif isinstance(o, URIRef): + object_type = "uri" + object_datatype = None + object_language = None + else: + object_type = "blank_node" + object_datatype = None + object_language = None + + triple = { + "subject": str(s), + "predicate": str(p), + "object": str(o), + "object_type": object_type, + "object_datatype": object_datatype, + "object_language": object_language, + } + current_chunk.append(triple) + + if len(current_chunk) >= chunk_size: + yield current_chunk + current_chunk = [] + + if current_chunk: + yield current_chunk + + class MetadataHandler: """Handle metadata for the HuggingFace dataset.""" @staticmethod @@ -555,138 +765,35 @@ class ConversionStrategy(ABC): raise def _stream_generic_rdf(self, config: ConversionConfig) -> Iterator[list[RDFTriple]]: - graph = Graph() - try: - graph.parse(str(config.input_path), format=config.rdf_format) - except Exception as e: - logger.error(f"Error parsing RDF file: {e}") - raise - current_chunk: list[RDFTriple] = [] - for s, p, o in graph: - current_chunk.append(extract_triple(s, p, o)) - if len(current_chunk) >= config.chunk_size: - yield current_chunk - current_chunk = [] - if current_chunk: - yield current_chunk + """Stream generic RDF file using standard parsing. + + Thin wrapper that delegates to module-level function. + """ + return stream_generic_rdf(config.input_path, config.rdf_format, config.chunk_size) def _stream_geonames_parallel(self, config: ConversionConfig, num_workers: int) -> Iterator[list[RDFTriple]]: - """Stream GeoNames RDF file with parallel processing.""" + """Stream GeoNames RDF file with parallel processing. + + Thin wrapper that prints progress then delegates to module-level function. + """ self.progress.print(f"[yellow]Using parallel GeoNames parser with {num_workers} workers[/yellow]") - - # Create batches of lines - batches = batch_file_lines(config.input_path, batch_size=config.chunk_size, format="geonames") - - # Process batches in parallel - try: - with mp.Pool(processes=num_workers) as pool: - for triples in pool.imap_unordered(process_geonames_lines, batches, chunksize=1): - if triples: - yield triples - except Exception as e: - logger.error(f"Error processing GeoNames file: {e}") - raise + return stream_geonames_parallel(config.input_path, config.chunk_size, num_workers) def _stream_ntriples_parallel(self, config: ConversionConfig, num_workers: int) -> Iterator[list[RDFTriple]]: - """Stream N-Triples file with parallel processing.""" + """Stream N-Triples file with parallel processing. + + Thin wrapper that prints progress then delegates to module-level function. + """ self.progress.print(f"[yellow]Using parallel N-Triples parser with {num_workers} workers[/yellow]") - - # Create batches of lines - batches = batch_file_lines(config.input_path, batch_size=config.chunk_size, format="ntriples") - - # Process batches in parallel - try: - with mp.Pool(processes=num_workers) as pool: - for triples in pool.imap_unordered(process_ntriples_lines, batches, chunksize=1): - if triples: - yield triples - except Exception as e: - logger.error(f"Error processing N-Triples file: {e}") - raise + return stream_ntriples_parallel(config.input_path, config.chunk_size, num_workers) def _stream_turtle_chunks(self, config: ConversionConfig) -> Iterator[list[RDFTriple]]: - """Stream Turtle file in chunks using incremental parsing.""" + """Stream Turtle file in chunks using incremental parsing. + + Thin wrapper that prints progress then delegates to module-level function. + """ self.progress.print("[yellow]Streaming Turtle file (line-by-line parser)[/yellow]") - - # Open file with compression support - file_obj = self.file_handler.open_file(config.input_path) - - try: - current_chunk: list[str] = [] - triple_count = 0 - line_count = 0 - - # Collect prefixes first - prefix_lines: list[str] = [] - in_prefixes = True - - for line_bytes in file_obj: - line_count += 1 - # Ensure we have a string - line = line_bytes if isinstance(line_bytes, str) else line_bytes.decode('utf-8', errors='ignore') - stripped = line.strip() - - # Skip empty lines and comments - if not stripped or stripped.startswith("#"): - continue - - # Collect prefix declarations - if in_prefixes and (stripped.startswith("@prefix") or stripped.startswith("@base")): - prefix_lines.append(line) - continue - elif in_prefixes: - # End of prefixes, now we're in the data - in_prefixes = False - - # Parse in batches of lines ending with '.' - if stripped.endswith("."): - triple_count += 1 - - if triple_count >= config.chunk_size: - # Try to parse this chunk - chunk_text = "".join(prefix_lines) + "\n" + "".join(current_chunk) + "\n" + line - - try: - graph = Graph() - graph.parse(data=chunk_text, format="turtle") - - # Extract triples - triples = [extract_triple(s, p, o) for s, p, o in graph] - if triples: - yield triples - - except Exception as e: - logger.error(f"Error parsing Turtle chunk: {e}") - - # Reset for next chunk - current_chunk = [] - triple_count = 0 - else: - current_chunk.append(line) - else: - # Part of a multi-line statement - current_chunk.append(line) - - # Log progress - if line_count % 100000 == 0: - self.progress.print(f"[dim]Processed {line_count:,} lines...[/dim]") - - # Process remaining lines - if current_chunk: - chunk_text = "".join(prefix_lines) + "\n" + "".join(current_chunk) - try: - graph = Graph() - graph.parse(data=chunk_text, format="turtle") - - triples = [extract_triple(s, p, o) for s, p, o in graph] - if triples: - yield triples - - except Exception as e: - logger.error(f"Error parsing final Turtle chunk: {e}") - - finally: - file_obj.close() + return stream_turtle_chunks(config.input_path, config.chunk_size) # ============================================================================= # CONCRETE STRATEGIES @@ -987,31 +1094,50 @@ class ParallelStreamingStrategy(ConversionStrategy): # Estimate total chunks based on file size (rough estimate) estimated_chunks = max(10, int(file_size_mb * 1024 * 1024 / (config.chunk_size * 100))) - # Create generator function that yields individual triples + # Determine format and print status messages BEFORE creating generator + # (Cannot print from inside generator due to pickling constraints) + format_type = None + if is_geonames and config.rdf_format in ("xml", "application/rdf+xml"): + # Check if it's the special GeoNames format (URLs followed by XML) + with open(config.input_path, encoding="utf-8", errors="ignore") as f: + first_line = f.readline().strip() + if first_line.startswith("http://") or first_line.startswith("https://"): + format_type = "geonames" + self.progress.print(f"[yellow]Using parallel GeoNames parser with {num_workers} workers[/yellow]") + else: + format_type = "generic" + self.progress.print(f"[yellow]Using standard RDF parser for {config.rdf_format} (single-threaded)[/yellow]") + elif config.rdf_format in ("nt", "ntriples"): + format_type = "ntriples" + self.progress.print(f"[yellow]Using parallel N-Triples parser with {num_workers} workers[/yellow]") + elif config.rdf_format in ("turtle", "ttl"): + format_type = "turtle" + self.progress.print("[yellow]Streaming Turtle file (line-by-line parser)[/yellow]") + else: + format_type = "generic" + self.progress.print(f"[yellow]Using standard RDF parser for {config.rdf_format} (single-threaded)[/yellow]") + + # Note: Cannot capture self or generators in dataset_generator closure. + # Dataset.from_generator() requires picklable generators. + # Must call streaming functions directly inside the generator. def dataset_generator(): triple_count = 0 chunk_count = 0 last_print_count = 0 - # Select appropriate streaming method based on format - if is_geonames and config.rdf_format in ("xml", "application/rdf+xml"): - # Check if it's the special GeoNames format (URLs followed by XML) - with open(config.input_path, encoding="utf-8", errors="ignore") as f: - first_line = f.readline().strip() - if first_line.startswith("http://") or first_line.startswith("https://"): - chunk_iter = self._stream_geonames_parallel(config, num_workers) - else: - chunk_iter = self._stream_generic_rdf(config) - elif config.rdf_format in ("nt", "ntriples"): - chunk_iter = self._stream_ntriples_parallel(config, num_workers) - elif config.rdf_format in ("turtle", "ttl"): - chunk_iter = self._stream_turtle_chunks(config) + # Create the chunk iterator inside the generator to avoid pickling issues + if format_type == "geonames": + chunk_source = stream_geonames_parallel(config.input_path, config.chunk_size, num_workers) + elif format_type == "ntriples": + chunk_source = stream_ntriples_parallel(config.input_path, config.chunk_size, num_workers) + elif format_type == "turtle": + chunk_source = stream_turtle_chunks(config.input_path, config.chunk_size) else: - # For other formats, use generic streaming - chunk_iter = self._stream_generic_rdf(config) + # Generic RDF parsing + chunk_source = stream_generic_rdf(config.input_path, config.rdf_format, config.chunk_size) # Yield individual triples from chunks - for chunk in chunk_iter: + for chunk in chunk_source: chunk_count += 1 for triple in chunk: triple_count += 1 @@ -1021,11 +1147,12 @@ class ParallelStreamingStrategy(ConversionStrategy): elapsed = time.time() - start_time rate = triple_count / elapsed if elapsed > 0 else 0 progress_pct = 10 + min(60, int((chunk_count / estimated_chunks) * 60)) - self.progress.print( + print( f" Processing: {chunk_count:,} chunks • " f"{triple_count:,} triples • {rate:.0f} triples/sec", + flush=True, ) - self.progress.emit_progress(progress_pct) + print(f"PROGRESS: {progress_pct}", flush=True) last_print_count = triple_count # Create dataset using from_generator for true streaming diff --git a/scripts/upload_all_datasets.py b/scripts/upload_all_datasets.py index 65b4b6a..5fd481d 100755 --- a/scripts/upload_all_datasets.py +++ b/scripts/upload_all_datasets.py @@ -55,6 +55,7 @@ DIRECTORY STRUCTURE: from __future__ import annotations import argparse +import os import shutil import subprocess import sys @@ -1203,6 +1204,10 @@ def main() -> int: """ args, parser = parse_args() + # Fix for "AF_UNIX path too long" error in multiprocessing + # This forces the temporary directory to be /tmp (short path) instead of a potentially deep workspace path + os.environ["TMPDIR"] = "/tmp" + start_time = time.monotonic() # List datasets to process -- 2.52.0 From 1d33b67198fab4fe9d9346e576a78178d0355d5d Mon Sep 17 00:00:00 2001 From: Brent Edwards Date: Fri, 26 Dec 2025 18:49:33 -0800 Subject: [PATCH 2/8] A few minor fixes. --- scripts/convert_rdf_to_hf_dataset.py | 1292 ----------------- .../convert_rdf_to_hf_dataset_streaming.py | 720 --------- ...rt_rdf_to_hf_dataset_streaming_parallel.py | 862 ----------- ...vert_rdf_to_hf_dataset_streaming_simple.py | 672 --------- ...vert_rdf_to_hf_dataset_streaming_turtle.py | 496 ------- scripts/convert_rdf_to_hf_dataset_unified.py | 11 +- 6 files changed, 8 insertions(+), 4045 deletions(-) delete mode 100755 scripts/convert_rdf_to_hf_dataset.py delete mode 100755 scripts/convert_rdf_to_hf_dataset_streaming.py delete mode 100755 scripts/convert_rdf_to_hf_dataset_streaming_parallel.py delete mode 100755 scripts/convert_rdf_to_hf_dataset_streaming_simple.py delete mode 100755 scripts/convert_rdf_to_hf_dataset_streaming_turtle.py diff --git a/scripts/convert_rdf_to_hf_dataset.py b/scripts/convert_rdf_to_hf_dataset.py deleted file mode 100755 index 932c953..0000000 --- a/scripts/convert_rdf_to_hf_dataset.py +++ /dev/null @@ -1,1292 +0,0 @@ -#!/usr/bin/env python3 -"""Convert RDF datasets to HuggingFace dataset format. - -This script parses RDF files in various formats (Turtle, N-Triples, RDF/XML, etc.) -and converts them to HuggingFace dataset format with lossless preservation of all -semantic information. Part of the CleverErnie dataset preparation pipeline. - -USAGE: - # Basic conversion - python scripts/convert_rdf_to_hf_dataset.py [options] - - # Examples - Different RDF formats - python scripts/convert_rdf_to_hf_dataset.py \ - raw_datasets/wordnet/english-wordnet-2024.ttl \ - hf_datasets/wordnet \ - --format turtle - - python scripts/convert_rdf_to_hf_dataset.py \ - raw_datasets/dbpedia/mappingbased_literals_en.nt \ - hf_datasets/dbpedia \ - --format nt - - python scripts/convert_rdf_to_hf_dataset.py \ - raw_datasets/schema/schemaorg.rdf \ - hf_datasets/schema \ - --format xml - - # With metadata - python scripts/convert_rdf_to_hf_dataset.py \ - raw_datasets/wordnet/english-wordnet-2024.ttl \ - hf_datasets/wordnet \ - --format turtle \ - --description "English WordNet 2024 lexical database" \ - --homepage "https://en-word.net/" \ - --license "CC BY 4.0" \ - --citation "English WordNet 2024 (Fellbaum et al.)" - - # Verbose mode (show library warnings) - python scripts/convert_rdf_to_hf_dataset.py input.ttl output/ --format turtle -v - -ARGUMENTS: - input Path to input RDF file (any supported format) - output Path to output directory for HuggingFace dataset - -f, --format FORMAT RDF format: turtle, nt, ntriples, xml, n3, trig, nquads - (default: turtle) - --description TEXT Dataset description (added to metadata) - --citation TEXT Dataset citation (added to metadata) - --homepage URL Dataset homepage URL (added to metadata) - --license TEXT Dataset license (added to metadata) - -v, --verbose Show verbose output including library warnings - (default: only errors) - -RDF FORMATS SUPPORTED: - turtle .ttl Terse RDF Triple Language (human-readable) - nt .nt N-Triples (line-based, fast parsing) - ntriples .nt N-Triples (alias) - xml .rdf RDF/XML (XML-based, verbose) - n3 .n3 Notation3 (superset of Turtle) - trig .trig TriG (Turtle with named graphs) - nquads .nq N-Quads (N-Triples with named graphs) - -OUTPUT SCHEMA: - The script creates a HuggingFace dataset with the following structure: - { - "subject": "http://example.org/resource/123", - "predicate": "http://www.w3.org/1999/02/22-rdf-syntax-ns#type", - "object": "http://example.org/class/Person", - "object_type": "uri", # uri, literal, or blank_node - "object_datatype": "xsd:string", # XSD datatype or None - "object_language": "en" # Language tag or None - } - -LOSSLESS CONVERSION: - All semantic information from the original RDF is preserved: - ✓ All triple relationships (subject-predicate-object) - ✓ URI references and blank nodes - ✓ Literal datatypes (strings, integers, dates, etc.) - ✓ Language tags for multilingual text - ✓ RDF namespace prefixes (expanded to full URIs) - - Round-trip guarantee: RDF → HuggingFace → RDF produces semantically - identical graphs. - -PERFORMANCE: - Parsing speed varies by format: - - N-Triples: ~50K triples/second (fastest) - - Turtle: ~10-20K triples/second - - RDF/XML: ~5-10K triples/second (slowest) - - Large datasets (>1M triples) may take several minutes to hours. - The script provides progress indicators for all phases: - 1. Reading file from disk - 2. Parsing RDF syntax and building graph - 3. Converting triples to dataset format - -OUTPUT STRUCTURE: - hf_datasets/wordnet/ - ├── data/ # Dataset split directory - │ ├── train-00000-of-00001.parquet - │ └── ... - ├── dataset_info.json # HuggingFace dataset metadata - └── state.json # Dataset state information - -REQUIREMENTS: - pip install rdflib datasets rich - -NOTES: - - Large datasets require significant RAM (2-3x uncompressed file size) - - Parsing is the most time-consuming step (can take hours for large datasets) - - Default mode suppresses library warnings for clean output (use -v for debugging) - - Compressed files (.gz, .bz2) are supported automatically - -RELATED: - - rdf_dataset_downloader.py: Download raw RDF datasets - - upload_all_datasets.py: Full pipeline (download + convert + upload) - - dataset_registry.py: Registry of available datasets - -For detailed documentation, see: docs/rdf_dataset_scripts.md -""" - -from __future__ import annotations - -import argparse -import bz2 -import gzip -import logging -import multiprocessing as mp -import sys -import time -from contextlib import contextmanager -from io import BytesIO -from multiprocessing import Pool -from pathlib import Path -from typing import Any - -from datasets import Dataset, DatasetDict, Features, Value -from rdflib import Graph, Literal, URIRef -from rich.console import Console -from rich.progress import ( - BarColumn, - Progress, - SpinnerColumn, - TextColumn, - TimeElapsedColumn, - TimeRemainingColumn, -) - -# Logger will be configured in main() based on --verbose flag -logger = logging.getLogger(__name__) - - -def parse_tsv_triples(file_path: Path) -> list[dict[str, str]]: - """Parse TSV triples from benchmark datasets like WN18RR. - - Args: - file_path: Path to TSV file with tab-separated triples - - Returns: - List of triple dictionaries - """ - console = Console() - triples = [] - - try: - with Progress( - SpinnerColumn(), - TextColumn("[progress.description]{task.description}"), - BarColumn(), - TextColumn("[progress.percentage]{task.percentage:>3.0f}%"), - TimeElapsedColumn(), - console=console, - transient=False, - ) as progress: - # Count lines for progress bar - try: - with open(file_path, encoding="utf-8") as f: - total_lines = sum(1 for _ in f) - except OSError as e: - console.print(f"[red]Error reading file for line count: {e}[/red]") - raise - - task = progress.add_task( - "[yellow]Parsing TSV triples...[/yellow]", - total=total_lines, - ) - - with open(file_path, encoding="utf-8") as f: - for i, line in enumerate(f): - line = line.strip() - if not line: - continue - - parts = line.split("\t") - if len(parts) != 3: - continue - - subject, predicate, obj = parts - triple = { - "subject": subject.strip(), - "predicate": predicate.strip(), - "object": obj.strip(), - "object_type": "literal", - "object_datatype": None, - "object_language": None, - } - triples.append(triple) - - if i % 1000 == 0: - progress.update(task, advance=1000) - - progress.update(task, completed=total_lines) - except FileNotFoundError: - console.print(f"[red]TSV file not found: {file_path}[/red]") - console.print( - "[yellow]Make sure the dataset has been downloaded correctly[/yellow]" - ) - raise - except OSError as e: - console.print(f"[red]Error reading TSV file: {e}[/red]") - raise - - console.print(f"[green]✓ Parsed {len(triples):,} TSV triples[/green]\n") - return triples - - -class ProgressFileReader: - """File reader wrapper that tracks read progress for progress bar.""" - - def __init__(self, file_obj, progress_callback, total_size): - """Initialize progress file reader. - - Args: - file_obj: File object to wrap - progress_callback: Callback function(bytes_read) to update progress - total_size: Total file size in bytes - """ - self.file_obj = file_obj - self.progress_callback = progress_callback - self.total_size = total_size - self.bytes_read = 0 - - def read(self, size=-1): - """Read from file and update progress.""" - data = self.file_obj.read(size) - if data: - self.bytes_read += len(data) - self.progress_callback(self.bytes_read) - return data - - def readline(self, size=-1): - """Read line from file and update progress.""" - data = self.file_obj.readline(size) - if data: - self.bytes_read += len(data) - self.progress_callback(self.bytes_read) - return data - - def __enter__(self): - """Context manager entry.""" - return self - - def __exit__(self, exc_type, exc_val, exc_tb): - """Context manager exit.""" - if hasattr(self.file_obj, "__exit__"): - return self.file_obj.__exit__(exc_type, exc_val, exc_tb) - return False - - def __getattr__(self, name): - """Delegate other attributes to wrapped file object.""" - return getattr(self.file_obj, name) - - -def get_file_chunks(file_path: Path, num_workers: int) -> list[tuple[int, int]]: - """Divide file into byte-range chunks for parallel reading. - - Returns list of (start_byte, end_byte) tuples, one per worker. - Each chunk is aligned to document boundaries. - - Args: - file_path: Path to file to chunk - num_workers: Number of chunks to create - - Returns: - List of (start_byte, end_byte) tuples - """ - file_size = file_path.stat().st_size - chunk_size = file_size // num_workers - - chunks = [] - - with open(file_path, "rb") as f: - for i in range(num_workers): - start = i * chunk_size - - if i == num_workers - 1: - # Last chunk gets remainder - end = file_size - else: - # Seek to approximate end position - end = start + chunk_size - f.seek(end) - - # Read forward until we find a URL line (document boundary) - # This ensures we split between documents, not mid-document - while end < file_size: - line = f.readline() - end = f.tell() - decoded = line.decode("utf-8", errors="ignore") - if decoded.startswith("http://") or decoded.startswith("https://"): - break - - chunks.append((start, end)) - - return chunks - - -def parse_file_chunk(args: tuple) -> tuple[list[dict[str, str]], int]: - """Worker function: parse documents from a byte-range chunk of the file. - - Each worker independently: - 1. Opens the file - 2. Seeks to its start position - 3. Streams and parses documents until end position - 4. Returns all triples from its chunk - - Args: - args: (file_path, start_byte, end_byte, worker_id, progress_dict, lock) - - Returns: - Tuple of (list of triple dictionaries, worker_id) - """ - from rdflib import Graph, Literal, URIRef - - file_path, start_byte, end_byte, worker_id, progress_dict, lock = args - - triples = [] - current_xml = [] - bytes_read = start_byte - doc_count = 0 - chunk_size = end_byte - start_byte - - with open(file_path, encoding="utf-8") as f: - # Seek to start position - f.seek(start_byte) - - # If not at file start, read to next document boundary - if start_byte > 0: - # Skip partial document at beginning - while bytes_read < end_byte: - line = f.readline() - if not line: - break - bytes_read += len(line.encode("utf-8")) - if line.startswith("http://") or line.startswith("https://"): - break - - # Process documents in this chunk - last_report = bytes_read - while bytes_read < end_byte: - line = f.readline() - if not line: - break - - bytes_read += len(line.encode("utf-8")) - - # Report progress every 10MB - if bytes_read - last_report > 10 * 1024 * 1024: - progress_pct = ((bytes_read - start_byte) / chunk_size) * 100 - with lock: - progress_dict[worker_id] = progress_pct - last_report = bytes_read - - # Document boundary - if line.startswith("http://") or line.startswith("https://"): - if current_xml: - # Parse accumulated document - xml_str = "".join(current_xml) - try: - graph = Graph() - graph.parse(data=xml_str, format="xml") - - for s, p, o in graph: - # Extract object metadata like the main parser does - if isinstance(o, Literal): - object_type = "literal" - object_datatype = ( - str(o.datatype) if o.datatype else None - ) - object_language = o.language if o.language else None - elif isinstance(o, URIRef): - object_type = "uri" - object_datatype = None - object_language = None - else: - object_type = "blank_node" - object_datatype = None - object_language = None - - triples.append( - { - "subject": str(s), - "predicate": str(p), - "object": str(o), - "object_type": object_type, - "object_datatype": object_datatype, - "object_language": object_language, - } - ) - - doc_count += 1 - except Exception: - # Skip malformed documents - pass - - current_xml = [] - continue - - current_xml.append(line) - - # Parse final document in chunk - if current_xml: - xml_str = "".join(current_xml) - try: - graph = Graph() - graph.parse(data=xml_str, format="xml") - - for s, p, o in graph: - # Extract object metadata like the main parser does - if isinstance(o, Literal): - object_type = "literal" - object_datatype = str(o.datatype) if o.datatype else None - object_language = o.language if o.language else None - elif isinstance(o, URIRef): - object_type = "uri" - object_datatype = None - object_language = None - else: - object_type = "blank_node" - object_datatype = None - object_language = None - - triples.append( - { - "subject": str(s), - "predicate": str(p), - "object": str(o), - "object_type": object_type, - "object_datatype": object_datatype, - "object_language": object_language, - } - ) - - doc_count += 1 - except Exception: - pass - - # Mark this worker as 100% complete - with lock: - progress_dict[worker_id] = 100.0 - - return (triples, worker_id) - - -def parse_geonames_file( - file_path: Path, num_workers: int | None = None -) -> list[dict[str, str]]: - """Parse GeoNames RDF file with streaming parallel processing. - - GeoNames format has lines like "https://sws.geonames.org/3/" before each " - "XML document. - Each XML document is a complete RDF/XML graph that needs to be parsed separately. - - Strategy: - - Divide file into N byte-range chunks (one per CPU core) - - Each worker independently streams and parses its chunk - - No memory bottleneck - each worker processes incrementally - - Pure CPU parallelization with disk I/O in parallel - - Args: - file_path: Path to GeoNames RDF file - num_workers: Number of worker processes (default: CPU count) - - Returns: - List of triple dictionaries - """ - console = Console() - - if num_workers is None: - num_workers = mp.cpu_count() - - file_size_bytes = file_path.stat().st_size - file_size_mb = file_size_bytes / (1024 * 1024) - - console.print( - f"[yellow]Parsing GeoNames RDF file (streaming parallel with " - f"{num_workers} workers)...[/yellow]" - ) - console.print(f"[dim]File size: {file_size_mb:.2f} MB[/dim]\n") - - # Phase 1: Calculate file chunks (fast) - console.print("[cyan]Phase 1: Dividing file into chunks...[/cyan]") - chunks = get_file_chunks(file_path, num_workers) - - for i, (start, end) in enumerate(chunks): - chunk_mb = (end - start) / (1024 * 1024) - console.print( - f"[dim] Worker {i}: bytes {start:,} to {end:,} ({chunk_mb:.1f} MB)[/dim]" - ) - - console.print(f"[green]✓ File divided into {len(chunks)} chunks[/green]\n") - - # Phase 2: Parse chunks in parallel - console.print( - f"[cyan]Phase 2: Parsing chunks with {num_workers} parallel workers...[/cyan]" - ) - console.print("[dim]Progress updates every 10MB per worker[/dim]\n") - - # Create shared progress tracking - from multiprocessing import Manager - - manager = Manager() - progress_dict = manager.dict() - lock = manager.Lock() - - # Initialize progress for all workers - for i in range(num_workers): - progress_dict[i] = 0.0 - - # Prepare worker arguments with shared progress - worker_args = [ - (file_path, start, end, i, progress_dict, lock) - for i, (start, end) in enumerate(chunks) - ] - - all_triples = [] - - with Progress( - SpinnerColumn(), - TextColumn("[progress.description]{task.description}"), - BarColumn(), - TextColumn("[progress.percentage]{task.percentage:>3.0f}%"), - TimeElapsedColumn(), - console=console, - transient=False, - refresh_per_second=4, # Update display frequently - ) as progress: - task = progress.add_task("[yellow]Processing chunks...[/yellow]", total=100.0) - - with Pool(processes=num_workers) as pool: - # Start all workers - result = pool.map_async(parse_file_chunk, worker_args) - - # Poll for progress updates - import time - - while not result.ready(): - time.sleep(0.25) # Poll 4 times per second - - # Calculate overall progress (average of all workers) - with lock: - if progress_dict: - avg_progress = sum(progress_dict.values()) / len(progress_dict) - progress.update( - task, - completed=avg_progress, - description=( - f"[yellow]Processing chunks (avg " - f"{avg_progress:.1f}%)...[/yellow]" - ), - ) - - # Get all results - worker_results = result.get() - - # Collect all triples (results are now tuples of (triples, worker_id)) - for triples, _worker_id in worker_results: - all_triples.extend(triples) - - # Update to 100% complete - progress.update( - task, - completed=100.0, - description=( - f"[green]✓ All workers complete " - f"({len(all_triples):,} triples)[/green]" - ), - ) - - console.print( - f"\n[green]✓ Successfully extracted {len(all_triples):,} triples[/green]\n" - ) - return all_triples - - -def parse_rdf_file(file_path: Path, format: str = "turtle") -> list[dict[str, str]]: - """Parse an RDF file and extract triples. - - Args: - file_path: Path to RDF file - format: RDF format (turtle, nt, xml, etc.) - - Returns: - List of triple dictionaries - - Progress stages (for wrapper scripts): - 0-10%: Initialization and format detection - 10-50%: Parsing RDF file to graph - 50-90%: Converting graph to dataset format - 90-100%: Final preparation - """ - console = Console() - - # Get file size for status messages - try: - file_size_bytes = file_path.stat().st_size - file_size_mb = file_size_bytes / (1024 * 1024) - except OSError as e: - console.print(f"[red]Error accessing file: {e}[/red]") - console.print(f"[yellow]File: {file_path}[/yellow]") - raise - - console.print(f"[cyan]Parsing RDF file: {file_path}[/cyan]") - # Stage 1: Initialization (0-10%) - # Don't emit progress here, let the wrapper handle it - - # Handle TSV/TXT benchmark datasets like WN18RR - if format == "tsv" or ( - file_path.suffix in {".txt", ".tsv"} and "train" in file_path.name - ): - console.print( - "[yellow]Detected TSV/benchmark format. Parsing as " - "tab-separated triples.[/yellow]" - ) - console.print(f"[dim]File size: {file_size_mb:.2f} MB[/dim]\n") - # Don't emit progress here, it conflicts with the wrapper - return parse_tsv_triples(file_path) - - console.print(f"[dim]File size: {file_size_mb:.2f} MB ({format} format)[/dim]\n") - - # Check if this is a GeoNames file (special format with URL prefixes) - is_geonames = "geonames" in str(file_path).lower() - - if is_geonames and format in ("xml", "application/rdf+xml"): - # Check first few lines to confirm GeoNames format - with open(file_path, encoding="utf-8") as f: - first_line = f.readline().strip() - if first_line.startswith("http://") or first_line.startswith("https://"): - # This is GeoNames format, use special parser - return parse_geonames_file(file_path) - - graph = Graph() - - # Parse with progress bar showing file read progress - start_time = time.time() - - # For XML/RDF formats, parse directly from file - # (BytesIO causes format detection issues) - # For other formats, we can use BytesIO with progress tracking - if format in ("xml", "application/rdf+xml"): - console.print("[dim]Phase 1+2: Parsing XML/RDF directly from file[/dim]") - console.print( - "[yellow]Note: Large OWL/XML files can take several minutes " - "to parse[/yellow]\n" - ) - - # Estimate parse time based on file size (XML is slower: ~10-20 seconds per MB) - estimated_parse_time = file_size_mb * 15 # seconds - - with Progress( - SpinnerColumn(), - TextColumn("[progress.description]{task.description}"), - BarColumn(), - TextColumn("[progress.percentage]{task.percentage:>3.0f}%"), - TimeElapsedColumn(), - console=console, - transient=False, - ) as progress: - parse_task = progress.add_task( - "[yellow]Parsing RDF/XML triples into graph...[/yellow]", - total=100, # Percentage-based - ) - - parse_start = time.time() - - # Update progress in background thread - import threading - - stop_updates = threading.Event() - - def update_parse_progress(): - """Update progress bar based on elapsed time vs estimated time.""" - last_printed_progress = 0 - while not stop_updates.is_set(): - elapsed = time.time() - parse_start - # Calculate progress as percentage of estimated time - parse_progress = min( - 95, (elapsed / estimated_parse_time) * 100 - ) # Cap at 95% until done - progress.update(parse_task, completed=parse_progress) - - # Map to overall progress: parsing is 10-50% of total conversion - # So parse_progress 0-100% maps to 10-50% overall - overall_progress = 10 + (parse_progress / 100) * 40 - - # Emit progress marker for wrapper scripts - if overall_progress >= last_printed_progress + 2: # Update every 2% - print(f"\nPROGRESS: {overall_progress:.0f}", flush=True) - last_printed_progress = overall_progress - - time.sleep(0.5) # Update every 500ms - - # Start progress updater thread - progress_thread = threading.Thread( - target=update_parse_progress, daemon=True - ) - progress_thread.start() - - # Parse directly from file path for XML - this is more reliable - try: - graph.parse(str(file_path), format=format) - except Exception as e: - console.print(f"[red]Error parsing RDF file: {e}[/red]") - console.print(f"[yellow]File: {file_path}[/yellow]") - console.print(f"[yellow]Format: {format}[/yellow]") - console.print( - "[yellow]The file may be corrupted or in a different " - "format[/yellow]" - ) - raise - - # Stop progress updates - stop_updates.set() - progress_thread.join(timeout=1.0) - - parse_time = time.time() - parse_start - - # Set to 100% complete - progress.update( - parse_task, - completed=100, - description=( - f"[green]✓ Parsed {len(graph):,} triples in " - f"{parse_time:.1f}s[/green]" - ), - ) - time.sleep(0.5) - else: - # For non-XML formats, use BytesIO approach with file read progress - with Progress( - SpinnerColumn(), - TextColumn("[progress.description]{task.description}"), - BarColumn(), - TextColumn("[progress.percentage]{task.percentage:>3.0f}%"), - TextColumn("•"), - TextColumn("{task.completed:.1f}/{task.total:.1f} MB"), - TimeElapsedColumn(), - console=console, - transient=False, - ) as progress: - # Create task with file size in MB - read_task = progress.add_task( - "[yellow]Reading RDF file...[/yellow]", - total=file_size_mb, - ) - - # Progress callback updates progress bar - def update_progress(bytes_read): - mb_read = bytes_read / (1024 * 1024) - progress.update(read_task, completed=mb_read) - - # Read file content into memory first - console.print("[dim]Phase 1: Reading file from disk[/dim]") - try: - if file_path.suffix == ".gz": - progress.update( - read_task, - description=( - "[yellow]Reading compressed file (gzip)...[/yellow]" - ), - ) - with ( - gzip.open(file_path, "rb") as f, - ProgressFileReader(f, update_progress, file_size_bytes) as pf, - ): - file_content = pf.read() - elif file_path.suffix == ".bz2" or str(file_path).endswith(".ttl.bz2"): - progress.update( - read_task, - description="[yellow]Reading compressed file (bz2)...[/yellow]", - ) - with ( - bz2.open(file_path, "rb") as f, - ProgressFileReader(f, update_progress, file_size_bytes) as pf, - ): - file_content = pf.read() - else: - progress.update( - read_task, description="[yellow]Reading file...[/yellow]" - ) - with ( - file_path.open("rb") as f, - ProgressFileReader(f, update_progress, file_size_bytes) as pf, - ): - file_content = pf.read() - except (gzip.BadGzipFile, bz2.BadGzipFile) as e: - console.print(f"[red]Invalid or corrupted compressed file: {e}[/red]") - console.print(f"[yellow]File: {file_path}[/yellow]") - console.print("[yellow]Please re-download the dataset[/yellow]") - raise - except OSError as e: - console.print(f"[red]Error reading file: {e}[/red]") - console.print(f"[yellow]File: {file_path}[/yellow]") - console.print("[yellow]Check file permissions and disk space[/yellow]") - raise - - read_time = time.time() - start_time - progress.update( - read_task, - completed=file_size_mb, - description=( - f"[green]✓ Read {file_size_mb:.1f} MB in " - f"{read_time:.1f}s[/green]" - ), - ) - time.sleep(0.3) - - # Now parse the content into triples (this is what takes time) - console.print("\n[dim]Phase 2: Parsing RDF syntax and building graph[/dim]") - - # Estimate parse time based on file size (rough heuristic: ~2-4 seconds per MB) - estimated_parse_time = file_size_mb * 2.5 # seconds - - with Progress( - SpinnerColumn(), - TextColumn("[progress.description]{task.description}"), - BarColumn(), - TextColumn("[progress.percentage]{task.percentage:>3.0f}%"), - TimeElapsedColumn(), - console=console, - transient=False, - ) as progress: - parse_task = progress.add_task( - "[yellow]Parsing RDF triples into graph...[/yellow]", - total=100, # Percentage-based - ) - - parse_start = time.time() - - # Update progress in background thread - import threading - - stop_updates = threading.Event() - - def update_parse_progress(): - """Update progress bar based on elapsed time vs estimated time.""" - last_printed_progress = 0 - while not stop_updates.is_set(): - elapsed = time.time() - parse_start - # Calculate progress as percentage of estimated time - parse_progress = min( - 95, (elapsed / estimated_parse_time) * 100 - ) # Cap at 95% until done - progress.update(parse_task, completed=parse_progress) - - # Map to overall progress: parsing is 10-50% of total conversion - # So parse_progress 0-100% maps to 10-50% overall - overall_progress = 10 + (parse_progress / 100) * 40 - - # Emit progress marker for wrapper scripts - if overall_progress >= last_printed_progress + 2: # Update every 2% - print(f"\nPROGRESS: {overall_progress:.0f}", flush=True) - last_printed_progress = overall_progress - - time.sleep(0.5) # Update every 500ms - - # Start progress updater thread - progress_thread = threading.Thread( - target=update_parse_progress, daemon=True - ) - progress_thread.start() - - # This is the slow part - parsing and building graph structure - try: - graph.parse(BytesIO(file_content), format=format) - except Exception as e: - console.print(f"[red]Error parsing RDF content: {e}[/red]") - console.print(f"[yellow]File: {file_path}[/yellow]") - console.print(f"[yellow]Format: {format}[/yellow]") - console.print( - "[yellow]The file may be corrupted or in a different " - "format[/yellow]" - ) - raise - - # Stop progress updates - stop_updates.set() - progress_thread.join(timeout=1.0) - - parse_time = time.time() - parse_start - - # Set to 100% complete - progress.update( - parse_task, - completed=100, - description=( - f"[green]✓ Parsed {len(graph):,} triples in " - f"{parse_time:.1f}s[/green]" - ), - ) - time.sleep(0.5) - - # Calculate total elapsed time - elapsed = time.time() - start_time - - console.print( - f"\n[bold green]✓ Successfully loaded {len(graph):,} triples[/bold green]" - ) - console.print( - f"[dim]Parse speed: {len(graph) / elapsed:.0f} triples/second[/dim]\n" - ) - - # Prepare graph for iteration (this can take time for large graphs) - console.print("[cyan]Preparing graph for conversion...[/cyan]") - with Progress( - SpinnerColumn(), - TextColumn("[progress.description]{task.description}"), - TimeElapsedColumn(), - console=console, - transient=True, # This progress bar disappears after completion - ) as prep_progress: - prep_task = prep_progress.add_task( - "[yellow]Indexing triples for conversion...[/yellow]", total=None - ) - - # Create list from graph iterator (this is what takes time) - graph_list = list(graph) - - prep_progress.update( - prep_task, - description=( - f"[green]✓ Ready to convert {len(graph_list):,} triples[/green]" - ), - ) - time.sleep(0.3) # Brief pause so user sees completion - - # Convert to list of dictionaries with detailed progress - triples = [] - - # Emit progress at 50% (start of conversion) - print("\nPROGRESS: 50", flush=True) - sys.stdout.flush() # Force flush - - with Progress( - TextColumn("[progress.description]{task.description}"), - BarColumn(), - TextColumn("[progress.percentage]{task.percentage:>3.0f}%"), - TextColumn("•"), - TextColumn("{task.completed:,}/{task.total:,} triples"), - TimeRemainingColumn(), - console=console, - ) as progress: - task = progress.add_task( - "[cyan]Converting triples to dataset format...", total=len(graph_list) - ) - - # Calculate update frequency - total_count = len(graph_list) - # Update at least 20 times during conversion, - # but not more than every 100 triples - update_interval = max(1, min(100, total_count // 20)) - last_progress_emit = 50 - - for idx, (subj, pred, obj) in enumerate(graph_list): - # Convert URIRefs and Literals to strings - subject = str(subj) - predicate = str(pred) - object_value = str(obj) - - # Extract object type and datatype if available - if isinstance(obj, Literal): - object_type = "literal" - object_datatype = str(obj.datatype) if obj.datatype else None - object_language = obj.language if obj.language else None - elif isinstance(obj, URIRef): - object_type = "uri" - object_datatype = None - object_language = None - else: - object_type = "blank_node" - object_datatype = None - object_language = None - - triple = { - "subject": subject, - "predicate": predicate, - "object": object_value, - "object_type": object_type, - "object_datatype": object_datatype, - "object_language": object_language, - } - triples.append(triple) - - # Update progress bar and emit progress markers - if (idx + 1) % update_interval == 0 or idx == total_count - 1: - # Update visual progress bar - progress.update(task, completed=idx + 1) - - # Calculate and emit overall progress (50-90% range) - conversion_pct = ((idx + 1) / total_count) * 100 - overall_progress = 50 + (conversion_pct / 100) * 40 - - # Only emit if we've made significant progress - if overall_progress >= last_progress_emit + 2 and idx < total_count - 1: - print(f"\nPROGRESS: {overall_progress:.0f}", flush=True) - sys.stdout.flush() # Force flush for subprocess - last_progress_emit = overall_progress - - # Emit final conversion progress - print("\nPROGRESS: 90", flush=True) - sys.stdout.flush() - - return triples - - -def create_hf_dataset(triples: list[dict[str, str]]) -> DatasetDict: - """Create HuggingFace dataset from triples. - - Args: - triples: List of triple dictionaries - - Returns: - DatasetDict with single 'data' split containing all triples - """ - console = Console() - triple_count = len(triples) - console.print( - f"[cyan]Creating HuggingFace dataset from {triple_count} triples[/cyan]" - ) - - # Define schema - features = Features( - { - "subject": Value("string"), - "predicate": Value("string"), - "object": Value("string"), - "object_type": Value("string"), - "object_datatype": Value("string"), - "object_language": Value("string"), - } - ) - - # Build dataset with visible progress feedback (heavy operation for large graphs) - dataset_container: dict[str, Dataset] = {} - error_container: list[BaseException] = [] - - @contextmanager - def dataset_build_progress(stage: str): - with Progress( - SpinnerColumn(), - TextColumn("[progress.description]{task.description}"), - TimeElapsedColumn(), - console=console, - transient=False, - refresh_per_second=6, - ) as progress: - task = progress.add_task( - f"[yellow]{stage}[/yellow]", - total=None, - ) - start_time = time.time() - try: - yield progress, task - except BaseException: - progress.update(task, description="[red]Dataset build failed[/red]") - raise - else: - elapsed = time.time() - start_time - progress.update( - task, - description=f"[green]✓ Dataset built in {elapsed:.1f}s[/green]", - ) - time.sleep(0.3) - - def build_dataset() -> None: - try: - dataset_container["data"] = Dataset.from_list(triples, features=features) - except MemoryError as exc: - console.print("[red]Out of memory creating dataset[/red]") - console.print(f"[yellow]Dataset has {len(triples):,} triples[/yellow]") - console.print( - "[yellow]Try using the streaming converter for large datasets[/yellow]" - ) - error_container.append(exc) - except BaseException as exc: # Capture any issue (including KeyboardInterrupt) - error_container.append(exc) - - with dataset_build_progress( - "Building HuggingFace dataset (this can take several minutes)..." - ) as ( - progress, - task, - ): - import threading - - thread = threading.Thread(target=build_dataset, daemon=True) - thread.start() - - while thread.is_alive(): - progress.update( - task, - description=( - "[yellow]Building dataset...[/yellow] " - f"[dim]{triple_count:,} triples[/dim]" - ), - ) - time.sleep(0.5) - - thread.join() - - if error_container: - raise error_container[0] - - dataset = dataset_container["data"] - - # Return as DatasetDict with single 'data' split - dataset_dict = DatasetDict({"data": dataset}) - - console.print(f"[green]Created dataset with {len(dataset)} triples[/green]") - - return dataset_dict - - -def add_dataset_metadata( - dataset_dict: DatasetDict, metadata: dict[str, Any] -) -> DatasetDict: - """Add metadata to the dataset. - - Args: - dataset_dict: HuggingFace dataset - metadata: Metadata dictionary - - Returns: - Dataset with metadata - """ - # Add metadata to dataset info - for split in dataset_dict: - dataset_dict[split].info.description = metadata.get("description", "") - dataset_dict[split].info.citation = metadata.get("citation", "") - dataset_dict[split].info.homepage = metadata.get("homepage", "") - dataset_dict[split].info.license = metadata.get("license", "") - - return dataset_dict - - -def convert_rdf_to_hf( - input_path: Path, - output_path: Path, - rdf_format: str = "turtle", - metadata: dict[str, Any] | None = None, -) -> None: - """Convert an RDF file to HuggingFace dataset format. - - Args: - input_path: Path to input RDF file - output_path: Path to output directory - rdf_format: RDF format (turtle, nt, xml, tsv, etc.) - metadata: Optional metadata dictionary - """ - console = Console() - - console.print("\n[bold cyan]Converting RDF to HuggingFace Dataset[/bold cyan]") - console.print(f"Input: {input_path}") - console.print(f"Output: {output_path}\n") - - # Parse file based on format - if rdf_format == "tsv": - # Use TSV parser for benchmark datasets - triples = parse_tsv_triples(input_path) - else: - # Use RDF parser for standard RDF formats - triples = parse_rdf_file(input_path, format=rdf_format) - - if not triples: - console.print("[red]No triples found in file[/red]") - console.print(f"[yellow]File: {input_path}[/yellow]") - console.print(f"[yellow]Format: {rdf_format}[/yellow]") - console.print("[yellow]Check file format and content[/yellow]") - return - - # Create HuggingFace dataset - dataset_dict = create_hf_dataset(triples) - - # Add metadata if provided - if metadata: - dataset_dict = add_dataset_metadata(dataset_dict, metadata) - - # Save dataset - try: - output_path.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_path}[/yellow]") - console.print("[yellow]Check permissions and disk space[/yellow]") - raise - - try: - dataset_dict.save_to_disk(str(output_path)) - except Exception as e: - console.print(f"[red]Error saving dataset: {e}[/red]") - console.print(f"[yellow]Output directory: {output_path}[/yellow]") - console.print("[yellow]Check disk space and permissions[/yellow]") - raise - - console.print(f"\n[bold green]✓ Dataset saved to {output_path}[/bold green]") - - # Print statistics - console.print("\n[bold]Dataset Statistics:[/bold]") - console.print(f" • Total triples: {len(triples)}") - - -def main(): - """Main entry point.""" - parser = argparse.ArgumentParser( - description="Convert RDF files to HuggingFace dataset format" - ) - parser.add_argument("input", type=Path, help="Input RDF file") - parser.add_argument( - "output", type=Path, help="Output directory for HuggingFace dataset" - ) - parser.add_argument( - "--format", - "-f", - default="turtle", - choices=["turtle", "nt", "ntriples", "xml", "n3", "trig", "nquads", "tsv"], - help="RDF format (default: turtle)", - ) - parser.add_argument("--description", type=str, help="Dataset description") - parser.add_argument("--citation", type=str, help="Dataset citation") - parser.add_argument("--homepage", type=str, help="Dataset homepage URL") - parser.add_argument("--license", type=str, help="Dataset license") - parser.add_argument( - "--verbose", - "-v", - action="store_true", - help=( - "Show verbose output including warnings from libraries " - "(default: only errors)" - ), - ) - - args = parser.parse_args() - - # Configure logging based on verbosity - # Default: only show ERROR level (suppress all warnings and info) - # Verbose: show WARNING level and above - log_level = logging.WARNING if args.verbose else logging.ERROR - - # Configure root logger to suppress all third-party library output - logging.basicConfig( - level=log_level, - format="%(levelname)s:%(name)s:%(message)s", - force=True, # Override any existing configuration - ) - - # Also set specific loggers that might be noisy - logging.getLogger("rdflib").setLevel(log_level) - logging.getLogger("datasets").setLevel(log_level) - logging.getLogger("urllib3").setLevel(log_level) - logging.getLogger("filelock").setLevel(log_level) - - # Build metadata - metadata = {} - if args.description: - metadata["description"] = args.description - if args.citation: - metadata["citation"] = args.citation - if args.homepage: - metadata["homepage"] = args.homepage - if args.license: - metadata["license"] = args.license - - # Convert - convert_rdf_to_hf( - input_path=args.input, - output_path=args.output, - rdf_format=args.format, - metadata=metadata if metadata else None, - ) - - -if __name__ == "__main__": - main() diff --git a/scripts/convert_rdf_to_hf_dataset_streaming.py b/scripts/convert_rdf_to_hf_dataset_streaming.py deleted file mode 100755 index d15485c..0000000 --- a/scripts/convert_rdf_to_hf_dataset_streaming.py +++ /dev/null @@ -1,720 +0,0 @@ -#!/usr/bin/env python3 -"""Convert RDF datasets to HuggingFace dataset format using streaming. - -This minimizes memory usage. - -This script parses RDF files in various formats (Turtle, N-Triples, RDF/XML, etc.) -and converts them to HuggingFace dataset format with lossless preservation of all -semantic information. Uses streaming and chunking to handle large datasets with -minimal memory footprint. - -USAGE: - # Basic conversion - python scripts/convert_rdf_to_hf_dataset_streaming.py \\ - [options] - - # Examples - Different RDF formats - python scripts/convert_rdf_to_hf_dataset_streaming.py \ - raw_datasets/geonames/geonames.rdf \ - hf_datasets/geonames \ - --format xml \ - --chunk-size 10000 - -ARGUMENTS: - input Path to input RDF file (any supported format) - output Path to output directory for HuggingFace dataset - -f, --format FORMAT RDF format: turtle, nt, ntriples, xml, n3, trig, nquads - (default: turtle) - --chunk-size SIZE Number of triples to process at once (default: 10000) - --max-workers NUM Maximum number of worker processes (default: CPU count) - --description TEXT Dataset description (added to metadata) - --citation TEXT Dataset citation (added to metadata) - --homepage URL Dataset homepage URL (added to metadata) - --license TEXT Dataset license (added to metadata) - -v, --verbose Show verbose output including library warnings - -MEMORY OPTIMIZATIONS: - - Streams RDF file instead of loading entirely into memory - - Processes triples in configurable chunks - - Writes to Parquet files incrementally - - Uses generators to avoid materializing large lists - - Parallel processing for GeoNames format without memory accumulation - -OUTPUT: - Creates HuggingFace dataset in Arrow/Parquet format with minimal memory usage. -""" - -from __future__ import annotations - -import argparse -import gzip -import json -import logging -import multiprocessing as mp -import shutil -import time -from collections.abc import Iterator -from pathlib import Path -from typing import Any, cast - -import pyarrow as pa -import pyarrow.parquet as pq -from datasets import Dataset, DatasetDict -from rdflib import Graph, Literal, URIRef -from rich.console import Console -from rich.progress import ( - BarColumn, - Progress, - SpinnerColumn, - TextColumn, - TimeElapsedColumn, -) - -# Logger will be configured in main() based on --verbose flag -logger = logging.getLogger(__name__) - - -def _process_ntriples_file( - file_obj, chunk_size: int -) -> Iterator[list[dict[str, str]]]: - """Process N-Triples file and yield chunks of triples. - - This helper function processes the file object line by line, parsing each - line as an N-Triple and accumulating triples into chunks. - - Args: - file_obj: File-like object (already opened) - chunk_size: Number of triples per chunk - - Yields: - Chunks of triple dictionaries - """ - from rdflib import Graph - - current_chunk = [] - - for line_no, line in enumerate(file_obj, 1): - line = line.strip() - # Skip empty lines and comments - if not line or line.startswith("#"): - continue - - try: - # Parse single line as N-Triple - mini_graph = Graph() - mini_graph.parse(data=line, format="nt") - - for s, p, o in mini_graph: - # Convert RDFLib objects to dictionary format - if isinstance(o, Literal): - object_type = "literal" - object_datatype = ( - str(o.datatype) if o.datatype else None - ) - object_language = o.language if o.language else None - elif isinstance(o, URIRef): - object_type = "uri" - object_datatype = None - object_language = None - else: - object_type = "blank_node" - object_datatype = None - object_language = None - - triple = { - "subject": str(s), - "predicate": str(p), - "object": str(o), - "object_type": object_type, - "object_datatype": object_datatype, - "object_language": object_language, - } - current_chunk.append(triple) - - # Yield chunk when it reaches the target size - if len(current_chunk) >= chunk_size: - yield current_chunk - current_chunk = [] - - except Exception as e: - # Skip malformed lines and continue processing - logger.debug(f"Skipping malformed line {line_no}: {e}") - continue - - # Yield any remaining triples that didn't fill a complete chunk - if current_chunk: - yield current_chunk - - -def stream_ntriples( - file_path: Path, chunk_size: int = 10000 -) -> Iterator[list[dict[str, str]]]: - """Stream N-Triples file in chunks without loading entire file into memory. - - Args: - file_path: Path to N-Triples file - chunk_size: Number of triples per chunk - - Yields: - Chunks of triple dictionaries - """ - # Open file with appropriate context manager based on compression - # Context managers ensure files are properly closed even if exceptions occur - if file_path.suffix == ".gz": - with gzip.open(file_path, "rt", encoding="utf-8") as file_obj: - yield from _process_ntriples_file(file_obj, chunk_size) - else: - with open(file_path, encoding="utf-8") as file_obj: - yield from _process_ntriples_file(file_obj, chunk_size) - - -def stream_turtle_chunks( - file_path: Path, chunk_size: int = 10000 -) -> Iterator[list[dict[str, str]]]: - """Stream Turtle file in chunks using incremental parsing. - - Args: - file_path: Path to Turtle file - chunk_size: Number of triples per chunk - - Yields: - Chunks of triple dictionaries - """ - from rdflib import Graph - - # For Turtle, we need to parse the entire file due to prefix definitions - # But we can yield results in chunks to avoid keeping all in memory - graph = Graph() - - # Parse file - if file_path.suffix == ".gz": - with gzip.open(file_path, "rt", encoding="utf-8") as f: - graph.parse(f, format="turtle") - else: - graph.parse(str(file_path), format="turtle") - - # Yield triples in chunks - current_chunk = [] - for s, p, o in graph: - # Convert to dictionary format - if isinstance(o, Literal): - object_type = "literal" - object_datatype = str(o.datatype) if o.datatype else None - object_language = o.language if o.language else None - elif isinstance(o, URIRef): - object_type = "uri" - object_datatype = None - object_language = None - else: - object_type = "blank_node" - object_datatype = None - object_language = None - - triple = { - "subject": str(s), - "predicate": str(p), - "object": str(o), - "object_type": object_type, - "object_datatype": object_datatype, - "object_language": object_language, - } - current_chunk.append(triple) - - if len(current_chunk) >= chunk_size: - yield current_chunk - current_chunk = [] - - # Yield remaining triples - if current_chunk: - yield current_chunk - - -def process_geonames_segment(args): - """Process a file segment for GeoNames parallel processing. - - This function is defined at module level so it can be pickled for multiprocessing. - - Args: - args: Tuple of (file_path, start_byte, end_byte, chunk_size) - - Returns: - List of chunks of triple dictionaries - """ - from rdflib import Graph, Literal, URIRef - - file_path, start_byte, end_byte, chunk_size = args - chunks = [] - current_chunk = [] - current_xml = [] - - with open(file_path, encoding="utf-8") as f: - f.seek(start_byte) - - # Skip to next document boundary if not at start - if start_byte > 0: - while f.tell() < end_byte: - line = f.readline() - if not line: - break - if line.startswith("http://") or line.startswith("https://"): - break - - # Process documents in segment - while f.tell() < end_byte: - line = f.readline() - if not line: - break - - # Document boundary - if line.startswith("http://") or line.startswith("https://"): - if current_xml: - # Parse accumulated document - xml_str = "".join(current_xml) - try: - graph = Graph() - graph.parse(data=xml_str, format="xml") - - for s, p, o in graph: - # Convert to dictionary format - if isinstance(o, Literal): - object_type = "literal" - object_datatype = ( - str(o.datatype) if o.datatype else None - ) - object_language = o.language if o.language else None - elif isinstance(o, URIRef): - object_type = "uri" - object_datatype = None - object_language = None - else: - object_type = "blank_node" - object_datatype = None - object_language = None - - triple = { - "subject": str(s), - "predicate": str(p), - "object": str(o), - "object_type": object_type, - "object_datatype": object_datatype, - "object_language": object_language, - } - current_chunk.append(triple) - - if len(current_chunk) >= chunk_size: - chunks.append(current_chunk) - current_chunk = [] - except Exception: - pass - - current_xml = [] - continue - - current_xml.append(line) - - # Handle remaining data - if current_xml: - xml_str = "".join(current_xml) - try: - graph = Graph() - graph.parse(data=xml_str, format="xml") - - for s, p, o in graph: - # Convert to dictionary format - if isinstance(o, Literal): - object_type = "literal" - object_datatype = str(o.datatype) if o.datatype else None - object_language = o.language if o.language else None - elif isinstance(o, URIRef): - object_type = "uri" - object_datatype = None - object_language = None - else: - object_type = "blank_node" - object_datatype = None - object_language = None - - triple = { - "subject": str(s), - "predicate": str(p), - "object": str(o), - "object_type": object_type, - "object_datatype": object_datatype, - "object_language": object_language, - } - current_chunk.append(triple) - - if len(current_chunk) >= chunk_size: - chunks.append(current_chunk) - current_chunk = [] - except Exception: - pass - - if current_chunk: - chunks.append(current_chunk) - - return chunks - - -def stream_geonames_chunks( - file_path: Path, chunk_size: int = 10000, num_workers: int | None = None -) -> Iterator[list[dict[str, str]]]: - """Stream GeoNames RDF file using parallel chunk processing. - - Each worker processes a file segment independently and yields chunks. - - Args: - file_path: Path to GeoNames RDF file - chunk_size: Number of triples per chunk - num_workers: Number of worker processes - - Yields: - Chunks of triple dictionaries - """ - if num_workers is None: - num_workers = mp.cpu_count() - - from multiprocessing import Pool - - # Calculate file segments - file_size = file_path.stat().st_size - segment_size = file_size // num_workers - segments = [] - - with open(file_path, "rb") as f: - for i in range(num_workers): - start = i * segment_size - if i == num_workers - 1: - end = file_size - else: - end = start + segment_size - f.seek(end) - # Find next document boundary - while end < file_size: - line = f.readline() - end = f.tell() - decoded = line.decode("utf-8", errors="ignore") - if decoded.startswith("http://") or decoded.startswith("https://"): - break - segments.append((str(file_path), start, end, chunk_size)) - - # Process segments in parallel using the module-level function - with Pool(processes=num_workers) as pool: - results = pool.map(process_geonames_segment, segments) - - # Yield chunks from all workers - for worker_chunks in results: - yield from worker_chunks - - -def stream_rdf_chunks( - file_path: Path, - format: str = "turtle", - chunk_size: int = 10000, - num_workers: int | None = None, -) -> Iterator[list[dict[str, str]]]: - """Stream RDF file in chunks based on format. - - Args: - file_path: Path to RDF file - format: RDF format - chunk_size: Number of triples per chunk - num_workers: Number of workers for parallel formats - - Yields: - Chunks of triple dictionaries - """ - console = Console() - - # Check if this is GeoNames format - is_geonames = "geonames" in str(file_path).lower() - if is_geonames and format in ("xml", "application/rdf+xml"): - with open(file_path, encoding="utf-8") as f: - first_line = f.readline().strip() - if first_line.startswith("http://") or first_line.startswith("https://"): - console.print("[yellow]Using GeoNames streaming parser[/yellow]") - yield from stream_geonames_chunks(file_path, chunk_size, num_workers) - return - - # Use appropriate streaming method based on format - if format in ("nt", "ntriples"): - console.print("[yellow]Using N-Triples streaming parser[/yellow]") - yield from stream_ntriples(file_path, chunk_size) - elif format == "turtle": - console.print("[yellow]Using Turtle chunk parser[/yellow]") - yield from stream_turtle_chunks(file_path, chunk_size) - else: - # For other formats, fall back to regular parsing with chunked output - console.print(f"[yellow]Using generic RDF parser for {format}[/yellow]") - graph = Graph() - graph.parse(str(file_path), format=format) - - current_chunk = [] - for s, p, o in graph: - # Convert to dictionary format - if isinstance(o, Literal): - object_type = "literal" - object_datatype = str(o.datatype) if o.datatype else None - object_language = o.language if o.language else None - elif isinstance(o, URIRef): - object_type = "uri" - object_datatype = None - object_language = None - else: - object_type = "blank_node" - object_datatype = None - object_language = None - - triple = { - "subject": str(s), - "predicate": str(p), - "object": str(o), - "object_type": object_type, - "object_datatype": object_datatype, - "object_language": object_language, - } - current_chunk.append(triple) - - if len(current_chunk) >= chunk_size: - yield current_chunk - current_chunk = [] - - if current_chunk: - yield current_chunk - - -def convert_rdf_to_hf_streaming( - input_path: Path, - output_path: Path, - rdf_format: str = "turtle", - chunk_size: int = 10000, - num_workers: int | None = None, - metadata: dict[str, Any] | None = None, -) -> None: - """Convert an RDF file to HuggingFace dataset format using streaming. - - Args: - input_path: Path to input RDF file - output_path: Path to output directory - rdf_format: RDF format (turtle, nt, xml, etc.) - chunk_size: Number of triples to process at once - num_workers: Number of worker processes for parallel formats - metadata: Optional metadata dictionary - """ - console = Console() - - # Create output directory - output_path.mkdir(parents=True, exist_ok=True) - - # Get file size for progress tracking - file_size_mb = input_path.stat().st_size / (1024 * 1024) - console.print(f"[cyan]Converting RDF file: {input_path}[/cyan]") - console.print(f"[dim]File size: {file_size_mb:.2f} MB ({rdf_format} format)[/dim]") - console.print(f"[dim]Chunk size: {chunk_size:,} triples[/dim]") - console.print(f"[dim]Output: {output_path}[/dim]\n") - - # Define schema - schema = pa.schema( - [ - ("subject", pa.string()), - ("predicate", pa.string()), - ("object", pa.string()), - ("object_type", pa.string()), - ("object_datatype", pa.string()), - ("object_language", pa.string()), - ] - ) - - # Process and write chunks - data_dir = output_path / "data" - data_dir.mkdir(exist_ok=True) - - total_triples = 0 - chunk_count = 0 - parquet_files = [] - - with Progress( - SpinnerColumn(), - TextColumn("[progress.description]{task.description}"), - BarColumn(), - TextColumn("{task.completed:,} chunks"), - TextColumn("•"), - TextColumn("{task.fields[triples]:,} triples"), - TimeElapsedColumn(), - console=console, - transient=False, - ) as progress: - task = progress.add_task( - "[yellow]Processing RDF chunks...[/yellow]", - total=None, - triples=0, - ) - - # Stream and process chunks - for chunk in stream_rdf_chunks(input_path, rdf_format, chunk_size, num_workers): - # Convert chunk to Arrow table - table = pa.Table.from_pylist(chunk, schema=schema) - - # Write chunk to Parquet file - chunk_file = data_dir / f"chunk-{chunk_count:05d}.parquet" - pq.write_table(table, chunk_file) - parquet_files.append(chunk_file) - - # Update progress - total_triples += len(chunk) - chunk_count += 1 - progress.update( - task, - advance=1, - completed=chunk_count, - triples=total_triples, - ) - - console.print( - f"\n[green]✓ Processed {total_triples:,} triples in " - f"{chunk_count} chunks[/green]" - ) - - # Merge Parquet files into final dataset structure - console.print("\n[yellow]Merging chunks into final dataset...[/yellow]") - - # Read all Parquet files as a single dataset - # Note: from_parquet can return different types, but with a glob pattern - # it returns a Dataset that supports train_test_split - dataset = cast(Dataset, Dataset.from_parquet(str(data_dir / "*.parquet"))) - - # Split into train/test - console.print("[dim]Creating train/test split (95%/5%)...[/dim]") - train_test = dataset.train_test_split(test_size=0.05, seed=42) - - # Create DatasetDict - dataset_dict = DatasetDict( - { - "train": train_test["train"], - "test": train_test["test"], - } - ) - - # Add metadata - if metadata: - for split in dataset_dict: - dataset_dict[split].info.description = metadata.get("description", "") - dataset_dict[split].info.citation = metadata.get("citation", "") - dataset_dict[split].info.homepage = metadata.get("homepage", "") - dataset_dict[split].info.license = metadata.get("license", "") - - # Save final dataset - console.print("[yellow]Saving final dataset...[/yellow]") - dataset_dict.save_to_disk(str(output_path)) - - # Clean up temporary chunk files - console.print("[dim]Cleaning up temporary files...[/dim]") - shutil.rmtree(data_dir) - - # Save dataset info - info = { - "format": "parquet", - "total_triples": total_triples, - "train_size": len(dataset_dict["train"]), - "test_size": len(dataset_dict["test"]), - "source_format": rdf_format, - "chunk_size": chunk_size, - "conversion_date": time.strftime("%Y-%m-%d %H:%M:%S"), - } - - with open(output_path / "dataset_info.json", "w") as f: - json.dump(info, f, indent=2) - - console.print( - "\n[bold green]✓ Successfully converted to HuggingFace dataset[/bold green]" - ) - console.print(f"[green] Train: {len(dataset_dict['train']):,} triples[/green]") - console.print(f"[green] Test: {len(dataset_dict['test']):,} triples[/green]") - console.print(f"[green] Location: {output_path}[/green]") - - -def main(): - """Main entry point.""" - parser = argparse.ArgumentParser( - description="Convert RDF to HuggingFace dataset format with streaming", - formatter_class=argparse.RawDescriptionHelpFormatter, - ) - - parser.add_argument("input", type=Path, help="Input RDF file") - parser.add_argument( - "output", type=Path, help="Output directory for HuggingFace dataset" - ) - parser.add_argument( - "-f", - "--format", - default="turtle", - choices=["turtle", "nt", "ntriples", "xml", "n3", "trig", "nquads"], - help="RDF format (default: turtle)", - ) - parser.add_argument( - "--chunk-size", - type=int, - default=10000, - help="Number of triples to process at once (default: 10000)", - ) - parser.add_argument( - "--max-workers", - type=int, - default=None, - help="Maximum number of worker processes (default: CPU count)", - ) - parser.add_argument("--description", help="Dataset description") - parser.add_argument("--citation", help="Dataset citation") - parser.add_argument("--homepage", help="Dataset homepage URL") - parser.add_argument("--license", help="Dataset license") - parser.add_argument( - "-v", "--verbose", action="store_true", help="Show verbose output" - ) - - args = parser.parse_args() - - # Configure logging - if args.verbose: - logging.basicConfig(level=logging.INFO) - else: - logging.basicConfig(level=logging.ERROR) - - # Validate input - if not args.input.exists(): - print(f"Error: Input file not found: {args.input}") - return 1 - - # Prepare metadata - metadata = {} - if args.description: - metadata["description"] = args.description - if args.citation: - metadata["citation"] = args.citation - if args.homepage: - metadata["homepage"] = args.homepage - if args.license: - metadata["license"] = args.license - - # Convert - try: - convert_rdf_to_hf_streaming( - input_path=args.input, - output_path=args.output, - rdf_format=args.format, - chunk_size=args.chunk_size, - num_workers=args.max_workers, - metadata=metadata, - ) - return 0 - except Exception as e: - console = Console() - console.print(f"[red]Error: {e}[/red]") - if args.verbose: - import traceback - - traceback.print_exc() - return 1 - - -if __name__ == "__main__": - exit(main()) diff --git a/scripts/convert_rdf_to_hf_dataset_streaming_parallel.py b/scripts/convert_rdf_to_hf_dataset_streaming_parallel.py deleted file mode 100755 index 4a488f5..0000000 --- a/scripts/convert_rdf_to_hf_dataset_streaming_parallel.py +++ /dev/null @@ -1,862 +0,0 @@ -#!/usr/bin/env python3 -"""Parallel streaming converter for RDF to HuggingFace dataset format with minimal memory usage. - -This version uses multiprocessing.Pool.imap for simpler and more reliable parallelization. -""" - -from __future__ import annotations - -import argparse -import gzip -import json -import logging -import multiprocessing as mp -import shutil -import sys -import time -from collections.abc import Iterator -from pathlib import Path -from typing import Any - -import pyarrow as pa -import pyarrow.parquet as pq -from datasets import Dataset, DatasetDict, Features, Value -from rdflib import Graph, Literal, URIRef -from rich.console import Console -from rich.progress import ( - BarColumn, - Progress, - SpinnerColumn, - TextColumn, - TimeElapsedColumn, -) - -logger = logging.getLogger(__name__) -logger.addHandler(logging.StreamHandler(sys.stdout)) - - -def process_geonames_lines(lines): - """Process a batch of lines from GeoNames file. - - Module-level function for multiprocessing. - - Args: - lines: List of lines containing one or more XML documents - - Returns: - List of triple dictionaries - """ - from rdflib import Graph, Literal, URIRef - - triples = [] - current_xml = [] - - for line in lines: - # Check if this is a document boundary (URL line) - if line.startswith("http://") or line.startswith("https://"): - # Process previous document if exists - if current_xml: - xml_str = "".join(current_xml) - try: - # Parse the RDF/XML document - graph = Graph() - graph.parse(data=xml_str, format="xml") - - # Extract triples - for s, p, o in graph: - # Convert to dictionary format - if isinstance(o, Literal): - object_type = "literal" - object_datatype = str(o.datatype) if o.datatype else None - object_language = o.language if o.language else None - elif isinstance(o, URIRef): - object_type = "uri" - object_datatype = None - object_language = None - else: - object_type = "blank_node" - object_datatype = None - object_language = None - - triple = { - "subject": str(s), - "predicate": str(p), - "object": str(o), - "object_type": object_type, - "object_datatype": object_datatype, - "object_language": object_language, - } - triples.append(triple) - - except Exception as e: - # Skip malformed documents - logger.error(f"Error parsing XML document: {e}. Skipping document.") - pass - - # Reset for next document - current_xml = [] - - # Skip the URL line itself - continue - - # Accumulate XML content - current_xml.append(line) - - # Process final document if exists - if current_xml: - xml_str = "".join(current_xml) - try: - graph = Graph() - graph.parse(data=xml_str, format="xml") - - for s, p, o in graph: - if isinstance(o, Literal): - object_type = "literal" - object_datatype = str(o.datatype) if o.datatype else None - object_language = o.language if o.language else None - elif isinstance(o, URIRef): - object_type = "uri" - object_datatype = None - object_language = None - else: - object_type = "blank_node" - object_datatype = None - object_language = None - - triple = { - "subject": str(s), - "predicate": str(p), - "object": str(o), - "object_type": object_type, - "object_datatype": object_datatype, - "object_language": object_language, - } - triples.append(triple) - except Exception as e: - logger.error(f"Error parsing final XML document: {e}. Skipping document.") - pass - - return triples - - -def process_ntriples_lines(lines): - """Process a batch of N-Triples lines. - - Args: - lines: List of N-Triple lines - - Returns: - List of triple dictionaries - """ - from rdflib import Graph, Literal, URIRef - - triples = [] - - for line in lines: - line = line.strip() - if not line or line.startswith("#"): - continue - - try: - # Parse single N-Triple line - mini_graph = Graph() - mini_graph.parse(data=line, format="nt") - - for s, p, o in mini_graph: - # Convert to dictionary format - if isinstance(o, Literal): - object_type = "literal" - object_datatype = str(o.datatype) if o.datatype else None - object_language = o.language if o.language else None - elif isinstance(o, URIRef): - object_type = "uri" - object_datatype = None - object_language = None - else: - object_type = "blank_node" - object_datatype = None - object_language = None - - triple = { - "subject": str(s), - "predicate": str(p), - "object": str(o), - "object_type": object_type, - "object_datatype": object_datatype, - "object_language": object_language, - } - triples.append(triple) - - except Exception as e: - # Skip malformed lines - logger.error(f"Error parsing N-Triple line: {line}. Skipping line. Error: {e}") - continue - - return triples - - -def batch_file_lines(file_path: Path, batch_size: int = 1000, format: str = "geonames"): - """Generator that yields batches of lines from file. - - For GeoNames format, ensures complete XML documents are kept together. - - Args: - file_path: Path to file - batch_size: Number of lines per batch - format: File format (geonames or ntriples) - - Yields: - Batches of lines - """ - # Handle compressed files - if file_path.suffix == ".gz": - try: - file_obj = gzip.open(file_path, "rt", encoding="utf-8") - # Test read to check if file is valid - test_line = file_obj.readline() - if not test_line and file_path.stat().st_size > 0: - raise EOFError("Compressed file appears to be empty or corrupted") - file_obj.seek(0) - except (gzip.BadGzipFile, EOFError, OSError) as e: - from rich.console import Console - - console = Console() - console.print(f"[red]Error: Compressed file is corrupted or incomplete: {e}[/red]") - console.print(f"[yellow]File: {file_path}[/yellow]") - console.print("[yellow]Please re-download the dataset or use an uncompressed version[/yellow]") - raise - else: - file_obj = open(file_path, encoding="utf-8") - - try: - current_batch = [] - current_doc = [] - - if format == "geonames": - # For GeoNames, keep documents together - for line in file_obj: - # Check if this is a document boundary - if line.startswith("http://") or line.startswith("https://"): - if current_doc: - # Add completed document to batch - current_batch.extend(current_doc) - current_doc = [line] - - # Yield batch if large enough - if len(current_batch) >= batch_size: - yield current_batch - current_batch = [] - else: - current_doc = [line] - else: - current_doc.append(line) - - # Add final document - if current_doc: - current_batch.extend(current_doc) - - # Yield final batch - if current_batch: - yield current_batch - - else: - # For N-Triples, just batch lines - for line in file_obj: - current_batch.append(line) - - if len(current_batch) >= batch_size: - yield current_batch - current_batch = [] - - # Yield remaining lines - if current_batch: - yield current_batch - - finally: - file_obj.close() - - -def stream_geonames_parallel( - file_path: Path, chunk_size: int = 10000, num_workers: int | None = None -) -> Iterator[list[dict[str, str]]]: - """Stream GeoNames RDF file with parallel processing using imap. - - Args: - file_path: Path to GeoNames RDF file - chunk_size: Number of lines per batch - num_workers: Number of worker processes - - Yields: - Chunks of triple dictionaries - """ - if num_workers is None: - num_workers = max(1, mp.cpu_count() - 1) - - console = Console() - console.print(f"[yellow]Using parallel GeoNames parser with {num_workers} workers[/yellow]") - - # Create batches of lines - batches = batch_file_lines(file_path, batch_size=chunk_size, format="geonames") - - # Process batches in parallel - try: - with mp.Pool(processes=num_workers) as pool: - # Use imap_unordered for better performance (order doesn't matter) - for triples in pool.imap_unordered(process_geonames_lines, batches, chunksize=1): - if triples: - yield triples - except Exception as e: - logger.error(f"Error processing GeoNames file: {e}") - raise e - - -def stream_ntriples_parallel( - file_path: Path, chunk_size: int = 10000, num_workers: int | None= None -) -> Iterator[list[dict[str, str]]]: - """Stream N-Triples file with parallel processing using imap. - - Args: - file_path: Path to N-Triples file - chunk_size: Number of lines per batch - num_workers: Number of worker processes - - Yields: - Chunks of triple dictionaries - """ - if num_workers is None: - num_workers = max(1, mp.cpu_count() - 1) - - console = Console() - console.print(f"[yellow]Using parallel N-Triples parser with {num_workers} workers[/yellow]") - - # Create batches of lines - batches = batch_file_lines(file_path, batch_size=chunk_size, format="ntriples") - - # Process batches in parallel - try: - with mp.Pool(processes=num_workers) as pool: - for triples in pool.imap_unordered(process_ntriples_lines, batches, chunksize=1): - if triples: - yield triples - except Exception as e: - logger.error(f"Error processing N-Triples file: {e}") - raise e - - -def stream_turtle_chunks(file_path: Path, chunk_size: int = 10000) -> Iterator[list[dict[str, str]]]: - """Stream Turtle file in chunks using incremental parsing. - - This is a simpler approach that reads the file line by line and yields chunks - without trying to parse the entire file at once. - - Args: - file_path: Path to Turtle file - chunk_size: Number of triples per chunk - - Yields: - Chunks of triple dictionaries - """ - console = Console() - console.print("[yellow]Streaming Turtle file (line-by-line parser)[/yellow]") - - # Detect if file is compressed - is_gzipped = file_path.suffix == ".gz" - is_bz2 = file_path.suffix == ".bz2" or str(file_path).endswith(".ttl.bz2") - - if is_bz2: - import bz2 - - file_obj = bz2.open(file_path, "rt", encoding="utf-8", errors="ignore") - elif is_gzipped: - import gzip - - file_obj = gzip.open(file_path, "rt", encoding="utf-8", errors="ignore") - else: - file_obj = open(file_path, encoding="utf-8", errors="ignore") - - try: - current_chunk = [] - triple_count = 0 - line_count = 0 - - # Collect prefixes first - prefix_lines = [] - in_prefixes = True - - for line in file_obj: - line_count += 1 - stripped = line.strip() - - # Skip empty lines and comments - if not stripped or stripped.startswith("#"): - continue - - # Collect prefix declarations - if in_prefixes and (stripped.startswith("@prefix") or stripped.startswith("@base")): - prefix_lines.append(line) - continue - elif in_prefixes: - # End of prefixes, now we're in the data - in_prefixes = False - - # For simplicity, we'll parse in batches of lines ending with '.' - # This is a heuristic - proper Turtle parsing would need full context - if stripped.endswith("."): - # This might be end of a triple - triple_count += 1 - - if triple_count >= chunk_size: - # Try to parse this chunk - chunk_text = "".join(prefix_lines) + "\n" + "".join(current_chunk) + "\n" + line - - try: - graph = Graph() - graph.parse(data=chunk_text, format="turtle") - - # Extract triples - triples = [] - for s, p, o in graph: - if isinstance(o, Literal): - object_type = "literal" - object_datatype = str(o.datatype) if o.datatype else None - object_language = o.language if o.language else None - elif isinstance(o, URIRef): - object_type = "uri" - object_datatype = None - object_language = None - else: - object_type = "blank_node" - object_datatype = None - object_language = None - - triple = { - "subject": str(s), - "predicate": str(p), - "object": str(o), - "object_type": object_type, - "object_datatype": object_datatype, - "object_language": object_language, - } - triples.append(triple) - - if triples: - yield triples - - except Exception as e: - # If parsing fails, skip this chunk - logger.error(f"Error parsing Turtle chunk: {e}") - - # Reset for next chunk - current_chunk = [] - triple_count = 0 - else: - current_chunk.append(line) - else: - # Part of a multi-line statement - current_chunk.append(line) - - # Log progress - if line_count % 100000 == 0: - console.print(f"[dim]Processed {line_count:,} lines...[/dim]") - - # Process remaining lines - if current_chunk: - chunk_text = "".join(prefix_lines) + "\n" + "".join(current_chunk) - try: - graph = Graph() - graph.parse(data=chunk_text, format="turtle") - - triples = [] - for s, p, o in graph: - if isinstance(o, Literal): - object_type = "literal" - object_datatype = str(o.datatype) if o.datatype else None - object_language = o.language if o.language else None - elif isinstance(o, URIRef): - object_type = "uri" - object_datatype = None - object_language = None - else: - object_type = "blank_node" - object_datatype = None - object_language = None - - triple = { - "subject": str(s), - "predicate": str(p), - "object": str(o), - "object_type": object_type, - "object_datatype": object_datatype, - "object_language": object_language, - } - triples.append(triple) - - if triples: - yield triples - - except Exception as e: - logger.error(f"Error parsing final Turtle chunk: {e}") - - finally: - file_obj.close() - - -def stream_rdf_chunks( - file_path: Path, format: str = "turtle", chunk_size: int = 10000, num_workers: int | None = None -) -> Iterator[list[dict[str, str]]]: - """Stream RDF file in chunks based on format. - - Args: - file_path: Path to RDF file - format: RDF format - chunk_size: Number of triples per chunk - num_workers: Number of worker processes - - Yields: - Chunks of triple dictionaries - """ - console = Console() - - # Check if this is GeoNames format - is_geonames = "geonames" in str(file_path).lower() - if is_geonames and format in ("xml", "application/rdf+xml"): - with open(file_path, encoding="utf-8", errors="ignore") as f: - first_line = f.readline().strip() - if first_line.startswith("http://") or first_line.startswith("https://"): - yield from stream_geonames_parallel(file_path, chunk_size, num_workers) - return - - # Use appropriate streaming method based on format - if format in ("nt", "ntriples"): - yield from stream_ntriples_parallel(file_path, chunk_size, num_workers) - elif format in ("turtle", "ttl"): - # Use streaming Turtle parser - yield from stream_turtle_chunks(file_path, chunk_size) - else: - # For other formats that require full parsing, fall back to single-threaded - console.print(f"[yellow]Using standard RDF parser for {format} (single-threaded)[/yellow]") - console.print("[dim]Note: This may use significant memory for large files[/dim]") - - graph = Graph() - graph.parse(str(file_path), format=format) - - current_chunk = [] - for s, p, o in graph: - # Convert to dictionary format - if isinstance(o, Literal): - object_type = "literal" - object_datatype = str(o.datatype) if o.datatype else None - object_language = o.language if o.language else None - elif isinstance(o, URIRef): - object_type = "uri" - object_datatype = None - object_language = None - else: - object_type = "blank_node" - object_datatype = None - object_language = None - - triple = { - "subject": str(s), - "predicate": str(p), - "object": str(o), - "object_type": object_type, - "object_datatype": object_datatype, - "object_language": object_language, - } - current_chunk.append(triple) - - if len(current_chunk) >= chunk_size: - yield current_chunk - current_chunk = [] - - if current_chunk: - yield current_chunk - - -def convert_rdf_to_hf_streaming( - input_path: Path, - output_path: Path | None, - rdf_format: str = "turtle", - chunk_size: int = 10000, - num_workers: int | None = None, - metadata: dict[str, Any] | None = None, - push_to_hub: bool = False, - hub_repo_id: str | None = None, -) -> None: - """Convert an RDF file to HuggingFace dataset format using parallel streaming. - - Args: - input_path: Path to input RDF file - output_path: Path to output directory (optional if push_to_hub is True) - rdf_format: RDF format (turtle, nt, xml, etc.) - chunk_size: Number of triples to process at once - num_workers: Number of worker processes - metadata: Optional metadata dictionary - push_to_hub: If True, upload directly to HuggingFace Hub - hub_repo_id: Repository ID for HuggingFace Hub (required if push_to_hub is True) - """ - console = Console() - - if num_workers is None: - num_workers = mp.cpu_count() - - if push_to_hub and not hub_repo_id: - raise ValueError("hub_repo_id is required when push_to_hub is True") - - if not push_to_hub and output_path is None: - raise ValueError("output_path is required when push_to_hub is False") - - if output_path and not push_to_hub: - output_path.mkdir(parents=True, exist_ok=True) - - # Get file size for progress tracking - try: - file_size_mb = input_path.stat().st_size / (1024 * 1024) - except OSError as e: - console.print(f"[red]Error reading file size: {e}[/red]") - raise e - console.print(f"[cyan]Converting RDF file: {input_path.name}[/cyan]") - console.print(f"[dim]File size: {file_size_mb:.2f} MB ({rdf_format} format)[/dim]") - console.print(f"[dim]Chunk size: {chunk_size:,} lines per batch[/dim]") - console.print(f"[dim]Workers: {num_workers} CPU cores[/dim]") - if push_to_hub: - console.print(f"[dim]Destination: {hub_repo_id} (HuggingFace Hub)[/dim]\n") - else: - console.print(f"[dim]Output: {output_path}[/dim]\n") - - # Emit initial progress (streaming uses 0-100% scale directly) - print("\nPROGRESS: 5", flush=True) - - # Define schema - schema = pa.schema( - [ - ("subject", pa.string()), - ("predicate", pa.string()), - ("object", pa.string()), - ("object_type", pa.string()), - ("object_datatype", pa.string()), - ("object_language", pa.string()), - ] - ) - - # Process chunks and write directly to dataset - total_triples = 0 - chunk_idx = 0 - start_time = time.time() - - # Estimate total chunks based on file size (rough estimate) - estimated_chunks = max(10, int(file_size_mb * 1024 * 1024 / (chunk_size * 100))) - - console.print("\n[yellow]Creating dataset from streamed chunks...[/yellow]") - print("\nPROGRESS: 10", flush=True) - - # Note: Cannot use Rich progress bar inside generator - would capture - # unpicklable objects. Dataset.from_generator() requires picklable - # generators, so using simple print statements - def dataset_generator(): - triple_count = 0 - chunk_count = 0 - last_print_count = 0 - - for chunk in stream_rdf_chunks(input_path, rdf_format, chunk_size, num_workers): - chunk_count += 1 - for triple in chunk: - triple_count += 1 - yield triple - - if triple_count - last_print_count >= 100000: - elapsed = time.time() - start_time - rate = triple_count / elapsed if elapsed > 0 else 0 - progress_pct = 10 + min(60, int((chunk_count / estimated_chunks) * 60)) - print( - f" Processing: {chunk_count:,} chunks • " - f"{triple_count:,} triples • {rate:.0f} triples/sec", - flush=True, - ) - print(f"PROGRESS: {progress_pct}", flush=True) - last_print_count = triple_count - - # Create dataset using from_generator for true streaming - # dataset = Dataset.from_generator(dataset_generator) - features = Features({ - 'subject': Value('string'), - 'predicate': Value('string'), - 'object': Value('string'), - 'object_type': Value('string'), - 'object_datatype': Value('string'), - 'object_language': Value('string'), - }) - - dataset = Dataset.from_generator(dataset_generator, features=features) - - print("\nPROGRESS: 75", flush=True) - console.print("[green]✓ Dataset generation complete[/green]") - - # Get statistics after generation - total_triples = len(dataset) - - # Wrap in DatasetDict - dataset_dict = DatasetDict({"data": dataset}) - - # Add metadata - if metadata: - for split in dataset_dict: - dataset_dict[split].info.description = metadata.get("description", "") - dataset_dict[split].info.citation = metadata.get("citation", "") - dataset_dict[split].info.homepage = metadata.get("homepage", "") - dataset_dict[split].info.license = metadata.get("license", "") - - elapsed_time = time.time() - start_time - - if push_to_hub: - console.print(f"[yellow]Uploading to HuggingFace Hub: {hub_repo_id}...[/yellow]") - print("\nPROGRESS: 95", flush=True) - try: - dataset_dict.push_to_hub(hub_repo_id, private=False) - console.print(f"[bold green]✓ Successfully uploaded to {hub_repo_id}[/bold green]") - except Exception as e: - console.print(f"[red]Error uploading to Hub: {e}[/red]") - console.print(f"[yellow]Repository: {hub_repo_id}[/yellow]") - console.print("[yellow]Check authentication and permissions[/yellow]") - raise e - print("\nPROGRESS: 100", flush=True) - else: - assert output_path is not None - console.print("[yellow]Saving final dataset...[/yellow]") - print("\nPROGRESS: 95", flush=True) - try: - dataset_dict.save_to_disk(str(output_path), num_proc=num_workers) - except pa.ArrowInvalid as e: - console.print(f"[red]Error saving dataset: {e}[/red]") - console.print(f"[yellow]Output path: {output_path}[/yellow]") - console.print("[yellow]Check disk space and permissions[/yellow]") - raise e - print("\nPROGRESS: 100", flush=True) - - info = { - "format": "parquet", - "total_triples": total_triples, - "data_size": len(dataset_dict["data"]), - "source_format": rdf_format, - "chunk_size": chunk_size, - "num_workers": num_workers, - "processing_time_seconds": round(elapsed_time, 2), - "triples_per_second": round(total_triples / elapsed_time, 2) if elapsed_time > 0 else 0, - "conversion_date": time.strftime("%Y-%m-%d %H:%M:%S"), - } - - with open(output_path / "dataset_info.json", "w") as f: - json.dump(info, f, indent=2) - - console.print("\n[bold green]✓ Successfully converted to HuggingFace dataset[/bold green]") - console.print(f"[green] Data split: {len(dataset_dict['data']):,} triples[/green]") - console.print(f"[green] Processing time: {elapsed_time:.1f} seconds[/green]") - console.print(f"[green] Speed: {total_triples / elapsed_time:.0f} triples/second[/green]") - if push_to_hub: - console.print(f"[green] Repository: {hub_repo_id}[/green]") - else: - console.print(f"[green] Location: {output_path}[/green]") - - -def main(): - """Main entry point.""" - parser = argparse.ArgumentParser( - description="Convert RDF to HuggingFace dataset format with parallel streaming", - formatter_class=argparse.RawDescriptionHelpFormatter, - ) - - parser.add_argument("input", type=Path, help="Input RDF file") - parser.add_argument("output", nargs="?", type=Path, help="Output directory for HuggingFace dataset (optional if --push-to-hub)") - parser.add_argument( - "-f", - "--format", - default="turtle", - choices=["turtle", "nt", "ntriples", "xml", "n3", "trig", "nquads"], - help="RDF format (default: turtle)", - ) - parser.add_argument( - "--chunk-size", - type=int, - default=10000, - help="Number of lines per batch (default: 10000)", - ) - parser.add_argument( - "--num-workers", - type=int, - default=None, - help="Number of worker processes (default: CPU count)", - ) - parser.add_argument("--description", help="Dataset description") - parser.add_argument("--citation", help="Dataset citation") - parser.add_argument("--homepage", help="Dataset homepage URL") - parser.add_argument("--license", help="Dataset license") - parser.add_argument( - "--push-to-hub", - action="store_true", - help="Upload directly to HuggingFace Hub without saving to disk", - ) - parser.add_argument( - "--hub-repo-id", - type=str, - help="HuggingFace Hub repository ID (required if --push-to-hub)", - ) - parser.add_argument("-v", "--verbose", action="store_true", help="Show verbose output") - - args = parser.parse_args() - - # Configure logging - if args.verbose: - logging.basicConfig(level=logging.INFO) - else: - logging.basicConfig(level=logging.ERROR) - - # Validate input - if not args.input.exists(): - print(f"Error: Input file not found: {args.input}") - return 1 - - if args.push_to_hub and not args.hub_repo_id: - print("Error: --hub-repo-id is required when using --push-to-hub") - return 1 - - if not args.push_to_hub and not args.output: - print("Error: output directory is required when not using --push-to-hub") - return 1 - - # Prepare metadata - metadata = {} - if args.description: - metadata["description"] = args.description - if args.citation: - metadata["citation"] = args.citation - if args.homepage: - metadata["homepage"] = args.homepage - if args.license: - metadata["license"] = args.license - - # Convert - try: - convert_rdf_to_hf_streaming( - input_path=args.input, - output_path=args.output, - rdf_format=args.format, - chunk_size=args.chunk_size, - num_workers=args.num_workers, - metadata=metadata, - push_to_hub=args.push_to_hub, - hub_repo_id=args.hub_repo_id, - ) - return 0 - except Exception as e: - console = Console() - console.print(f"[red]Error: {e}[/red]") - if args.verbose: - import traceback - - traceback.print_exc() - return 1 - - -if __name__ == "__main__": - exit(main()) diff --git a/scripts/convert_rdf_to_hf_dataset_streaming_simple.py b/scripts/convert_rdf_to_hf_dataset_streaming_simple.py deleted file mode 100755 index 471c17f..0000000 --- a/scripts/convert_rdf_to_hf_dataset_streaming_simple.py +++ /dev/null @@ -1,672 +0,0 @@ -#!/usr/bin/env python3 -"""Simple streaming converter for RDF to HuggingFace dataset format with -minimal memory usage. - -This version uses a simpler single-threaded approach that's more reliable and still -memory-efficient. -""" - -from __future__ import annotations - -import argparse -import gzip -import json -import logging -import shutil -import time -from collections.abc import Iterator -from pathlib import Path -from typing import Any - -import pyarrow as pa -import pyarrow.parquet as pq -from datasets import Dataset, DatasetDict -from rdflib import Graph, Literal, URIRef -from rich.console import Console -from rich.progress import ( - BarColumn, - Progress, - SpinnerColumn, - TextColumn, - TimeElapsedColumn, -) - -logger = logging.getLogger(__name__) - - -def stream_geonames_simple( - file_path: Path, chunk_size: int = 10000 -) -> Iterator[list[dict[str, str]]]: - """Stream GeoNames RDF file with simple line-by-line processing. - - Args: - file_path: Path to GeoNames RDF file - chunk_size: Number of triples per chunk - - Yields: - Chunks of triple dictionaries - """ - current_chunk = [] - current_xml = [] - doc_count = 0 - - # Handle compressed files - if file_path.suffix == ".gz": - with gzip.open(file_path, "rt", encoding="utf-8") as file_obj: - for _line_no, line in enumerate(file_obj, 1): - # Check if this is a document boundary (URL line) - if line.startswith("http://") or line.startswith("https://"): - # Process previous document if exists - if current_xml: - xml_str = "".join(current_xml) - try: - # Parse the RDF/XML document - graph = Graph() - graph.parse(data=xml_str, format="xml") - - # Extract triples - for s, p, o in graph: - # Convert to dictionary format - if isinstance(o, Literal): - object_type = "literal" - object_datatype = ( - str(o.datatype) if o.datatype else None - ) - object_language = o.language if o.language else None - elif isinstance(o, URIRef): - object_type = "uri" - object_datatype = None - object_language = None - else: - object_type = "blank_node" - object_datatype = None - object_language = None - - triple = { - "subject": str(s), - "predicate": str(p), - "object": str(o), - "object_type": object_type, - "object_datatype": object_datatype, - "object_language": object_language, - } - current_chunk.append(triple) - - # Yield chunk if it's full - if len(current_chunk) >= chunk_size: - yield current_chunk - current_chunk = [] - - doc_count += 1 - if doc_count % 1000 == 0: - logger.info(f"Processed {doc_count} documents") - - except Exception as e: - # Skip malformed documents - logger.debug(f"Error parsing document {doc_count}: {e}") - - # Reset for next document - current_xml = [] - - # Skip the URL line itself - continue - - # Accumulate XML content - current_xml.append(line) - - # Process final document if exists - if current_xml: - xml_str = "".join(current_xml) - try: - graph = Graph() - graph.parse(data=xml_str, format="xml") - - for s, p, o in graph: - if isinstance(o, Literal): - object_type = "literal" - object_datatype = str(o.datatype) if o.datatype else None - object_language = o.language if o.language else None - elif isinstance(o, URIRef): - object_type = "uri" - object_datatype = None - object_language = None - else: - object_type = "blank_node" - object_datatype = None - object_language = None - - triple = { - "subject": str(s), - "predicate": str(p), - "object": str(o), - "object_type": object_type, - "object_datatype": object_datatype, - "object_language": object_language, - } - current_chunk.append(triple) - - if len(current_chunk) >= chunk_size: - yield current_chunk - current_chunk = [] - except Exception as e: - logger.debug(f"Error parsing final document: {e}") - - # Yield remaining triples - if current_chunk: - yield current_chunk - else: - with open(file_path, encoding="utf-8") as file_obj: - for _line_no, line in enumerate(file_obj, 1): - # Check if this is a document boundary (URL line) - if line.startswith("http://") or line.startswith("https://"): - # Process previous document if exists - if current_xml: - xml_str = "".join(current_xml) - try: - # Parse the RDF/XML document - graph = Graph() - graph.parse(data=xml_str, format="xml") - - # Extract triples - for s, p, o in graph: - # Convert to dictionary format - if isinstance(o, Literal): - object_type = "literal" - object_datatype = ( - str(o.datatype) if o.datatype else None - ) - object_language = ( - o.language if o.language else None - ) - elif isinstance(o, URIRef): - object_type = "uri" - object_datatype = None - object_language = None - else: - object_type = "blank_node" - object_datatype = None - object_language = None - - triple = { - "subject": str(s), - "predicate": str(p), - "object": str(o), - "object_type": object_type, - "object_datatype": object_datatype, - "object_language": object_language, - } - current_chunk.append(triple) - - # Yield chunk if it's full - if len(current_chunk) >= chunk_size: - yield current_chunk - current_chunk = [] - - doc_count += 1 - if doc_count % 1000 == 0: - logger.info(f"Processed {doc_count} documents") - - except Exception as e: - # Skip malformed documents - logger.debug(f"Error parsing document {doc_count}: {e}") - - # Reset for next document - current_xml = [] - - # Skip the URL line itself - continue - - # Accumulate XML content - current_xml.append(line) - - # Process final document if exists - if current_xml: - xml_str = "".join(current_xml) - try: - graph = Graph() - graph.parse(data=xml_str, format="xml") - - for s, p, o in graph: - if isinstance(o, Literal): - object_type = "literal" - object_datatype = str(o.datatype) if o.datatype else None - object_language = o.language if o.language else None - elif isinstance(o, URIRef): - object_type = "uri" - object_datatype = None - object_language = None - else: - object_type = "blank_node" - object_datatype = None - object_language = None - - triple = { - "subject": str(s), - "predicate": str(p), - "object": str(o), - "object_type": object_type, - "object_datatype": object_datatype, - "object_language": object_language, - } - current_chunk.append(triple) - - if len(current_chunk) >= chunk_size: - yield current_chunk - current_chunk = [] - except Exception as e: - logger.debug(f"Error parsing final document: {e}") - - # Yield remaining triples - if current_chunk: - yield current_chunk - - -def stream_ntriples_simple( - file_path: Path, chunk_size: int = 10000 -) -> Iterator[list[dict[str, str]]]: - """Stream N-Triples file line by line. - - Args: - file_path: Path to N-Triples file - chunk_size: Number of triples per chunk - - Yields: - Chunks of triple dictionaries - """ - current_chunk = [] - - # Handle compressed files - if file_path.suffix == ".gz": - with gzip.open(file_path, "rt", encoding="utf-8") as file_obj: - # Create a mini graph for parsing individual lines - for _line_no, line in enumerate(file_obj, 1): - line = line.strip() - if not line or line.startswith("#"): - continue - - try: - # Parse single N-Triple line - mini_graph = Graph() - mini_graph.parse(data=line, format="nt") - - for s, p, o in mini_graph: - # Convert to dictionary format - if isinstance(o, Literal): - object_type = "literal" - object_datatype = str(o.datatype) if o.datatype else None - object_language = o.language if o.language else None - elif isinstance(o, URIRef): - object_type = "uri" - object_datatype = None - object_language = None - else: - object_type = "blank_node" - object_datatype = None - object_language = None - - triple = { - "subject": str(s), - "predicate": str(p), - "object": str(o), - "object_type": object_type, - "object_datatype": object_datatype, - "object_language": object_language, - } - current_chunk.append(triple) - - if len(current_chunk) >= chunk_size: - yield current_chunk - current_chunk = [] - - except Exception as e: - # Skip malformed lines - logger.debug(f"Skipping malformed line {_line_no}: {e}") - continue - - # Yield remaining triples - if current_chunk: - yield current_chunk - else: - with open(file_path, encoding="utf-8") as file_obj: - # Create a mini graph for parsing individual lines - for _line_no, line in enumerate(file_obj, 1): - line = line.strip() - if not line or line.startswith("#"): - continue - - try: - # Parse single N-Triple line - mini_graph = Graph() - mini_graph.parse(data=line, format="nt") - - for s, p, o in mini_graph: - # Convert to dictionary format - if isinstance(o, Literal): - object_type = "literal" - object_datatype = str(o.datatype) if o.datatype else None - object_language = o.language if o.language else None - elif isinstance(o, URIRef): - object_type = "uri" - object_datatype = None - object_language = None - else: - object_type = "blank_node" - object_datatype = None - object_language = None - - triple = { - "subject": str(s), - "predicate": str(p), - "object": str(o), - "object_type": object_type, - "object_datatype": object_datatype, - "object_language": object_language, - } - current_chunk.append(triple) - - if len(current_chunk) >= chunk_size: - yield current_chunk - current_chunk = [] - - except Exception as e: - # Skip malformed lines - logger.debug(f"Skipping malformed line {_line_no}: {e}") - continue - - # Yield remaining triples - if current_chunk: - yield current_chunk - - -def stream_rdf_chunks( - file_path: Path, format: str = "turtle", chunk_size: int = 10000 -) -> Iterator[list[dict[str, str]]]: - """Stream RDF file in chunks based on format. - - Args: - file_path: Path to RDF file - format: RDF format - chunk_size: Number of triples per chunk - - Yields: - Chunks of triple dictionaries - """ - console = Console() - - # Check if this is GeoNames format - is_geonames = "geonames" in str(file_path).lower() - if is_geonames and format in ("xml", "application/rdf+xml"): - with open(file_path, encoding="utf-8", errors="ignore") as f: - first_line = f.readline().strip() - if first_line.startswith("http://") or first_line.startswith("https://"): - console.print( - "[yellow]Using GeoNames streaming parser (single-threaded)[/yellow]" - ) - yield from stream_geonames_simple(file_path, chunk_size) - return - - # Use appropriate streaming method based on format - if format in ("nt", "ntriples"): - console.print("[yellow]Using N-Triples streaming parser[/yellow]") - yield from stream_ntriples_simple(file_path, chunk_size) - else: - # For other formats, we need to parse the whole file but yield in chunks - console.print(f"[yellow]Using standard RDF parser for {format}[/yellow]") - console.print( - "[dim]Note: Non-streaming formats require loading the full graph[/dim]" - ) - - graph = Graph() - graph.parse(str(file_path), format=format) - - current_chunk = [] - for s, p, o in graph: - # Convert to dictionary format - if isinstance(o, Literal): - object_type = "literal" - object_datatype = str(o.datatype) if o.datatype else None - object_language = o.language if o.language else None - elif isinstance(o, URIRef): - object_type = "uri" - object_datatype = None - object_language = None - else: - object_type = "blank_node" - object_datatype = None - object_language = None - - triple = { - "subject": str(s), - "predicate": str(p), - "object": str(o), - "object_type": object_type, - "object_datatype": object_datatype, - "object_language": object_language, - } - current_chunk.append(triple) - - if len(current_chunk) >= chunk_size: - yield current_chunk - current_chunk = [] - - if current_chunk: - yield current_chunk - - -def convert_rdf_to_hf_streaming( - input_path: Path, - output_path: Path, - rdf_format: str = "turtle", - chunk_size: int = 10000, - metadata: dict[str, Any] | None = None, -) -> None: - """Convert an RDF file to HuggingFace dataset format using streaming. - - Args: - input_path: Path to input RDF file - output_path: Path to output directory - rdf_format: RDF format (turtle, nt, xml, etc.) - chunk_size: Number of triples to process at once - metadata: Optional metadata dictionary - """ - console = Console() - - # Create output directory - output_path.mkdir(parents=True, exist_ok=True) - - # Get file size for progress tracking - file_size_mb = input_path.stat().st_size / (1024 * 1024) - console.print(f"[cyan]Converting RDF file: {input_path.name}[/cyan]") - console.print(f"[dim]File size: {file_size_mb:.2f} MB ({rdf_format} format)[/dim]") - console.print(f"[dim]Chunk size: {chunk_size:,} triples[/dim]") - console.print(f"[dim]Output: {output_path}[/dim]\n") - - # Define schema - schema = pa.schema( - [ - ("subject", pa.string()), - ("predicate", pa.string()), - ("object", pa.string()), - ("object_type", pa.string()), - ("object_datatype", pa.string()), - ("object_language", pa.string()), - ] - ) - - # Process and write chunks - temp_chunks_dir = output_path / "temp_chunks" - temp_chunks_dir.mkdir(exist_ok=True) - - total_triples = 0 - chunk_count = 0 - parquet_files = [] - - with Progress( - SpinnerColumn(), - TextColumn("[progress.description]{task.description}"), - BarColumn(), - TextColumn("{task.completed:,} chunks"), - TextColumn("•"), - TextColumn("{task.fields[triples]:,} triples"), - TimeElapsedColumn(), - console=console, - transient=False, - ) as progress: - task = progress.add_task( - "[yellow]Processing RDF chunks...[/yellow]", - total=None, - triples=0, - ) - - # Stream and process chunks - for chunk in stream_rdf_chunks(input_path, rdf_format, chunk_size): - # Convert chunk to Arrow table - table = pa.Table.from_pylist(chunk, schema=schema) - - # Write chunk to Parquet file - chunk_file = temp_chunks_dir / f"chunk-{chunk_count:05d}.parquet" - pq.write_table(table, chunk_file) - parquet_files.append(chunk_file) - - # Update progress - total_triples += len(chunk) - chunk_count += 1 - progress.update( - task, - advance=1, - completed=chunk_count, - triples=total_triples, - description=( - f"[yellow]Processing chunks ({file_size_mb:.1f} MB " - f"file)...[/yellow]" - ), - ) - - console.print( - f"\n[green]✓ Processed {total_triples:,} triples in " - f"{chunk_count} chunks[/green]" - ) - - # Merge Parquet files into final dataset structure - console.print("\n[yellow]Merging chunks into final dataset...[/yellow]") - - # Read all Parquet files as a single dataset - dataset = Dataset.from_parquet(str(temp_chunks_dir / "*.parquet")) - - # Create DatasetDict with single 'data' split - dataset_dict = DatasetDict({"data": dataset}) - - # Add metadata - if metadata: - for split in dataset_dict: - dataset_dict[split].info.description = metadata.get("description", "") - dataset_dict[split].info.citation = metadata.get("citation", "") - dataset_dict[split].info.homepage = metadata.get("homepage", "") - dataset_dict[split].info.license = metadata.get("license", "") - - # Save final dataset - console.print("[yellow]Saving final dataset...[/yellow]") - dataset_dict.save_to_disk(str(output_path)) - - # Clean up temporary chunk files - console.print("[dim]Cleaning up temporary files...[/dim]") - shutil.rmtree(temp_chunks_dir) - - # Save dataset info - info = { - "format": "parquet", - "total_triples": total_triples, - "data_size": len(dataset_dict["data"]), - "source_format": rdf_format, - "chunk_size": chunk_size, - "conversion_date": time.strftime("%Y-%m-%d %H:%M:%S"), - } - - with open(output_path / "dataset_info.json", "w") as f: - json.dump(info, f, indent=2) - - console.print( - "\n[bold green]✓ Successfully converted to HuggingFace dataset[/bold green]" - ) - console.print(f"[green] Data split: {len(dataset_dict['data']):,} triples[/green]") - console.print(f"[green] Location: {output_path}[/green]") - - -def main(): - """Main entry point.""" - parser = argparse.ArgumentParser( - description="Convert RDF to HuggingFace dataset format with streaming", - formatter_class=argparse.RawDescriptionHelpFormatter, - ) - - parser.add_argument("input", type=Path, help="Input RDF file") - parser.add_argument( - "output", type=Path, help="Output directory for HuggingFace dataset" - ) - parser.add_argument( - "-f", - "--format", - default="turtle", - choices=["turtle", "nt", "ntriples", "xml", "n3", "trig", "nquads"], - help="RDF format (default: turtle)", - ) - parser.add_argument( - "--chunk-size", - type=int, - default=10000, - help="Number of triples to process at once (default: 10000)", - ) - parser.add_argument("--description", help="Dataset description") - parser.add_argument("--citation", help="Dataset citation") - parser.add_argument("--homepage", help="Dataset homepage URL") - parser.add_argument("--license", help="Dataset license") - parser.add_argument( - "-v", "--verbose", action="store_true", help="Show verbose output" - ) - - args = parser.parse_args() - - # Configure logging - if args.verbose: - logging.basicConfig(level=logging.INFO) - else: - logging.basicConfig(level=logging.ERROR) - - # Validate input - if not args.input.exists(): - print(f"Error: Input file not found: {args.input}") - return 1 - - # Prepare metadata - metadata = {} - if args.description: - metadata["description"] = args.description - if args.citation: - metadata["citation"] = args.citation - if args.homepage: - metadata["homepage"] = args.homepage - if args.license: - metadata["license"] = args.license - - # Convert - try: - convert_rdf_to_hf_streaming( - input_path=args.input, - output_path=args.output, - rdf_format=args.format, - chunk_size=args.chunk_size, - metadata=metadata, - ) - return 0 - except Exception as e: - console = Console() - console.print(f"[red]Error: {e}[/red]") - if args.verbose: - import traceback - - traceback.print_exc() - return 1 - - -if __name__ == "__main__": - exit(main()) diff --git a/scripts/convert_rdf_to_hf_dataset_streaming_turtle.py b/scripts/convert_rdf_to_hf_dataset_streaming_turtle.py deleted file mode 100755 index f906eaa..0000000 --- a/scripts/convert_rdf_to_hf_dataset_streaming_turtle.py +++ /dev/null @@ -1,496 +0,0 @@ -#!/usr/bin/env python3 -"""Streaming converter for RDF Turtle files to HuggingFace dataset format. - -This version implements proper streaming for Turtle files by parsing them in chunks. -Turtle format requires more careful parsing than N-Triples due to prefixes -and multi-line statements. -""" - -from __future__ import annotations - -import argparse -import gzip -import json -import logging -import re -import shutil -import time -from collections.abc import Iterator -from pathlib import Path -from typing import Any, cast - -import pyarrow as pa -import pyarrow.parquet as pq -from datasets import Dataset, DatasetDict -from rdflib import Graph, Literal, URIRef -from rich.console import Console -from rich.progress import ( - BarColumn, - Progress, - SpinnerColumn, - TextColumn, - TimeElapsedColumn, -) - -logger = logging.getLogger(__name__) - - -def parse_turtle_chunk( - chunk_text: str, prefixes: dict[str, str] -) -> list[dict[str, str]]: - """Parse a chunk of Turtle text with given prefixes. - - Args: - chunk_text: Turtle text chunk to parse - prefixes: Dictionary of prefix mappings - - Returns: - List of triple dictionaries - """ - triples = [] - - # Build complete Turtle document with prefixes - doc_lines = [] - for prefix, uri in prefixes.items(): - doc_lines.append(f"@prefix {prefix}: <{uri}> .") - doc_lines.append("") - doc_lines.append(chunk_text) - - try: - # Parse the chunk - graph = Graph() - graph.parse(data="\n".join(doc_lines), format="turtle") - - # Extract triples - for s, p, o in graph: - if isinstance(o, Literal): - object_type = "literal" - object_datatype = str(o.datatype) if o.datatype else None - object_language = o.language if o.language else None - elif isinstance(o, URIRef): - object_type = "uri" - object_datatype = None - object_language = None - else: - object_type = "blank_node" - object_datatype = None - object_language = None - - triple = { - "subject": str(s), - "predicate": str(p), - "object": str(o), - "object_type": object_type, - "object_datatype": object_datatype, - "object_language": object_language, - } - triples.append(triple) - except Exception as e: - # Log parse errors but continue - logger.debug(f"Error parsing chunk: {e}") - - return triples - - -def stream_turtle_file( - file_path: Path, chunk_size: int = 10000 -) -> Iterator[list[dict[str, str]]]: - """Stream Turtle file in chunks, handling prefixes and multi-line statements. - - Args: - file_path: Path to Turtle file - chunk_size: Approximate number of statements per chunk - - Yields: - Chunks of triple dictionaries - """ - console = Console() - console.print(f"[yellow]Streaming Turtle file: {file_path.name}[/yellow]") - - # Detect if file is gzipped - is_gzipped = file_path.suffix == ".gz" - - if is_gzipped: - with gzip.open( - file_path, "rt", encoding="utf-8", errors="ignore" - ) as file_obj: - prefixes = {} - current_chunk = [] - statement_count = 0 - line_count = 0 - in_multiline = False - multiline_buffer = [] - - for line in file_obj: - line_count += 1 - line = line.strip() - - # Skip empty lines and comments - if not line or line.startswith("#"): - continue - - # Handle prefix declarations - if line.startswith("@prefix"): - # Extract prefix and URI - match = re.match(r"@prefix\s+(\w+):\s*<([^>]+)>", line) - if match: - prefixes[match.group(1)] = match.group(2) - continue - elif line.startswith("@base"): - # Handle base declarations - continue - - # Handle multi-line statements (simple heuristic) - if in_multiline: - multiline_buffer.append(line) - if line.endswith("."): - # End of statement - complete_statement = " ".join(multiline_buffer) - current_chunk.append(complete_statement) - statement_count += 1 - in_multiline = False - multiline_buffer = [] - else: - # Check if this is start of a multi-line statement - if ( - not line.endswith(".") - and not line.endswith(";") - and not line.endswith(",") - ): - # Start of multi-line - in_multiline = True - multiline_buffer = [line] - else: - # Complete single-line statement - current_chunk.append(line) - statement_count += 1 - - # Process chunk when it reaches target size - if statement_count >= chunk_size: - # Parse and yield the chunk - chunk_text = "\n".join(current_chunk) - triples = parse_turtle_chunk(chunk_text, prefixes) - if triples: - yield triples - - # Reset for next chunk - current_chunk = [] - statement_count = 0 - - # Log progress periodically - if line_count % 100000 == 0: - console.print(f"[dim]Processed {line_count:,} lines...[/dim]") - - # Process remaining statements - if current_chunk: - chunk_text = "\n".join(current_chunk) - triples = parse_turtle_chunk(chunk_text, prefixes) - if triples: - yield triples - else: - with open(file_path, encoding="utf-8", errors="ignore") as file_obj: - prefixes = {} - current_chunk = [] - statement_count = 0 - line_count = 0 - in_multiline = False - multiline_buffer = [] - - for line in file_obj: - line_count += 1 - line = line.strip() - - # Skip empty lines and comments - if not line or line.startswith("#"): - continue - - # Handle prefix declarations - if line.startswith("@prefix"): - # Extract prefix and URI - match = re.match(r"@prefix\s+(\w+):\s*<([^>]+)>", line) - if match: - prefixes[match.group(1)] = match.group(2) - continue - elif line.startswith("@base"): - # Handle base declarations - continue - - # Handle multi-line statements (simple heuristic) - if in_multiline: - multiline_buffer.append(line) - if line.endswith("."): - # End of statement - complete_statement = " ".join(multiline_buffer) - current_chunk.append(complete_statement) - statement_count += 1 - in_multiline = False - multiline_buffer = [] - else: - # Check if this is start of a multi-line statement - if ( - not line.endswith(".") - and not line.endswith(";") - and not line.endswith(",") - ): - # Start of multi-line - in_multiline = True - multiline_buffer = [line] - else: - # Complete single-line statement - current_chunk.append(line) - statement_count += 1 - - # Process chunk when it reaches target size - if statement_count >= chunk_size: - # Parse and yield the chunk - chunk_text = "\n".join(current_chunk) - triples = parse_turtle_chunk(chunk_text, prefixes) - if triples: - yield triples - - # Reset for next chunk - current_chunk = [] - statement_count = 0 - - # Log progress periodically - if line_count % 100000 == 0: - console.print(f"[dim]Processed {line_count:,} lines...[/dim]") - - # Process remaining statements - if current_chunk: - chunk_text = "\n".join(current_chunk) - triples = parse_turtle_chunk(chunk_text, prefixes) - if triples: - yield triples - - -def process_turtle_chunk_parallel(args): - """Process a Turtle chunk in parallel (for future optimization). - - Module-level function for multiprocessing. - """ - chunk_text, prefixes = args - return parse_turtle_chunk(chunk_text, prefixes) - - -def convert_turtle_to_hf_streaming( - input_path: Path, - output_path: Path, - chunk_size: int = 10000, - metadata: dict[str, Any] | None = None, -) -> None: - """Convert a Turtle RDF file to HuggingFace dataset format using streaming. - - Args: - input_path: Path to input Turtle file - output_path: Path to output directory - chunk_size: Number of statements to process at once - metadata: Optional metadata dictionary - """ - console = Console() - - # Create output directory - output_path.mkdir(parents=True, exist_ok=True) - - # Get file size for progress tracking - file_size_mb = input_path.stat().st_size / (1024 * 1024) - console.print(f"[cyan]Converting Turtle RDF file: {input_path.name}[/cyan]") - console.print(f"[dim]File size: {file_size_mb:.2f} MB[/dim]") - console.print(f"[dim]Chunk size: {chunk_size:,} statements per batch[/dim]") - console.print(f"[dim]Output: {output_path}[/dim]\n") - - # Emit initial progress - print("\nPROGRESS: 5", flush=True) - - # Define schema - schema = pa.schema( - [ - ("subject", pa.string()), - ("predicate", pa.string()), - ("object", pa.string()), - ("object_type", pa.string()), - ("object_datatype", pa.string()), - ("object_language", pa.string()), - ] - ) - - # Process and write chunks - temp_chunks_dir = output_path / "temp_chunks" - temp_chunks_dir.mkdir(exist_ok=True) - - total_triples = 0 - chunk_count = 0 - parquet_files = [] - - with Progress( - SpinnerColumn(), - TextColumn("[progress.description]{task.description}"), - BarColumn(), - TextColumn("{task.completed:,} chunks"), - TextColumn("•"), - TextColumn("{task.fields[triples]:,} triples"), - TextColumn("•"), - TextColumn("{task.fields[rate]:.0f} triples/sec"), - TimeElapsedColumn(), - console=console, - transient=False, - ) as progress: - task = progress.add_task( - "[yellow]Processing Turtle chunks...[/yellow]", - total=None, - triples=0, - rate=0, - ) - - start_time = time.time() - - # Estimate total chunks based on file size - estimated_chunks = max(10, int(file_size_mb * 100)) # Rough estimate - - # Stream and process chunks - for _chunk_idx, chunk in enumerate( - stream_turtle_file(input_path, chunk_size) - ): - # Convert chunk to Arrow table - table = pa.Table.from_pylist(chunk, schema=schema) - - # Write chunk to Parquet file - chunk_file = temp_chunks_dir / f"chunk-{chunk_count:05d}.parquet" - pq.write_table(table, chunk_file) - parquet_files.append(chunk_file) - - # Update progress - total_triples += len(chunk) - chunk_count += 1 - elapsed = time.time() - start_time - rate = total_triples / elapsed if elapsed > 0 else 0 - - # Emit progress marker for wrapper scripts - progress_pct = 10 + min(70, (chunk_count / estimated_chunks) * 70) - if chunk_count % max(1, estimated_chunks // 20) == 0: - print(f"\nPROGRESS: {progress_pct:.0f}", flush=True) - - progress.update( - task, - advance=1, - completed=chunk_count, - triples=total_triples, - rate=rate, - description="[yellow]Processing Turtle file...[/yellow]", - ) - - console.print( - f"\n[green]✓ Processed {total_triples:,} triples in " - f"{chunk_count} chunks[/green]" - ) - print("\nPROGRESS: 80", flush=True) - - # Merge Parquet files into final dataset structure - console.print("\n[yellow]Merging chunks into final dataset...[/yellow]") - print("\nPROGRESS: 85", flush=True) - - # Read all Parquet files as a single dataset - # Note: from_parquet can return different types, but with a glob pattern - # it returns a Dataset that can be used in DatasetDict - dataset = cast(Dataset, Dataset.from_parquet(str(temp_chunks_dir / "*.parquet"))) - print("\nPROGRESS: 90", flush=True) - - # Create DatasetDict with single 'data' split - dataset_dict = DatasetDict({"data": dataset}) - - # Add metadata - if metadata: - for split in dataset_dict: - dataset_dict[split].info.description = metadata.get("description", "") - dataset_dict[split].info.citation = metadata.get("citation", "") - dataset_dict[split].info.homepage = metadata.get("homepage", "") - dataset_dict[split].info.license = metadata.get("license", "") - - # Save final dataset - console.print("[yellow]Saving final dataset...[/yellow]") - print("\nPROGRESS: 95", flush=True) - dataset_dict.save_to_disk(str(output_path)) - - # Clean up temporary chunk files - console.print("[dim]Cleaning up temporary files...[/dim]") - shutil.rmtree(temp_chunks_dir) - print("\nPROGRESS: 100", flush=True) - - # Save dataset info - elapsed_time = time.time() - start_time - info = { - "format": "parquet", - "total_triples": total_triples, - "data_size": len(dataset_dict["data"]), - "source_format": "turtle", - "chunk_size": chunk_size, - "processing_time_seconds": round(elapsed_time, 2), - "triples_per_second": round(total_triples / elapsed_time, 2) - if elapsed_time > 0 - else 0, - "conversion_date": time.strftime("%Y-%m-%d %H:%M:%S"), - } - - with open(output_path / "dataset_info.json", "w") as f: - json.dump(info, f, indent=2) - - console.print( - "\n[bold green]✓ Successfully converted to HuggingFace dataset[/bold green]" - ) - console.print(f"[green] Data split: {len(dataset_dict['data']):,} triples[/green]") - console.print(f"[green] Processing time: {elapsed_time:.1f} seconds[/green]") - console.print(f"[green] Output: {output_path}[/green]") - - -def main(): - """Main entry point for CLI usage.""" - parser = argparse.ArgumentParser( - description=( - "Convert RDF Turtle files to HuggingFace dataset format " - "with streaming" - ) - ) - parser.add_argument("input", type=Path, help="Input Turtle file path") - parser.add_argument("output", type=Path, help="Output dataset directory") - parser.add_argument( - "--chunk-size", - type=int, - default=10000, - help="Number of statements to process at once (default: 10000)", - ) - parser.add_argument("--description", type=str, help="Dataset description") - parser.add_argument("--homepage", type=str, help="Dataset homepage") - parser.add_argument("--license", type=str, help="Dataset license") - - args = parser.parse_args() - - # Validate input - if not args.input.exists(): - print(f"Error: Input file '{args.input}' not found") - return 1 - - # Prepare metadata - metadata = {} - if args.description: - metadata["description"] = args.description - if args.homepage: - metadata["homepage"] = args.homepage - if args.license: - metadata["license"] = args.license - - # Run conversion - try: - convert_turtle_to_hf_streaming( - args.input, - args.output, - chunk_size=args.chunk_size, - metadata=metadata, - ) - return 0 - except Exception as e: - print(f"Error: {e}") - return 1 - - -if __name__ == "__main__": - exit(main()) diff --git a/scripts/convert_rdf_to_hf_dataset_unified.py b/scripts/convert_rdf_to_hf_dataset_unified.py index 934d0de..b674ff7 100644 --- a/scripts/convert_rdf_to_hf_dataset_unified.py +++ b/scripts/convert_rdf_to_hf_dataset_unified.py @@ -51,13 +51,14 @@ from abc import ABC, abstractmethod from collections.abc import Iterator from contextlib import contextmanager from dataclasses import dataclass -from io import BytesIO +from io import BytesIO, TextIOWrapper, _WrappedBuffer from pathlib import Path from typing import IO, Any, TextIO, TypedDict, cast import pyarrow as pa import pyarrow.parquet as pq from datasets import Dataset, DatasetDict, Features, Value +from datasets.io.parquet import ParquetDatasetReader from rdflib import Graph, Literal, URIRef from rdflib.term import Node from rich.console import Console @@ -273,7 +274,7 @@ class SchemaManager: class FileHandler: """Handle file I/O with support for compressed files and error reporting.""" @staticmethod - def open_file(file_path: Path, mode: str = "r") -> IO[str] | IO[bytes]: + def open_file(file_path: Path, mode: str = "r") -> IO[str] | IO[bytes] | gzip.GzipFile | TextIOWrapper[_WrappedBuffer]: suffix = file_path.suffix.lower() try: if suffix == ".gz": @@ -707,6 +708,8 @@ class ConversionStrategy(ABC): logger.error(f"Error reading parquet chunks: {e}") raise + assert isinstance(dataset, Dataset) + if create_splits and config.create_train_test_split: train_test = dataset.train_test_split(test_size=config.test_size, seed=42) return DatasetDict({"train": train_test["train"], "test": train_test["test"]}) @@ -732,6 +735,7 @@ class ConversionStrategy(ABC): if extra_info: info.update(extra_info) try: + assert isinstance(config.output_path, Path) with open(config.output_path / "dataset_info.json", "w") as f: json.dump(info, f, indent=2) except Exception as e: @@ -835,7 +839,7 @@ class StandardStrategy(ConversionStrategy): strategy_used=self.name) self.progress.emit_progress(70) features = self.schema_manager.get_hf_features() - dataset = Dataset.from_list(triples, features=features) +git comm dataset = Dataset.from_list(triples, features=features) dataset_dict = DatasetDict({"data": dataset}) if config.metadata: dataset_dict = self.metadata_handler.add_metadata(dataset_dict, config.metadata) @@ -882,6 +886,7 @@ class StreamingStrategy(ConversionStrategy): if config.push_to_hub: temp_chunks_dir = Path(tempfile.mkdtemp(prefix="rdf_chunks_")) else: + assert isinstance(config.output_path, Path) config.output_path.mkdir(parents=True, exist_ok=True) temp_chunks_dir = config.output_path / "temp_chunks" temp_chunks_dir.mkdir(exist_ok=True) -- 2.52.0 From e7a540aeb378b50cd916ae0a923a1d1da6883116 Mon Sep 17 00:00:00 2001 From: Brent Edwards Date: Sun, 28 Dec 2025 17:13:34 -0800 Subject: [PATCH 3/8] More work toward `pyright` approval. --- ...rt_rdf_to_hf_dataset_streaming_parallel.py | 862 ++++++++++++++++++ scripts/convert_rdf_to_hf_dataset_unified.py | 21 +- scripts/dataset_downloader.py | 11 +- 3 files changed, 886 insertions(+), 8 deletions(-) create mode 100755 scripts/convert_rdf_to_hf_dataset_streaming_parallel.py diff --git a/scripts/convert_rdf_to_hf_dataset_streaming_parallel.py b/scripts/convert_rdf_to_hf_dataset_streaming_parallel.py new file mode 100755 index 0000000..4a488f5 --- /dev/null +++ b/scripts/convert_rdf_to_hf_dataset_streaming_parallel.py @@ -0,0 +1,862 @@ +#!/usr/bin/env python3 +"""Parallel streaming converter for RDF to HuggingFace dataset format with minimal memory usage. + +This version uses multiprocessing.Pool.imap for simpler and more reliable parallelization. +""" + +from __future__ import annotations + +import argparse +import gzip +import json +import logging +import multiprocessing as mp +import shutil +import sys +import time +from collections.abc import Iterator +from pathlib import Path +from typing import Any + +import pyarrow as pa +import pyarrow.parquet as pq +from datasets import Dataset, DatasetDict, Features, Value +from rdflib import Graph, Literal, URIRef +from rich.console import Console +from rich.progress import ( + BarColumn, + Progress, + SpinnerColumn, + TextColumn, + TimeElapsedColumn, +) + +logger = logging.getLogger(__name__) +logger.addHandler(logging.StreamHandler(sys.stdout)) + + +def process_geonames_lines(lines): + """Process a batch of lines from GeoNames file. + + Module-level function for multiprocessing. + + Args: + lines: List of lines containing one or more XML documents + + Returns: + List of triple dictionaries + """ + from rdflib import Graph, Literal, URIRef + + triples = [] + current_xml = [] + + for line in lines: + # Check if this is a document boundary (URL line) + if line.startswith("http://") or line.startswith("https://"): + # Process previous document if exists + if current_xml: + xml_str = "".join(current_xml) + try: + # Parse the RDF/XML document + graph = Graph() + graph.parse(data=xml_str, format="xml") + + # Extract triples + for s, p, o in graph: + # Convert to dictionary format + if isinstance(o, Literal): + object_type = "literal" + object_datatype = str(o.datatype) if o.datatype else None + object_language = o.language if o.language else None + elif isinstance(o, URIRef): + object_type = "uri" + object_datatype = None + object_language = None + else: + object_type = "blank_node" + object_datatype = None + object_language = None + + triple = { + "subject": str(s), + "predicate": str(p), + "object": str(o), + "object_type": object_type, + "object_datatype": object_datatype, + "object_language": object_language, + } + triples.append(triple) + + except Exception as e: + # Skip malformed documents + logger.error(f"Error parsing XML document: {e}. Skipping document.") + pass + + # Reset for next document + current_xml = [] + + # Skip the URL line itself + continue + + # Accumulate XML content + current_xml.append(line) + + # Process final document if exists + if current_xml: + xml_str = "".join(current_xml) + try: + graph = Graph() + graph.parse(data=xml_str, format="xml") + + for s, p, o in graph: + if isinstance(o, Literal): + object_type = "literal" + object_datatype = str(o.datatype) if o.datatype else None + object_language = o.language if o.language else None + elif isinstance(o, URIRef): + object_type = "uri" + object_datatype = None + object_language = None + else: + object_type = "blank_node" + object_datatype = None + object_language = None + + triple = { + "subject": str(s), + "predicate": str(p), + "object": str(o), + "object_type": object_type, + "object_datatype": object_datatype, + "object_language": object_language, + } + triples.append(triple) + except Exception as e: + logger.error(f"Error parsing final XML document: {e}. Skipping document.") + pass + + return triples + + +def process_ntriples_lines(lines): + """Process a batch of N-Triples lines. + + Args: + lines: List of N-Triple lines + + Returns: + List of triple dictionaries + """ + from rdflib import Graph, Literal, URIRef + + triples = [] + + for line in lines: + line = line.strip() + if not line or line.startswith("#"): + continue + + try: + # Parse single N-Triple line + mini_graph = Graph() + mini_graph.parse(data=line, format="nt") + + for s, p, o in mini_graph: + # Convert to dictionary format + if isinstance(o, Literal): + object_type = "literal" + object_datatype = str(o.datatype) if o.datatype else None + object_language = o.language if o.language else None + elif isinstance(o, URIRef): + object_type = "uri" + object_datatype = None + object_language = None + else: + object_type = "blank_node" + object_datatype = None + object_language = None + + triple = { + "subject": str(s), + "predicate": str(p), + "object": str(o), + "object_type": object_type, + "object_datatype": object_datatype, + "object_language": object_language, + } + triples.append(triple) + + except Exception as e: + # Skip malformed lines + logger.error(f"Error parsing N-Triple line: {line}. Skipping line. Error: {e}") + continue + + return triples + + +def batch_file_lines(file_path: Path, batch_size: int = 1000, format: str = "geonames"): + """Generator that yields batches of lines from file. + + For GeoNames format, ensures complete XML documents are kept together. + + Args: + file_path: Path to file + batch_size: Number of lines per batch + format: File format (geonames or ntriples) + + Yields: + Batches of lines + """ + # Handle compressed files + if file_path.suffix == ".gz": + try: + file_obj = gzip.open(file_path, "rt", encoding="utf-8") + # Test read to check if file is valid + test_line = file_obj.readline() + if not test_line and file_path.stat().st_size > 0: + raise EOFError("Compressed file appears to be empty or corrupted") + file_obj.seek(0) + except (gzip.BadGzipFile, EOFError, OSError) as e: + from rich.console import Console + + console = Console() + console.print(f"[red]Error: Compressed file is corrupted or incomplete: {e}[/red]") + console.print(f"[yellow]File: {file_path}[/yellow]") + console.print("[yellow]Please re-download the dataset or use an uncompressed version[/yellow]") + raise + else: + file_obj = open(file_path, encoding="utf-8") + + try: + current_batch = [] + current_doc = [] + + if format == "geonames": + # For GeoNames, keep documents together + for line in file_obj: + # Check if this is a document boundary + if line.startswith("http://") or line.startswith("https://"): + if current_doc: + # Add completed document to batch + current_batch.extend(current_doc) + current_doc = [line] + + # Yield batch if large enough + if len(current_batch) >= batch_size: + yield current_batch + current_batch = [] + else: + current_doc = [line] + else: + current_doc.append(line) + + # Add final document + if current_doc: + current_batch.extend(current_doc) + + # Yield final batch + if current_batch: + yield current_batch + + else: + # For N-Triples, just batch lines + for line in file_obj: + current_batch.append(line) + + if len(current_batch) >= batch_size: + yield current_batch + current_batch = [] + + # Yield remaining lines + if current_batch: + yield current_batch + + finally: + file_obj.close() + + +def stream_geonames_parallel( + file_path: Path, chunk_size: int = 10000, num_workers: int | None = None +) -> Iterator[list[dict[str, str]]]: + """Stream GeoNames RDF file with parallel processing using imap. + + Args: + file_path: Path to GeoNames RDF file + chunk_size: Number of lines per batch + num_workers: Number of worker processes + + Yields: + Chunks of triple dictionaries + """ + if num_workers is None: + num_workers = max(1, mp.cpu_count() - 1) + + console = Console() + console.print(f"[yellow]Using parallel GeoNames parser with {num_workers} workers[/yellow]") + + # Create batches of lines + batches = batch_file_lines(file_path, batch_size=chunk_size, format="geonames") + + # Process batches in parallel + try: + with mp.Pool(processes=num_workers) as pool: + # Use imap_unordered for better performance (order doesn't matter) + for triples in pool.imap_unordered(process_geonames_lines, batches, chunksize=1): + if triples: + yield triples + except Exception as e: + logger.error(f"Error processing GeoNames file: {e}") + raise e + + +def stream_ntriples_parallel( + file_path: Path, chunk_size: int = 10000, num_workers: int | None= None +) -> Iterator[list[dict[str, str]]]: + """Stream N-Triples file with parallel processing using imap. + + Args: + file_path: Path to N-Triples file + chunk_size: Number of lines per batch + num_workers: Number of worker processes + + Yields: + Chunks of triple dictionaries + """ + if num_workers is None: + num_workers = max(1, mp.cpu_count() - 1) + + console = Console() + console.print(f"[yellow]Using parallel N-Triples parser with {num_workers} workers[/yellow]") + + # Create batches of lines + batches = batch_file_lines(file_path, batch_size=chunk_size, format="ntriples") + + # Process batches in parallel + try: + with mp.Pool(processes=num_workers) as pool: + for triples in pool.imap_unordered(process_ntriples_lines, batches, chunksize=1): + if triples: + yield triples + except Exception as e: + logger.error(f"Error processing N-Triples file: {e}") + raise e + + +def stream_turtle_chunks(file_path: Path, chunk_size: int = 10000) -> Iterator[list[dict[str, str]]]: + """Stream Turtle file in chunks using incremental parsing. + + This is a simpler approach that reads the file line by line and yields chunks + without trying to parse the entire file at once. + + Args: + file_path: Path to Turtle file + chunk_size: Number of triples per chunk + + Yields: + Chunks of triple dictionaries + """ + console = Console() + console.print("[yellow]Streaming Turtle file (line-by-line parser)[/yellow]") + + # Detect if file is compressed + is_gzipped = file_path.suffix == ".gz" + is_bz2 = file_path.suffix == ".bz2" or str(file_path).endswith(".ttl.bz2") + + if is_bz2: + import bz2 + + file_obj = bz2.open(file_path, "rt", encoding="utf-8", errors="ignore") + elif is_gzipped: + import gzip + + file_obj = gzip.open(file_path, "rt", encoding="utf-8", errors="ignore") + else: + file_obj = open(file_path, encoding="utf-8", errors="ignore") + + try: + current_chunk = [] + triple_count = 0 + line_count = 0 + + # Collect prefixes first + prefix_lines = [] + in_prefixes = True + + for line in file_obj: + line_count += 1 + stripped = line.strip() + + # Skip empty lines and comments + if not stripped or stripped.startswith("#"): + continue + + # Collect prefix declarations + if in_prefixes and (stripped.startswith("@prefix") or stripped.startswith("@base")): + prefix_lines.append(line) + continue + elif in_prefixes: + # End of prefixes, now we're in the data + in_prefixes = False + + # For simplicity, we'll parse in batches of lines ending with '.' + # This is a heuristic - proper Turtle parsing would need full context + if stripped.endswith("."): + # This might be end of a triple + triple_count += 1 + + if triple_count >= chunk_size: + # Try to parse this chunk + chunk_text = "".join(prefix_lines) + "\n" + "".join(current_chunk) + "\n" + line + + try: + graph = Graph() + graph.parse(data=chunk_text, format="turtle") + + # Extract triples + triples = [] + for s, p, o in graph: + if isinstance(o, Literal): + object_type = "literal" + object_datatype = str(o.datatype) if o.datatype else None + object_language = o.language if o.language else None + elif isinstance(o, URIRef): + object_type = "uri" + object_datatype = None + object_language = None + else: + object_type = "blank_node" + object_datatype = None + object_language = None + + triple = { + "subject": str(s), + "predicate": str(p), + "object": str(o), + "object_type": object_type, + "object_datatype": object_datatype, + "object_language": object_language, + } + triples.append(triple) + + if triples: + yield triples + + except Exception as e: + # If parsing fails, skip this chunk + logger.error(f"Error parsing Turtle chunk: {e}") + + # Reset for next chunk + current_chunk = [] + triple_count = 0 + else: + current_chunk.append(line) + else: + # Part of a multi-line statement + current_chunk.append(line) + + # Log progress + if line_count % 100000 == 0: + console.print(f"[dim]Processed {line_count:,} lines...[/dim]") + + # Process remaining lines + if current_chunk: + chunk_text = "".join(prefix_lines) + "\n" + "".join(current_chunk) + try: + graph = Graph() + graph.parse(data=chunk_text, format="turtle") + + triples = [] + for s, p, o in graph: + if isinstance(o, Literal): + object_type = "literal" + object_datatype = str(o.datatype) if o.datatype else None + object_language = o.language if o.language else None + elif isinstance(o, URIRef): + object_type = "uri" + object_datatype = None + object_language = None + else: + object_type = "blank_node" + object_datatype = None + object_language = None + + triple = { + "subject": str(s), + "predicate": str(p), + "object": str(o), + "object_type": object_type, + "object_datatype": object_datatype, + "object_language": object_language, + } + triples.append(triple) + + if triples: + yield triples + + except Exception as e: + logger.error(f"Error parsing final Turtle chunk: {e}") + + finally: + file_obj.close() + + +def stream_rdf_chunks( + file_path: Path, format: str = "turtle", chunk_size: int = 10000, num_workers: int | None = None +) -> Iterator[list[dict[str, str]]]: + """Stream RDF file in chunks based on format. + + Args: + file_path: Path to RDF file + format: RDF format + chunk_size: Number of triples per chunk + num_workers: Number of worker processes + + Yields: + Chunks of triple dictionaries + """ + console = Console() + + # Check if this is GeoNames format + is_geonames = "geonames" in str(file_path).lower() + if is_geonames and format in ("xml", "application/rdf+xml"): + with open(file_path, encoding="utf-8", errors="ignore") as f: + first_line = f.readline().strip() + if first_line.startswith("http://") or first_line.startswith("https://"): + yield from stream_geonames_parallel(file_path, chunk_size, num_workers) + return + + # Use appropriate streaming method based on format + if format in ("nt", "ntriples"): + yield from stream_ntriples_parallel(file_path, chunk_size, num_workers) + elif format in ("turtle", "ttl"): + # Use streaming Turtle parser + yield from stream_turtle_chunks(file_path, chunk_size) + else: + # For other formats that require full parsing, fall back to single-threaded + console.print(f"[yellow]Using standard RDF parser for {format} (single-threaded)[/yellow]") + console.print("[dim]Note: This may use significant memory for large files[/dim]") + + graph = Graph() + graph.parse(str(file_path), format=format) + + current_chunk = [] + for s, p, o in graph: + # Convert to dictionary format + if isinstance(o, Literal): + object_type = "literal" + object_datatype = str(o.datatype) if o.datatype else None + object_language = o.language if o.language else None + elif isinstance(o, URIRef): + object_type = "uri" + object_datatype = None + object_language = None + else: + object_type = "blank_node" + object_datatype = None + object_language = None + + triple = { + "subject": str(s), + "predicate": str(p), + "object": str(o), + "object_type": object_type, + "object_datatype": object_datatype, + "object_language": object_language, + } + current_chunk.append(triple) + + if len(current_chunk) >= chunk_size: + yield current_chunk + current_chunk = [] + + if current_chunk: + yield current_chunk + + +def convert_rdf_to_hf_streaming( + input_path: Path, + output_path: Path | None, + rdf_format: str = "turtle", + chunk_size: int = 10000, + num_workers: int | None = None, + metadata: dict[str, Any] | None = None, + push_to_hub: bool = False, + hub_repo_id: str | None = None, +) -> None: + """Convert an RDF file to HuggingFace dataset format using parallel streaming. + + Args: + input_path: Path to input RDF file + output_path: Path to output directory (optional if push_to_hub is True) + rdf_format: RDF format (turtle, nt, xml, etc.) + chunk_size: Number of triples to process at once + num_workers: Number of worker processes + metadata: Optional metadata dictionary + push_to_hub: If True, upload directly to HuggingFace Hub + hub_repo_id: Repository ID for HuggingFace Hub (required if push_to_hub is True) + """ + console = Console() + + if num_workers is None: + num_workers = mp.cpu_count() + + if push_to_hub and not hub_repo_id: + raise ValueError("hub_repo_id is required when push_to_hub is True") + + if not push_to_hub and output_path is None: + raise ValueError("output_path is required when push_to_hub is False") + + if output_path and not push_to_hub: + output_path.mkdir(parents=True, exist_ok=True) + + # Get file size for progress tracking + try: + file_size_mb = input_path.stat().st_size / (1024 * 1024) + except OSError as e: + console.print(f"[red]Error reading file size: {e}[/red]") + raise e + console.print(f"[cyan]Converting RDF file: {input_path.name}[/cyan]") + console.print(f"[dim]File size: {file_size_mb:.2f} MB ({rdf_format} format)[/dim]") + console.print(f"[dim]Chunk size: {chunk_size:,} lines per batch[/dim]") + console.print(f"[dim]Workers: {num_workers} CPU cores[/dim]") + if push_to_hub: + console.print(f"[dim]Destination: {hub_repo_id} (HuggingFace Hub)[/dim]\n") + else: + console.print(f"[dim]Output: {output_path}[/dim]\n") + + # Emit initial progress (streaming uses 0-100% scale directly) + print("\nPROGRESS: 5", flush=True) + + # Define schema + schema = pa.schema( + [ + ("subject", pa.string()), + ("predicate", pa.string()), + ("object", pa.string()), + ("object_type", pa.string()), + ("object_datatype", pa.string()), + ("object_language", pa.string()), + ] + ) + + # Process chunks and write directly to dataset + total_triples = 0 + chunk_idx = 0 + start_time = time.time() + + # Estimate total chunks based on file size (rough estimate) + estimated_chunks = max(10, int(file_size_mb * 1024 * 1024 / (chunk_size * 100))) + + console.print("\n[yellow]Creating dataset from streamed chunks...[/yellow]") + print("\nPROGRESS: 10", flush=True) + + # Note: Cannot use Rich progress bar inside generator - would capture + # unpicklable objects. Dataset.from_generator() requires picklable + # generators, so using simple print statements + def dataset_generator(): + triple_count = 0 + chunk_count = 0 + last_print_count = 0 + + for chunk in stream_rdf_chunks(input_path, rdf_format, chunk_size, num_workers): + chunk_count += 1 + for triple in chunk: + triple_count += 1 + yield triple + + if triple_count - last_print_count >= 100000: + elapsed = time.time() - start_time + rate = triple_count / elapsed if elapsed > 0 else 0 + progress_pct = 10 + min(60, int((chunk_count / estimated_chunks) * 60)) + print( + f" Processing: {chunk_count:,} chunks • " + f"{triple_count:,} triples • {rate:.0f} triples/sec", + flush=True, + ) + print(f"PROGRESS: {progress_pct}", flush=True) + last_print_count = triple_count + + # Create dataset using from_generator for true streaming + # dataset = Dataset.from_generator(dataset_generator) + features = Features({ + 'subject': Value('string'), + 'predicate': Value('string'), + 'object': Value('string'), + 'object_type': Value('string'), + 'object_datatype': Value('string'), + 'object_language': Value('string'), + }) + + dataset = Dataset.from_generator(dataset_generator, features=features) + + print("\nPROGRESS: 75", flush=True) + console.print("[green]✓ Dataset generation complete[/green]") + + # Get statistics after generation + total_triples = len(dataset) + + # Wrap in DatasetDict + dataset_dict = DatasetDict({"data": dataset}) + + # Add metadata + if metadata: + for split in dataset_dict: + dataset_dict[split].info.description = metadata.get("description", "") + dataset_dict[split].info.citation = metadata.get("citation", "") + dataset_dict[split].info.homepage = metadata.get("homepage", "") + dataset_dict[split].info.license = metadata.get("license", "") + + elapsed_time = time.time() - start_time + + if push_to_hub: + console.print(f"[yellow]Uploading to HuggingFace Hub: {hub_repo_id}...[/yellow]") + print("\nPROGRESS: 95", flush=True) + try: + dataset_dict.push_to_hub(hub_repo_id, private=False) + console.print(f"[bold green]✓ Successfully uploaded to {hub_repo_id}[/bold green]") + except Exception as e: + console.print(f"[red]Error uploading to Hub: {e}[/red]") + console.print(f"[yellow]Repository: {hub_repo_id}[/yellow]") + console.print("[yellow]Check authentication and permissions[/yellow]") + raise e + print("\nPROGRESS: 100", flush=True) + else: + assert output_path is not None + console.print("[yellow]Saving final dataset...[/yellow]") + print("\nPROGRESS: 95", flush=True) + try: + dataset_dict.save_to_disk(str(output_path), num_proc=num_workers) + except pa.ArrowInvalid as e: + console.print(f"[red]Error saving dataset: {e}[/red]") + console.print(f"[yellow]Output path: {output_path}[/yellow]") + console.print("[yellow]Check disk space and permissions[/yellow]") + raise e + print("\nPROGRESS: 100", flush=True) + + info = { + "format": "parquet", + "total_triples": total_triples, + "data_size": len(dataset_dict["data"]), + "source_format": rdf_format, + "chunk_size": chunk_size, + "num_workers": num_workers, + "processing_time_seconds": round(elapsed_time, 2), + "triples_per_second": round(total_triples / elapsed_time, 2) if elapsed_time > 0 else 0, + "conversion_date": time.strftime("%Y-%m-%d %H:%M:%S"), + } + + with open(output_path / "dataset_info.json", "w") as f: + json.dump(info, f, indent=2) + + console.print("\n[bold green]✓ Successfully converted to HuggingFace dataset[/bold green]") + console.print(f"[green] Data split: {len(dataset_dict['data']):,} triples[/green]") + console.print(f"[green] Processing time: {elapsed_time:.1f} seconds[/green]") + console.print(f"[green] Speed: {total_triples / elapsed_time:.0f} triples/second[/green]") + if push_to_hub: + console.print(f"[green] Repository: {hub_repo_id}[/green]") + else: + console.print(f"[green] Location: {output_path}[/green]") + + +def main(): + """Main entry point.""" + parser = argparse.ArgumentParser( + description="Convert RDF to HuggingFace dataset format with parallel streaming", + formatter_class=argparse.RawDescriptionHelpFormatter, + ) + + parser.add_argument("input", type=Path, help="Input RDF file") + parser.add_argument("output", nargs="?", type=Path, help="Output directory for HuggingFace dataset (optional if --push-to-hub)") + parser.add_argument( + "-f", + "--format", + default="turtle", + choices=["turtle", "nt", "ntriples", "xml", "n3", "trig", "nquads"], + help="RDF format (default: turtle)", + ) + parser.add_argument( + "--chunk-size", + type=int, + default=10000, + help="Number of lines per batch (default: 10000)", + ) + parser.add_argument( + "--num-workers", + type=int, + default=None, + help="Number of worker processes (default: CPU count)", + ) + parser.add_argument("--description", help="Dataset description") + parser.add_argument("--citation", help="Dataset citation") + parser.add_argument("--homepage", help="Dataset homepage URL") + parser.add_argument("--license", help="Dataset license") + parser.add_argument( + "--push-to-hub", + action="store_true", + help="Upload directly to HuggingFace Hub without saving to disk", + ) + parser.add_argument( + "--hub-repo-id", + type=str, + help="HuggingFace Hub repository ID (required if --push-to-hub)", + ) + parser.add_argument("-v", "--verbose", action="store_true", help="Show verbose output") + + args = parser.parse_args() + + # Configure logging + if args.verbose: + logging.basicConfig(level=logging.INFO) + else: + logging.basicConfig(level=logging.ERROR) + + # Validate input + if not args.input.exists(): + print(f"Error: Input file not found: {args.input}") + return 1 + + if args.push_to_hub and not args.hub_repo_id: + print("Error: --hub-repo-id is required when using --push-to-hub") + return 1 + + if not args.push_to_hub and not args.output: + print("Error: output directory is required when not using --push-to-hub") + return 1 + + # Prepare metadata + metadata = {} + if args.description: + metadata["description"] = args.description + if args.citation: + metadata["citation"] = args.citation + if args.homepage: + metadata["homepage"] = args.homepage + if args.license: + metadata["license"] = args.license + + # Convert + try: + convert_rdf_to_hf_streaming( + input_path=args.input, + output_path=args.output, + rdf_format=args.format, + chunk_size=args.chunk_size, + num_workers=args.num_workers, + metadata=metadata, + push_to_hub=args.push_to_hub, + hub_repo_id=args.hub_repo_id, + ) + return 0 + except Exception as e: + console = Console() + console.print(f"[red]Error: {e}[/red]") + if args.verbose: + import traceback + + traceback.print_exc() + return 1 + + +if __name__ == "__main__": + exit(main()) diff --git a/scripts/convert_rdf_to_hf_dataset_unified.py b/scripts/convert_rdf_to_hf_dataset_unified.py index b674ff7..8ddfb00 100644 --- a/scripts/convert_rdf_to_hf_dataset_unified.py +++ b/scripts/convert_rdf_to_hf_dataset_unified.py @@ -57,7 +57,7 @@ from typing import IO, Any, TextIO, TypedDict, cast import pyarrow as pa import pyarrow.parquet as pq -from datasets import Dataset, DatasetDict, Features, Value +from datasets import Dataset, DatasetDict, Features, Value, IterableDataset from datasets.io.parquet import ParquetDatasetReader from rdflib import Graph, Literal, URIRef from rdflib.term import Node @@ -839,7 +839,7 @@ class StandardStrategy(ConversionStrategy): strategy_used=self.name) self.progress.emit_progress(70) features = self.schema_manager.get_hf_features() -git comm dataset = Dataset.from_list(triples, features=features) + dataset = Dataset.from_list(triples, features=features) # pyright: ignore dataset_dict = DatasetDict({"data": dataset}) if config.metadata: dataset_dict = self.metadata_handler.add_metadata(dataset_dict, config.metadata) @@ -952,6 +952,8 @@ class StreamingTurtleStrategy(ConversionStrategy): if config.push_to_hub: temp_chunks_dir = Path(tempfile.mkdtemp(prefix="rdf_chunks_")) else: + assert config.output_path is not None + config.output_path.mkdir(parents=True, exist_ok=True) temp_chunks_dir = config.output_path / "temp_chunks" temp_chunks_dir.mkdir(exist_ok=True) @@ -1026,6 +1028,8 @@ class SimpleStreamingStrategy(ConversionStrategy): self.progress.print(f"Output: {config.output_path}", "dim") # Create temp directory for chunks + assert config.output_path is not None + if config.push_to_hub: temp_chunks_dir = Path(tempfile.mkdtemp(prefix="rdf_chunks_")) else: @@ -1126,7 +1130,7 @@ class ParallelStreamingStrategy(ConversionStrategy): # Dataset.from_generator() requires picklable generators. # Must call streaming functions directly inside the generator. def dataset_generator(): - triple_count = 0 + nonlocal triple_count chunk_count = 0 last_print_count = 0 @@ -1162,16 +1166,20 @@ class ParallelStreamingStrategy(ConversionStrategy): # Create dataset using from_generator for true streaming features = self.schema_manager.get_hf_features() + triple_count = 0 dataset = Dataset.from_generator(dataset_generator, features=features) self.progress.emit_progress(75) self.progress.print("[green]✓ Dataset generation complete[/green]") # Get statistics after generation - total_triples = len(dataset) + if not isinstance(dataset, IterableDataset): + total_triples = len(dataset) + else: + total_triples = triple_count # Wrap in DatasetDict - dataset_dict = DatasetDict({"data": dataset}) + dataset_dict = DatasetDict({"data": dataset}) # pyright: ignore # Add metadata if config.metadata: @@ -1188,6 +1196,9 @@ class ParallelStreamingStrategy(ConversionStrategy): else: self.progress.print("[yellow]Saving final dataset...[/yellow]") self.progress.emit_progress(95) + + assert config.output_path is not None + config.output_path.mkdir(parents=True, exist_ok=True) dataset_dict.save_to_disk(str(config.output_path), num_proc=num_workers) self.progress.emit_progress(100) diff --git a/scripts/dataset_downloader.py b/scripts/dataset_downloader.py index 39d4e02..eeb60db 100644 --- a/scripts/dataset_downloader.py +++ b/scripts/dataset_downloader.py @@ -133,9 +133,14 @@ class DatasetDownloader: for split_name, ds in dataset.items(): self.console.print(f" • {split_name}: {len(ds):,} rows") # type: ignore[arg-type] else: # Single Dataset - self.console.print( - f"\n[green]✓ Downloaded dataset with {len(dataset):,} rows[/green]" - ) # type: ignore[arg-type] + if isinstance(dataset, Dataset): + self.console.print( + f"\n[green]✓ Downloaded dataset with {len(dataset):,} rows[/green]" + ) # type: ignore[arg-type] + else: + self.console.print( + f"\n[green]✓ Downloaded dataset" + ) self.console.print(f"\n[cyan]Dataset cached at: {self.output_dir}[/cyan]") return dataset # type: ignore[return-value] -- 2.52.0 From 90d77a7077f2749d35708f0182a8d0844457ac65 Mon Sep 17 00:00:00 2001 From: Aditya Chhabra Date: Mon, 29 Dec 2025 18:47:22 +0530 Subject: [PATCH 4/8] fix: fix chunksize argument in multiprocessing as per review --- scripts/convert_rdf_to_hf_dataset_unified.py | 84 ++------------------ 1 file changed, 8 insertions(+), 76 deletions(-) diff --git a/scripts/convert_rdf_to_hf_dataset_unified.py b/scripts/convert_rdf_to_hf_dataset_unified.py index 8ddfb00..8d37ad7 100644 --- a/scripts/convert_rdf_to_hf_dataset_unified.py +++ b/scripts/convert_rdf_to_hf_dataset_unified.py @@ -281,7 +281,7 @@ class FileHandler: if "b" in mode: return gzip.open(file_path, mode) return gzip.open(file_path, "rt", encoding="utf-8") - elif suffix == ".bz2" or str(file_path).endswith(".ttl.bz2"): + elif suffix == ".bz2": if "b" in mode: return bz2.open(file_path, mode) return bz2.open(file_path, "rt", encoding="utf-8") @@ -390,7 +390,7 @@ def stream_geonames_parallel( try: with mp.Pool(processes=num_workers) as pool: - for triples in pool.imap_unordered(process_geonames_lines, batches, chunksize=1): + for triples in pool.imap_unordered(process_geonames_lines, batches, chunksize=10): if triples: yield triples except Exception as e: @@ -412,7 +412,7 @@ def stream_ntriples_parallel( try: with mp.Pool(processes=num_workers) as pool: - for triples in pool.imap_unordered(process_ntriples_lines, batches, chunksize=1): + for triples in pool.imap_unordered(process_ntriples_lines, batches, chunksize=10): if triples: yield triples except Exception as e: @@ -425,13 +425,11 @@ def stream_turtle_chunks(file_path: Path, chunk_size: int = 10000) -> Iterator[l Module-level function for use in picklable generators. """ - is_gzipped = file_path.suffix == ".gz" - is_bz2 = file_path.suffix == ".bz2" or str(file_path).endswith(".ttl.bz2") - if is_bz2: + if file_path.suffix == ".bz2": import bz2 file_obj = bz2.open(file_path, "rt", encoding="utf-8", errors="ignore") - elif is_gzipped: + elif file_path.suffix == ".gz": import gzip file_obj = gzip.open(file_path, "rt", encoding="utf-8", errors="ignore") else: @@ -464,30 +462,7 @@ def stream_turtle_chunks(file_path: Path, chunk_size: int = 10000) -> Iterator[l graph = Graph() graph.parse(data=chunk_text, format="turtle") - triples = [] - for s, p, o in graph: - if isinstance(o, Literal): - object_type = "literal" - object_datatype = str(o.datatype) if o.datatype else None - object_language = o.language if o.language else None - elif isinstance(o, URIRef): - object_type = "uri" - object_datatype = None - object_language = None - else: - object_type = "blank_node" - object_datatype = None - object_language = None - - triple = { - "subject": str(s), - "predicate": str(p), - "object": str(o), - "object_type": object_type, - "object_datatype": object_datatype, - "object_language": object_language, - } - triples.append(triple) + triples = extract_triples_from_graph(graph) if triples: yield triples @@ -509,30 +484,7 @@ def stream_turtle_chunks(file_path: Path, chunk_size: int = 10000) -> Iterator[l graph = Graph() graph.parse(data=chunk_text, format="turtle") - triples = [] - for s, p, o in graph: - if isinstance(o, Literal): - object_type = "literal" - object_datatype = str(o.datatype) if o.datatype else None - object_language = o.language if o.language else None - elif isinstance(o, URIRef): - object_type = "uri" - object_datatype = None - object_language = None - else: - object_type = "blank_node" - object_datatype = None - object_language = None - - triple = { - "subject": str(s), - "predicate": str(p), - "object": str(o), - "object_type": object_type, - "object_datatype": object_datatype, - "object_language": object_language, - } - triples.append(triple) + triples = extract_triples_from_graph(graph) if triples: yield triples @@ -554,27 +506,7 @@ def stream_generic_rdf(file_path: Path, rdf_format: str, chunk_size: int = 10000 current_chunk = [] for s, p, o in graph: - if isinstance(o, Literal): - object_type = "literal" - object_datatype = str(o.datatype) if o.datatype else None - object_language = o.language if o.language else None - elif isinstance(o, URIRef): - object_type = "uri" - object_datatype = None - object_language = None - else: - object_type = "blank_node" - object_datatype = None - object_language = None - - triple = { - "subject": str(s), - "predicate": str(p), - "object": str(o), - "object_type": object_type, - "object_datatype": object_datatype, - "object_language": object_language, - } + triple = extract_triple(s, p, o) current_chunk.append(triple) if len(current_chunk) >= chunk_size: -- 2.52.0 From 2b1537d4f9afc0ef65d28617cfee41d9a7c6011c Mon Sep 17 00:00:00 2001 From: Aditya Chhabra Date: Mon, 29 Dec 2025 18:49:57 +0530 Subject: [PATCH 5/8] fix: fix ruff check errors --- scripts/convert_rdf_to_hf_dataset_unified.py | 382 ++++++++++++++----- scripts/upload_all_datasets.py | 21 +- 2 files changed, 298 insertions(+), 105 deletions(-) diff --git a/scripts/convert_rdf_to_hf_dataset_unified.py b/scripts/convert_rdf_to_hf_dataset_unified.py index 8d37ad7..f32b6ff 100644 --- a/scripts/convert_rdf_to_hf_dataset_unified.py +++ b/scripts/convert_rdf_to_hf_dataset_unified.py @@ -41,24 +41,21 @@ import gzip import json import logging import multiprocessing as mp -import re import shutil import sys import tempfile -import threading import time from abc import ABC, abstractmethod from collections.abc import Iterator from contextlib import contextmanager from dataclasses import dataclass -from io import BytesIO, TextIOWrapper, _WrappedBuffer +from io import TextIOWrapper from pathlib import Path from typing import IO, Any, TextIO, TypedDict, cast import pyarrow as pa import pyarrow.parquet as pq -from datasets import Dataset, DatasetDict, Features, Value, IterableDataset -from datasets.io.parquet import ParquetDatasetReader +from datasets import Dataset, DatasetDict, Features, IterableDataset, Value from rdflib import Graph, Literal, URIRef from rdflib.term import Node from rich.console import Console @@ -87,7 +84,7 @@ def process_geonames_lines(lines): Returns: List of triple dictionaries """ - from rdflib import Graph, Literal, URIRef + from rdflib import Graph triples = [] current_xml = [] @@ -165,7 +162,10 @@ def process_ntriples_lines(lines): except Exception as e: # Skip malformed lines - logger.error(f"Error parsing N-Triple line: {line}. Skipping line. Error: {e}") + logger.error( + f"Error parsing N-Triple line: {line}. " + f"Skipping line. Error: {e}" + ) continue return triples @@ -274,7 +274,9 @@ class SchemaManager: class FileHandler: """Handle file I/O with support for compressed files and error reporting.""" @staticmethod - def open_file(file_path: Path, mode: str = "r") -> IO[str] | IO[bytes] | gzip.GzipFile | TextIOWrapper[_WrappedBuffer]: + def open_file( + file_path: Path, mode: str = "r" + ) -> IO[str] | IO[bytes] | gzip.GzipFile | TextIO: suffix = file_path.suffix.lower() try: if suffix == ".gz": @@ -325,9 +327,14 @@ def batch_file_lines(file_path: Path, batch_size: int = 1000, format: str = "geo file_obj.seek(0) except (gzip.BadGzipFile, EOFError, OSError) as e: console = Console() - console.print(f"[red]Error: Compressed file is corrupted or incomplete: {e}[/red]") + console.print( + f"[red]Error: Compressed file is corrupted or incomplete: {e}[/red]" + ) console.print(f"[yellow]File: {file_path}[/yellow]") - console.print("[yellow]Please re-download the dataset or use an uncompressed version[/yellow]") + console.print( + "[yellow]Please re-download the dataset or use an " + "uncompressed version[/yellow]" + ) raise else: file_obj = open(file_path, encoding="utf-8") @@ -390,7 +397,9 @@ def stream_geonames_parallel( try: with mp.Pool(processes=num_workers) as pool: - for triples in pool.imap_unordered(process_geonames_lines, batches, chunksize=10): + for triples in pool.imap_unordered( + process_geonames_lines, batches, chunksize=10 + ): if triples: yield triples except Exception as e: @@ -412,7 +421,9 @@ def stream_ntriples_parallel( try: with mp.Pool(processes=num_workers) as pool: - for triples in pool.imap_unordered(process_ntriples_lines, batches, chunksize=10): + for triples in pool.imap_unordered( + process_ntriples_lines, batches, chunksize=10 + ): if triples: yield triples except Exception as e: @@ -420,7 +431,9 @@ def stream_ntriples_parallel( raise -def stream_turtle_chunks(file_path: Path, chunk_size: int = 10000) -> Iterator[list[RDFTriple]]: +def stream_turtle_chunks( + file_path: Path, chunk_size: int = 10000 +) -> Iterator[list[RDFTriple]]: """Stream Turtle file in chunks using incremental parsing. Module-level function for use in picklable generators. @@ -446,7 +459,9 @@ def stream_turtle_chunks(file_path: Path, chunk_size: int = 10000) -> Iterator[l if not stripped or stripped.startswith("#"): continue - if in_prefixes and (stripped.startswith("@prefix") or stripped.startswith("@base")): + if in_prefixes and ( + stripped.startswith("@prefix") or stripped.startswith("@base") + ): prefix_lines.append(line) continue elif in_prefixes: @@ -456,7 +471,10 @@ def stream_turtle_chunks(file_path: Path, chunk_size: int = 10000) -> Iterator[l triple_count += 1 if triple_count >= chunk_size: - chunk_text = "".join(prefix_lines) + "\n" + "".join(current_chunk) + "\n" + line + chunk_text = ( + "".join(prefix_lines) + "\n" + + "".join(current_chunk) + "\n" + line + ) try: graph = Graph() @@ -496,7 +514,9 @@ def stream_turtle_chunks(file_path: Path, chunk_size: int = 10000) -> Iterator[l file_obj.close() -def stream_generic_rdf(file_path: Path, rdf_format: str, chunk_size: int = 10000) -> Iterator[list[RDFTriple]]: +def stream_generic_rdf( + file_path: Path, rdf_format: str, chunk_size: int = 10000 +) -> Iterator[list[RDFTriple]]: """Stream generic RDF file using standard parsing. Module-level function for use in picklable generators. @@ -520,7 +540,9 @@ def stream_generic_rdf(file_path: Path, rdf_format: str, chunk_size: int = 10000 class MetadataHandler: """Handle metadata for the HuggingFace dataset.""" @staticmethod - def add_metadata(dataset_dict: DatasetDict, metadata: dict[str, Any]) -> DatasetDict: + def add_metadata( + dataset_dict: DatasetDict, metadata: dict[str, Any] + ) -> DatasetDict: for split in dataset_dict: if "description" in metadata: dataset_dict[split].info.description = metadata["description"] @@ -577,7 +599,9 @@ class ProgressTracker: console=self.console, transient=False, ) as progress: - task = progress.add_task(f"[yellow]{description}[/yellow]", total=total, **extras) + task = progress.add_task( + f"[yellow]{description}[/yellow]", total=total, **extras + ) yield progress, task # ============================================================================= @@ -609,17 +633,26 @@ class ConversionStrategy(ABC): """Perform conversion given the configuration.""" pass - def _save_dataset(self, dataset_dict: DatasetDict, config: ConversionConfig) -> None: + def _save_dataset( + self, dataset_dict: DatasetDict, config: ConversionConfig + ) -> None: try: if config.push_to_hub: if not config.hub_repo_id: raise ValueError("hub_repo_id is required when push_to_hub is True") - self.progress.print(f"Uploading to HuggingFace Hub: {config.hub_repo_id}...", "yellow") + self.progress.print( + f"Uploading to HuggingFace Hub: {config.hub_repo_id}...", + "yellow" + ) dataset_dict.push_to_hub(config.hub_repo_id, private=False) - self.progress.print(f"✓ Successfully uploaded to {config.hub_repo_id}", "green") + self.progress.print( + f"✓ Successfully uploaded to {config.hub_repo_id}", "green" + ) else: if not config.output_path: - raise ValueError("output_path is required when push_to_hub is False") + raise ValueError( + "output_path is required when push_to_hub is False" + ) config.output_path.mkdir(parents=True, exist_ok=True) dataset_dict.save_to_disk(str(config.output_path)) self.progress.print(f"Dataset saved to {config.output_path}", "green") @@ -633,7 +666,9 @@ class ConversionStrategy(ABC): try: if config.clean_cache: with tempfile.TemporaryDirectory() as temp_cache_dir: - dataset = Dataset.from_parquet(str(chunks_dir / "*.parquet"), cache_dir=temp_cache_dir) + dataset = Dataset.from_parquet( + str(chunks_dir / "*.parquet"), cache_dir=temp_cache_dir + ) else: dataset = Dataset.from_parquet(str(chunks_dir / "*.parquet")) except Exception as e: @@ -644,7 +679,9 @@ class ConversionStrategy(ABC): if create_splits and config.create_train_test_split: train_test = dataset.train_test_split(test_size=config.test_size, seed=42) - return DatasetDict({"train": train_test["train"], "test": train_test["test"]}) + return DatasetDict({ + "train": train_test["train"], "test": train_test["test"] + }) return DatasetDict({"data": dataset}) def _save_dataset_info_json(self, config: ConversionConfig, total_triples: int, @@ -656,7 +693,9 @@ class ConversionStrategy(ABC): "source_format": config.rdf_format, "chunk_size": config.chunk_size, "processing_time_seconds": round(elapsed, 2), - "triples_per_second": round(total_triples / elapsed, 2) if elapsed > 0 else 0, + "triples_per_second": ( + round(total_triples / elapsed, 2) if elapsed > 0 else 0 + ), "conversion_date": time.strftime("%Y-%m-%d %H:%M:%S"), } if "train" in dataset_dict: @@ -700,35 +739,57 @@ class ConversionStrategy(ABC): logger.error(f"Error streaming N-Triples: {e}") raise - def _stream_generic_rdf(self, config: ConversionConfig) -> Iterator[list[RDFTriple]]: + def _stream_generic_rdf( + self, config: ConversionConfig + ) -> Iterator[list[RDFTriple]]: """Stream generic RDF file using standard parsing. Thin wrapper that delegates to module-level function. """ - return stream_generic_rdf(config.input_path, config.rdf_format, config.chunk_size) + return stream_generic_rdf( + config.input_path, config.rdf_format, config.chunk_size + ) - def _stream_geonames_parallel(self, config: ConversionConfig, num_workers: int) -> Iterator[list[RDFTriple]]: + def _stream_geonames_parallel( + self, config: ConversionConfig, num_workers: int + ) -> Iterator[list[RDFTriple]]: """Stream GeoNames RDF file with parallel processing. Thin wrapper that prints progress then delegates to module-level function. """ - self.progress.print(f"[yellow]Using parallel GeoNames parser with {num_workers} workers[/yellow]") - return stream_geonames_parallel(config.input_path, config.chunk_size, num_workers) + self.progress.print( + f"[yellow]Using parallel GeoNames parser with " + f"{num_workers} workers[/yellow]" + ) + return stream_geonames_parallel( + config.input_path, config.chunk_size, num_workers + ) - def _stream_ntriples_parallel(self, config: ConversionConfig, num_workers: int) -> Iterator[list[RDFTriple]]: + def _stream_ntriples_parallel( + self, config: ConversionConfig, num_workers: int + ) -> Iterator[list[RDFTriple]]: """Stream N-Triples file with parallel processing. Thin wrapper that prints progress then delegates to module-level function. """ - self.progress.print(f"[yellow]Using parallel N-Triples parser with {num_workers} workers[/yellow]") - return stream_ntriples_parallel(config.input_path, config.chunk_size, num_workers) + self.progress.print( + f"[yellow]Using parallel N-Triples parser with " + f"{num_workers} workers[/yellow]" + ) + return stream_ntriples_parallel( + config.input_path, config.chunk_size, num_workers + ) - def _stream_turtle_chunks(self, config: ConversionConfig) -> Iterator[list[RDFTriple]]: + def _stream_turtle_chunks( + self, config: ConversionConfig + ) -> Iterator[list[RDFTriple]]: """Stream Turtle file in chunks using incremental parsing. Thin wrapper that prints progress then delegates to module-level function. """ - self.progress.print("[yellow]Streaming Turtle file (line-by-line parser)[/yellow]") + self.progress.print( + "[yellow]Streaming Turtle file (line-by-line parser)[/yellow]" + ) return stream_turtle_chunks(config.input_path, config.chunk_size) # ============================================================================= @@ -751,7 +812,9 @@ class StandardStrategy(ConversionStrategy): self.progress.print(f"Converting RDF file: {config.input_path}", "cyan") self.progress.print(f"File size: {file_size_mb:.2f} MB", "dim") if config.push_to_hub: - self.progress.print(f"Destination: {config.hub_repo_id} (HuggingFace Hub)", "dim") + self.progress.print( + f"Destination: {config.hub_repo_id} (HuggingFace Hub)", "dim" + ) else: self.progress.print(f"Output: {config.output_path}", "dim") self.progress.emit_progress(5) @@ -760,7 +823,9 @@ class StandardStrategy(ConversionStrategy): if config.rdf_format == "tsv": # For TSV, parse line-by-line differently (omitted here for brevity) triples = [] # Implement specialized TSV parser if needed - raise NotImplementedError("TSV parsing not implemented in StandardStrategy.") + raise NotImplementedError( + "TSV parsing not implemented in StandardStrategy." + ) else: triples = self._stream_generic_rdf(config) triples = [triple for chunk in triples for triple in chunk] @@ -774,13 +839,20 @@ class StandardStrategy(ConversionStrategy): dataset = Dataset.from_list(triples, features=features) # pyright: ignore dataset_dict = DatasetDict({"data": dataset}) if config.metadata: - dataset_dict = self.metadata_handler.add_metadata(dataset_dict, config.metadata) + dataset_dict = self.metadata_handler.add_metadata( + dataset_dict, config.metadata + ) self._save_dataset(dataset_dict, config) self.progress.emit_progress(100) elapsed = time.time() - start_time if config.output_path and not config.push_to_hub: - self._save_dataset_info_json(config, len(triples), dataset_dict, elapsed) - self.progress.print(f"✓ Converted {len(triples):,} triples in {elapsed:.1f}s", "bold green") + self._save_dataset_info_json( + config, len(triples), dataset_dict, elapsed + ) + self.progress.print( + f"✓ Converted {len(triples):,} triples in {elapsed:.1f}s", + "bold green" + ) return ConversionResult(success=True, total_triples=len(triples), processing_time_seconds=elapsed, output_path=config.output_path, @@ -809,7 +881,9 @@ class StreamingStrategy(ConversionStrategy): self.progress.print(f"File size: {file_size_mb:.2f} MB", "dim") self.progress.print(f"Chunk size: {config.chunk_size:,} triples", "dim") if config.push_to_hub: - self.progress.print(f"Destination: {config.hub_repo_id} (HuggingFace Hub)", "dim") + self.progress.print( + f"Destination: {config.hub_repo_id} (HuggingFace Hub)", "dim" + ) else: self.progress.print(f"Output: {config.output_path}", "dim") self.progress.emit_progress(5) @@ -825,7 +899,9 @@ class StreamingStrategy(ConversionStrategy): total_triples = 0 chunk_count = 0 - with self.progress.progress_bar("Processing RDF chunks...", total=None) as (progress, task): + with self.progress.progress_bar( + "Processing RDF chunks...", total=None + ) as (progress, task): for chunk in self._stream_generic_rdf(config): schema = self.schema_manager.get_arrow_schema() table = pa.Table.from_pylist(chunk, schema=schema) @@ -834,9 +910,14 @@ class StreamingStrategy(ConversionStrategy): total_triples += len(chunk) chunk_count += 1 progress.update(task, advance=1) - self.progress.print(f"✓ Processed {total_triples:,} triples in {chunk_count} chunks", "green") + self.progress.print( + f"✓ Processed {total_triples:,} triples in {chunk_count} chunks", + "green" + ) self.progress.emit_progress(80) - dataset_dict = self._merge_parquet_chunks(temp_chunks_dir, config, create_splits=True) + dataset_dict = self._merge_parquet_chunks( + temp_chunks_dir, config, create_splits=True + ) shutil.rmtree(temp_chunks_dir) if config.push_to_hub: @@ -845,7 +926,9 @@ class StreamingStrategy(ConversionStrategy): self.progress.emit_progress(100) elapsed = time.time() - start_time if config.output_path and not config.push_to_hub: - self._save_dataset_info_json(config, total_triples, dataset_dict, elapsed) + self._save_dataset_info_json( + config, total_triples, dataset_dict, elapsed + ) self.progress.print("✓ Converted to HuggingFace dataset", "green") return ConversionResult(success=True, total_triples=total_triples, processing_time_seconds=elapsed, @@ -875,7 +958,9 @@ class StreamingTurtleStrategy(ConversionStrategy): self.progress.print(f"File size: {file_size_mb:.2f} MB", "dim") self.progress.print(f"Chunk size: {config.chunk_size:,} statements", "dim") if config.push_to_hub: - self.progress.print(f"Destination: {config.hub_repo_id} (HuggingFace Hub)", "dim") + self.progress.print( + f"Destination: {config.hub_repo_id} (HuggingFace Hub)", "dim" + ) else: self.progress.print(f"Output: {config.output_path}", "dim") self.progress.emit_progress(5) @@ -905,7 +990,9 @@ class StreamingTurtleStrategy(ConversionStrategy): current_chunk = [] if current_chunk: yield current_chunk - with self.progress.progress_bar("Processing Turtle chunks...", total=None) as (progress, task): + with self.progress.progress_bar( + "Processing Turtle chunks...", total=None + ) as (progress, task): for chunk in stream_turtle_chunks(): schema = self.schema_manager.get_arrow_schema() table = pa.Table.from_pylist(chunk, schema=schema) @@ -914,7 +1001,10 @@ class StreamingTurtleStrategy(ConversionStrategy): total_triples += len(chunk) chunk_count += 1 progress.update(task, advance=1) - self.progress.print(f"✓ Processed {total_triples:,} triples in {chunk_count} chunks", "green") + self.progress.print( + f"✓ Processed {total_triples:,} triples in {chunk_count} chunks", + "green" + ) self.progress.emit_progress(80) dataset_dict = self._merge_parquet_chunks(temp_chunks_dir, config) shutil.rmtree(temp_chunks_dir) @@ -925,7 +1015,9 @@ class StreamingTurtleStrategy(ConversionStrategy): self.progress.emit_progress(100) elapsed = time.time() - start_time if config.output_path and not config.push_to_hub: - self._save_dataset_info_json(config, total_triples, dataset_dict, elapsed) + self._save_dataset_info_json( + config, total_triples, dataset_dict, elapsed + ) self.progress.print("✓ Converted to HuggingFace dataset", "green") return ConversionResult(success=True, total_triples=total_triples, processing_time_seconds=elapsed, @@ -955,7 +1047,9 @@ class SimpleStreamingStrategy(ConversionStrategy): self.progress.print(f"File size: {file_size_mb:.2f} MB", "dim") self.progress.print(f"Chunk size: {config.chunk_size:,} triples", "dim") if config.push_to_hub: - self.progress.print(f"Destination: {config.hub_repo_id} (HuggingFace Hub)", "dim") + self.progress.print( + f"Destination: {config.hub_repo_id} (HuggingFace Hub)", "dim" + ) else: self.progress.print(f"Output: {config.output_path}", "dim") @@ -971,7 +1065,9 @@ class SimpleStreamingStrategy(ConversionStrategy): total_triples = 0 chunk_count = 0 - with self.progress.progress_bar("Processing simple streaming chunks...", total=None) as (progress, task): + with self.progress.progress_bar( + "Processing simple streaming chunks...", total=None + ) as (progress, task): for chunk in self._stream_ntriples(config): schema = self.schema_manager.get_arrow_schema() table = pa.Table.from_pylist(chunk, schema=schema) @@ -980,7 +1076,10 @@ class SimpleStreamingStrategy(ConversionStrategy): total_triples += len(chunk) chunk_count += 1 progress.update(task, advance=1) - self.progress.print(f"✓ Processed {total_triples:,} triples in {chunk_count} chunks", "green") + self.progress.print( + f"✓ Processed {total_triples:,} triples in {chunk_count} chunks", + "green" + ) dataset_dict = self._merge_parquet_chunks(temp_chunks_dir, config) shutil.rmtree(temp_chunks_dir) @@ -989,7 +1088,9 @@ class SimpleStreamingStrategy(ConversionStrategy): elapsed = time.time() - start_time if config.output_path and not config.push_to_hub: - self._save_dataset_info_json(config, total_triples, dataset_dict, elapsed) + self._save_dataset_info_json( + config, total_triples, dataset_dict, elapsed + ) self.progress.print("✓ Converted to HuggingFace dataset", "green") return ConversionResult(success=True, total_triples=total_triples, processing_time_seconds=elapsed, @@ -1009,7 +1110,10 @@ class ParallelStreamingStrategy(ConversionStrategy): @property def description(self) -> str: - return "Parallel streaming conversion using multiprocessing with Dataset.from_generator()." + return ( + "Parallel streaming conversion using multiprocessing with " + "Dataset.from_generator()." + ) def convert(self, config: ConversionConfig) -> ConversionResult: start_time = time.time() @@ -1018,10 +1122,14 @@ class ParallelStreamingStrategy(ConversionStrategy): file_size_mb = self.file_handler.get_file_size_mb(config.input_path) self.progress.print(f"Converting: {config.input_path.name}", "cyan") self.progress.print(f"File size: {file_size_mb:.2f} MB", "dim") - self.progress.print(f"Chunk size: {config.chunk_size:,} lines per batch", "dim") + self.progress.print( + f"Chunk size: {config.chunk_size:,} lines per batch", "dim" + ) self.progress.print(f"Workers: {num_workers} CPU cores", "dim") if config.push_to_hub: - self.progress.print(f"Destination: {config.hub_repo_id} (HuggingFace Hub)", "dim") + self.progress.print( + f"Destination: {config.hub_repo_id} (HuggingFace Hub)", "dim" + ) else: self.progress.print(f"Output: {config.output_path}", "dim") self.progress.emit_progress(5) @@ -1029,11 +1137,15 @@ class ParallelStreamingStrategy(ConversionStrategy): # Check if this is GeoNames format is_geonames = "geonames" in str(config.input_path).lower() - self.progress.print("\n[yellow]Creating dataset from streamed chunks...[/yellow]") + self.progress.print( + "\n[yellow]Creating dataset from streamed chunks...[/yellow]" + ) self.progress.emit_progress(10) # Estimate total chunks based on file size (rough estimate) - estimated_chunks = max(10, int(file_size_mb * 1024 * 1024 / (config.chunk_size * 100))) + estimated_chunks = max( + 10, int(file_size_mb * 1024 * 1024 / (config.chunk_size * 100)) + ) # Determine format and print status messages BEFORE creating generator # (Cannot print from inside generator due to pickling constraints) @@ -1044,42 +1156,63 @@ class ParallelStreamingStrategy(ConversionStrategy): first_line = f.readline().strip() if first_line.startswith("http://") or first_line.startswith("https://"): format_type = "geonames" - self.progress.print(f"[yellow]Using parallel GeoNames parser with {num_workers} workers[/yellow]") + self.progress.print( + f"[yellow]Using parallel GeoNames parser with " + f"{num_workers} workers[/yellow]" + ) else: format_type = "generic" - self.progress.print(f"[yellow]Using standard RDF parser for {config.rdf_format} (single-threaded)[/yellow]") + self.progress.print( + f"[yellow]Using standard RDF parser for " + f"{config.rdf_format} (single-threaded)[/yellow]" + ) elif config.rdf_format in ("nt", "ntriples"): format_type = "ntriples" - self.progress.print(f"[yellow]Using parallel N-Triples parser with {num_workers} workers[/yellow]") + self.progress.print( + f"[yellow]Using parallel N-Triples parser with " + f"{num_workers} workers[/yellow]" + ) elif config.rdf_format in ("turtle", "ttl"): format_type = "turtle" - self.progress.print("[yellow]Streaming Turtle file (line-by-line parser)[/yellow]") + self.progress.print( + "[yellow]Streaming Turtle file (line-by-line parser)[/yellow]" + ) else: format_type = "generic" - self.progress.print(f"[yellow]Using standard RDF parser for {config.rdf_format} (single-threaded)[/yellow]") + self.progress.print( + f"[yellow]Using standard RDF parser for " + f"{config.rdf_format} (single-threaded)[/yellow]" + ) # Note: Cannot capture self or generators in dataset_generator closure. # Dataset.from_generator() requires picklable generators. # Must call streaming functions directly inside the generator. def dataset_generator(): nonlocal triple_count - chunk_count = 0 last_print_count = 0 - # Create the chunk iterator inside the generator to avoid pickling issues + # Create the chunk iterator inside the generator to avoid + # pickling issues if format_type == "geonames": - chunk_source = stream_geonames_parallel(config.input_path, config.chunk_size, num_workers) + chunk_source = stream_geonames_parallel( + config.input_path, config.chunk_size, num_workers + ) elif format_type == "ntriples": - chunk_source = stream_ntriples_parallel(config.input_path, config.chunk_size, num_workers) + chunk_source = stream_ntriples_parallel( + config.input_path, config.chunk_size, num_workers + ) elif format_type == "turtle": - chunk_source = stream_turtle_chunks(config.input_path, config.chunk_size) + chunk_source = stream_turtle_chunks( + config.input_path, config.chunk_size + ) else: # Generic RDF parsing - chunk_source = stream_generic_rdf(config.input_path, config.rdf_format, config.chunk_size) + chunk_source = stream_generic_rdf( + config.input_path, config.rdf_format, config.chunk_size + ) # Yield individual triples from chunks - for chunk in chunk_source: - chunk_count += 1 + for chunk_count, chunk in enumerate(chunk_source): for triple in chunk: triple_count += 1 yield triple @@ -1087,7 +1220,9 @@ class ParallelStreamingStrategy(ConversionStrategy): if triple_count - last_print_count >= 100000: elapsed = time.time() - start_time rate = triple_count / elapsed if elapsed > 0 else 0 - progress_pct = 10 + min(60, int((chunk_count / estimated_chunks) * 60)) + progress_pct = 10 + min( + 60, int((chunk_count / estimated_chunks) * 60) + ) print( f" Processing: {chunk_count:,} chunks • " f"{triple_count:,} triples • {rate:.0f} triples/sec", @@ -1115,15 +1250,23 @@ class ParallelStreamingStrategy(ConversionStrategy): # Add metadata if config.metadata: - dataset_dict = self.metadata_handler.add_metadata(dataset_dict, config.metadata) + dataset_dict = self.metadata_handler.add_metadata( + dataset_dict, config.metadata + ) elapsed = time.time() - start_time if config.push_to_hub: - self.progress.print(f"[yellow]Uploading to HuggingFace Hub: {config.hub_repo_id}...[/yellow]") + self.progress.print( + f"[yellow]Uploading to HuggingFace Hub: " + f"{config.hub_repo_id}...[/yellow]" + ) self.progress.emit_progress(95) dataset_dict.push_to_hub(config.hub_repo_id, private=False) - self.progress.print(f"[bold green]✓ Successfully uploaded to {config.hub_repo_id}[/bold green]") + self.progress.print( + f"[bold green]✓ Successfully uploaded to " + f"{config.hub_repo_id}[/bold green]" + ) self.progress.emit_progress(100) else: self.progress.print("[yellow]Saving final dataset...[/yellow]") @@ -1144,18 +1287,31 @@ class ParallelStreamingStrategy(ConversionStrategy): "chunk_size": config.chunk_size, "num_workers": num_workers, "processing_time_seconds": round(elapsed, 2), - "triples_per_second": round(total_triples / elapsed, 2) if elapsed > 0 else 0, + "triples_per_second": ( + round(total_triples / elapsed, 2) if elapsed > 0 else 0 + ), "conversion_date": time.strftime("%Y-%m-%d %H:%M:%S"), } with open(config.output_path / "dataset_info.json", "w") as f: json.dump(info, f, indent=2) - self.progress.print("\n[bold green]✓ Successfully converted to HuggingFace dataset[/bold green]") - self.progress.print(f"[green] Data split: {len(dataset_dict['data']):,} triples[/green]") - self.progress.print(f"[green] Processing time: {elapsed:.1f} seconds[/green]") - self.progress.print(f"[green] Speed: {total_triples / elapsed:.0f} triples/second[/green]") + self.progress.print( + "\n[bold green]✓ Successfully converted to " + "HuggingFace dataset[/bold green]" + ) + self.progress.print( + f"[green] Data split: {len(dataset_dict['data']):,} triples[/green]" + ) + self.progress.print( + f"[green] Processing time: {elapsed:.1f} seconds[/green]" + ) + self.progress.print( + f"[green] Speed: {total_triples / elapsed:.0f} triples/second[/green]" + ) if config.push_to_hub: - self.progress.print(f"[green] Repository: {config.hub_repo_id}[/green]") + self.progress.print( + f"[green] Repository: {config.hub_repo_id}[/green]" + ) else: self.progress.print(f"[green] Location: {config.output_path}[/green]") @@ -1185,31 +1341,44 @@ class StrategySelector: def __init__(self, progress_tracker: ProgressTracker): self.progress = progress_tracker - self._strategies = {name: cls(progress_tracker) for name, cls in self.STRATEGIES.items()} + self._strategies = { + name: cls(progress_tracker) for name, cls in self.STRATEGIES.items() + } def select_auto(self, file_path: Path, rdf_format: str) -> ConversionStrategy: file_size_mb = FileHandler.get_file_size_mb(file_path) cpu_count = mp.cpu_count() is_geonames = "geonames" in str(file_path).lower() if is_geonames: - self.progress.print("Auto-selected: streaming-parallel (GeoNames detected)", "cyan") + self.progress.print( + "Auto-selected: streaming-parallel (GeoNames detected)", "cyan" + ) return self._strategies["streaming-parallel"] if rdf_format in ("turtle", "ttl") and file_size_mb > 100: - self.progress.print("Auto-selected: streaming-turtle (large Turtle file)", "cyan") + self.progress.print( + "Auto-selected: streaming-turtle (large Turtle file)", "cyan" + ) return self._strategies["streaming-turtle"] if file_size_mb < 100: self.progress.print("Auto-selected: standard (file < 100MB)", "cyan") return self._strategies["standard"] if file_size_mb > 1000 and cpu_count > 1: - self.progress.print(f"Auto-selected: streaming-parallel (> 1GB, {cpu_count} cores)", "cyan") + self.progress.print( + f"Auto-selected: streaming-parallel (> 1GB, {cpu_count} cores)", + "cyan" + ) return self._strategies["streaming-parallel"] - self.progress.print(f"Auto-selected: streaming (file {file_size_mb:.0f}MB)", "cyan") + self.progress.print( + f"Auto-selected: streaming (file {file_size_mb:.0f}MB)", "cyan" + ) return self._strategies["streaming"] def select_manual(self, strategy_name: str) -> ConversionStrategy: if strategy_name not in self._strategies: valid = ", ".join(self._strategies.keys()) - raise ValueError(f"Unknown strategy '{strategy_name}'. Valid strategies: {valid}") + raise ValueError( + f"Unknown strategy '{strategy_name}'. Valid strategies: {valid}" + ) return self._strategies[strategy_name] def list_strategies(self) -> list[tuple[str, str]]: @@ -1234,18 +1403,25 @@ Strategies: """ ) parser.add_argument("input", type=Path, help="Input RDF file") - parser.add_argument("output", nargs="?", type=Path, help="Output directory for HuggingFace dataset (optional if --push-to-hub)") - parser.add_argument("-f", "--format", default="turtle", - choices=["turtle", "nt", "ntriples", "xml", "n3", "trig", "nquads", "tsv"], - help="RDF format (default: turtle)") + parser.add_argument( + "output", nargs="?", type=Path, + help="Output directory for HuggingFace dataset (optional if --push-to-hub)" + ) + parser.add_argument( + "-f", "--format", default="turtle", + choices=["turtle", "nt", "ntriples", "xml", "n3", "trig", "nquads", "tsv"], + help="RDF format (default: turtle)" + ) parser.add_argument("--strategy", default="auto", choices=["auto", "standard", "streaming", "streaming-turtle", "streaming-simple", "streaming-parallel"], help="Conversion strategy (default: auto)") parser.add_argument("--chunk-size", type=int, default=10000, help="Number of triples/lines per chunk (default: 10000)") - parser.add_argument("--num-workers", type=int, default=None, - help="Number of workers for parallel strategies (default: CPU count)") + parser.add_argument( + "--num-workers", type=int, default=None, + help="Number of workers for parallel strategies (default: CPU count)" + ) parser.add_argument("--create-splits", action="store_true", help="Create train/test splits (95%%/5%%)") parser.add_argument("--test-size", type=float, default=0.05, @@ -1256,10 +1432,14 @@ Strategies: parser.add_argument("--citation", type=str, help="Dataset citation") parser.add_argument("--homepage", type=str, help="Dataset homepage URL") parser.add_argument("--license", type=str, help="Dataset license") - parser.add_argument("--push-to-hub", action="store_true", - help="Upload directly to HuggingFace Hub without saving to disk") - parser.add_argument("--hub-repo-id", type=str, - help="HuggingFace Hub repository ID (required if --push-to-hub)") + parser.add_argument( + "--push-to-hub", action="store_true", + help="Upload directly to HuggingFace Hub without saving to disk" + ) + parser.add_argument( + "--hub-repo-id", type=str, + help="HuggingFace Hub repository ID (required if --push-to-hub)" + ) parser.add_argument("-v", "--verbose", action="store_true", help="Verbose output including warnings") args = parser.parse_args() @@ -1271,7 +1451,9 @@ Strategies: parser.error("output directory is required when not using --push-to-hub") log_level = logging.WARNING if args.verbose else logging.ERROR - logging.basicConfig(level=log_level, format="%(levelname)s:%(name)s:%(message)s", force=True) + logging.basicConfig( + level=log_level, format="%(levelname)s:%(name)s:%(message)s", force=True + ) logging.getLogger("rdflib").setLevel(log_level) logging.getLogger("datasets").setLevel(log_level) @@ -1331,7 +1513,9 @@ Strategies: progress.print("Conversion Complete!", "bold green") progress.print("="*60, "green") progress.print(f"Total triples: {result.total_triples:,}", "green") - progress.print(f"Processing time: {result.processing_time_seconds:.1f}s", "green") + progress.print( + f"Processing time: {result.processing_time_seconds:.1f}s", "green" + ) if config.push_to_hub: progress.print(f"Repository: {config.hub_repo_id}", "green") else: diff --git a/scripts/upload_all_datasets.py b/scripts/upload_all_datasets.py index 5fd481d..8182dc4 100755 --- a/scripts/upload_all_datasets.py +++ b/scripts/upload_all_datasets.py @@ -1018,14 +1018,16 @@ def convert(args: argparse.Namespace, rdf_file: Path, dataset_id: str) -> None: / "convert_rdf_to_hf_dataset_unified.py" ) - # Unified script will auto-select best strategy based on file size and format + # Unified script will auto-select best strategy based on file size + # and format console.print( "[dim]Using unified converter (auto-selects best strategy)[/dim]" ) if args.dry_run: console.print( - f"[dim]Would run: python scripts/convert_rdf_to_hf_dataset_unified.py " + "[dim]Would run: python " + "scripts/convert_rdf_to_hf_dataset_unified.py " "{rdf_file} {hf_dataset_dir} --strategy auto[/dim]" ) else: @@ -1205,7 +1207,8 @@ def main() -> int: args, parser = parse_args() # Fix for "AF_UNIX path too long" error in multiprocessing - # This forces the temporary directory to be /tmp (short path) instead of a potentially deep workspace path + # This forces the temporary directory to be /tmp (short path) instead of + # a potentially deep workspace path os.environ["TMPDIR"] = "/tmp" start_time = time.monotonic() @@ -1260,7 +1263,9 @@ def main() -> int: return 1 one_dataset_download_time = time.monotonic() - one_dataset_download_duration = one_dataset_download_time - one_dataset_start_time + one_dataset_download_duration = ( + one_dataset_download_time - one_dataset_start_time + ) print(f"`download` took {one_dataset_download_duration:.2f} seconds:") # Decompress @@ -1276,7 +1281,9 @@ def main() -> int: return 1 one_dataset_decompress_time = time.monotonic() - one_dataset_decompress_duration = one_dataset_decompress_time - one_dataset_download_time + one_dataset_decompress_duration = ( + one_dataset_decompress_time - one_dataset_download_time + ) print(f"`decompress` took {one_dataset_decompress_duration:.2f} seconds:") # Convert @@ -1303,7 +1310,9 @@ def main() -> int: return 1 one_dataset_convert_time = time.monotonic() - one_dataset_convert_duration = one_dataset_convert_time - one_dataset_decompress_time + one_dataset_convert_duration = ( + one_dataset_convert_time - one_dataset_decompress_time + ) print(f"`convert` took {one_dataset_convert_duration:.2f} seconds:") # Upload -- 2.52.0 From fbaa30e9047a6d053063b3b15c1ef81be9af0b97 Mon Sep 17 00:00:00 2001 From: Brent Edwards Date: Mon, 29 Dec 2025 14:19:34 -0800 Subject: [PATCH 6/8] Adding shell.nix. --- shell.nix | 18 ++++++++++++++++++ 1 file changed, 18 insertions(+) create mode 100644 shell.nix diff --git a/shell.nix b/shell.nix new file mode 100644 index 0000000..5a9fa43 --- /dev/null +++ b/shell.nix @@ -0,0 +1,18 @@ +let + pkgs = import {}; + lib-path = with pkgs; + lib.makeLibraryPath [ + stdenv.cc.cc + ]; +in + pkgs.mkShell { + packages = with pkgs; [ + python312 + (poetry.override {python3 = python312;}) + uv + ]; + + shellHook = '' + export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:${lib-path} + ''; + } -- 2.52.0 From da9f78532658b5d2d7edb2fbcc62e8b0f684e12c Mon Sep 17 00:00:00 2001 From: Brent Edwards Date: Mon, 29 Dec 2025 18:09:54 -0800 Subject: [PATCH 7/8] Updating `pyproject` to require Python 3.10. --- pyproject.toml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pyproject.toml b/pyproject.toml index 545a8f8..4eecc8f 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -7,7 +7,7 @@ name = "app-boilerplate" version = "0.1.0" description = "Replace with a short summary of your project." readme = "README.md" -requires-python = ">=3.9" +requires-python = ">=3.10" license = {text = "MIT"} authors = [ {name = "Your Name", email = "your.email@example.com"}, -- 2.52.0 From f8a7739016d5464f9344d69b989dfd00c9c31131 Mon Sep 17 00:00:00 2001 From: Brent Edwards Date: Tue, 30 Dec 2025 18:06:34 -0800 Subject: [PATCH 8/8] Minor changes. --- features/steps/rdf_converter_steps.py | 21 +++++++++---------- ...rt_rdf_to_hf_dataset_streaming_parallel.py | 21 +++++++------------ scripts/convert_rdf_to_hf_dataset_unified.py | 1 - scripts/dataset_downloader.py | 2 +- scripts/rdf_to_hf_incremental.py | 21 +++++++++---------- scripts/upload_all_datasets.py | 2 ++ src/streaming/extract_stream.py | 2 +- src/streaming/http_stream.py | 2 +- 8 files changed, 33 insertions(+), 39 deletions(-) diff --git a/features/steps/rdf_converter_steps.py b/features/steps/rdf_converter_steps.py index fccc02f..e889d29 100644 --- a/features/steps/rdf_converter_steps.py +++ b/features/steps/rdf_converter_steps.py @@ -14,23 +14,21 @@ import bz2 import gzip import sys from pathlib import Path -from typing import TYPE_CHECKING, Any +from typing import TYPE_CHECKING from behave import given, then, when +from scripts.convert_rdf_to_hf_dataset_unified import ( + ConversionConfig, + ConversionResult, + ProgressTracker, + StrategySelector, +) + # Add scripts directory to path for imports SCRIPTS_DIR = Path(__file__).parent.parent.parent / "scripts" sys.path.insert(0, str(SCRIPTS_DIR)) -from convert_rdf_to_hf_dataset_unified import ( - ConversionConfig, - ConversionResult, - FileHandler, - ProgressTracker, - SchemaManager, - StrategySelector, -) - if TYPE_CHECKING: from behave.runner import Context @@ -726,7 +724,8 @@ def step_verify_strategies_include(context: Context, strategy_names: str) -> Non def _run_cli(args: list[str]) -> tuple[int, str, str]: """Run the CLI by calling main() directly for coverage tracking.""" import io - from contextlib import redirect_stdout, redirect_stderr + from contextlib import redirect_stderr, redirect_stdout + from convert_rdf_to_hf_dataset_unified import main old_argv = sys.argv diff --git a/scripts/convert_rdf_to_hf_dataset_streaming_parallel.py b/scripts/convert_rdf_to_hf_dataset_streaming_parallel.py index 4a488f5..d0d8aee 100755 --- a/scripts/convert_rdf_to_hf_dataset_streaming_parallel.py +++ b/scripts/convert_rdf_to_hf_dataset_streaming_parallel.py @@ -11,7 +11,6 @@ import gzip import json import logging import multiprocessing as mp -import shutil import sys import time from collections.abc import Iterator @@ -19,17 +18,9 @@ from pathlib import Path from typing import Any import pyarrow as pa -import pyarrow.parquet as pq -from datasets import Dataset, DatasetDict, Features, Value +from datasets import Dataset, DatasetDict, Features, IterableDataset, Value from rdflib import Graph, Literal, URIRef from rich.console import Console -from rich.progress import ( - BarColumn, - Progress, - SpinnerColumn, - TextColumn, - TimeElapsedColumn, -) logger = logging.getLogger(__name__) logger.addHandler(logging.StreamHandler(sys.stdout)) @@ -655,7 +646,7 @@ def convert_rdf_to_hf_streaming( # unpicklable objects. Dataset.from_generator() requires picklable # generators, so using simple print statements def dataset_generator(): - triple_count = 0 + nonlocal triple_count chunk_count = 0 last_print_count = 0 @@ -688,16 +679,20 @@ def convert_rdf_to_hf_streaming( 'object_language': Value('string'), }) + triple_count = 0 dataset = Dataset.from_generator(dataset_generator, features=features) print("\nPROGRESS: 75", flush=True) console.print("[green]✓ Dataset generation complete[/green]") # Get statistics after generation - total_triples = len(dataset) + if not isinstance(dataset, IterableDataset): + total_triples = len(dataset) + else: + total_triples = triple_count # Wrap in DatasetDict - dataset_dict = DatasetDict({"data": dataset}) + dataset_dict = DatasetDict({"data": dataset}) # pyright: ignore # Add metadata if metadata: diff --git a/scripts/convert_rdf_to_hf_dataset_unified.py b/scripts/convert_rdf_to_hf_dataset_unified.py index f32b6ff..3dabad6 100644 --- a/scripts/convert_rdf_to_hf_dataset_unified.py +++ b/scripts/convert_rdf_to_hf_dataset_unified.py @@ -49,7 +49,6 @@ from abc import ABC, abstractmethod from collections.abc import Iterator from contextlib import contextmanager from dataclasses import dataclass -from io import TextIOWrapper from pathlib import Path from typing import IO, Any, TextIO, TypedDict, cast diff --git a/scripts/dataset_downloader.py b/scripts/dataset_downloader.py index eeb60db..b472823 100644 --- a/scripts/dataset_downloader.py +++ b/scripts/dataset_downloader.py @@ -139,7 +139,7 @@ class DatasetDownloader: ) # type: ignore[arg-type] else: self.console.print( - f"\n[green]✓ Downloaded dataset" + "\n[green]✓ Downloaded dataset" ) self.console.print(f"\n[cyan]Dataset cached at: {self.output_dir}[/cyan]") diff --git a/scripts/rdf_to_hf_incremental.py b/scripts/rdf_to_hf_incremental.py index 1e18e09..54e9026 100644 --- a/scripts/rdf_to_hf_incremental.py +++ b/scripts/rdf_to_hf_incremental.py @@ -39,9 +39,8 @@ import random import sys import time import zipfile +from collections.abc import Iterator from pathlib import Path -from typing import Iterator -from urllib.parse import urlparse import pyarrow as pa import pyarrow.parquet as pq @@ -55,6 +54,8 @@ PROJECT_ROOT = Path(__file__).resolve().parents[1] if str(PROJECT_ROOT) not in sys.path: sys.path.insert(0, str(PROJECT_ROOT)) +from convert_rdf_to_hf_dataset_streaming_parallel import stream_rdf_chunks +from dataset_registry import DATASET_REGISTRY, DatasetInfo, get_dataset_config from src.streaming.extract_stream import ( ExtractConfig, _infer_name_from_url, @@ -64,8 +65,6 @@ from src.streaming.extract_stream import ( spool_zip_to_tempfile, ) from src.streaming.http_stream import HTTPStreamConfig, stream_http_bytes -from dataset_registry import DATASET_REGISTRY, DatasetInfo, get_dataset_config -from convert_rdf_to_hf_dataset_streaming_parallel import stream_rdf_chunks # Setup logger logger = logging.getLogger(__name__) @@ -393,7 +392,7 @@ def create_rdf_chunk_iterator( if input_type == "file": # Local file mode: use optimized file-based streaming with multiprocessing - console.print(f"[cyan]Mode: Local file streaming[/cyan]") + console.print("[cyan]Mode: Local file streaming[/cyan]") console.print(f"[dim]Using parallel file-based parser ({num_workers} workers)[/dim]") file_path = Path(input_value) chunk_iter = stream_rdf_chunks(file_path, rdf_format, DEFAULT_BATCH_SIZE, num_workers) @@ -401,8 +400,8 @@ def create_rdf_chunk_iterator( elif input_type == "registry": # Streaming download mode: use line-based streaming - console.print(f"[cyan]Mode: Streaming download from registry[/cyan]") - console.print(f"[dim]Using line-based streaming parser (no multiprocessing)[/dim]") + console.print("[cyan]Mode: Streaming download from registry[/cyan]") + console.print("[dim]Using line-based streaming parser (no multiprocessing)[/dim]") ds = DATASET_REGISTRY[input_value] url = ds.url @@ -665,7 +664,7 @@ def list_datasets() -> None: by_category[info.category] = by_category.get(info.category, 0) + 1 console.print(f"\n[bold]Total:[/bold] {total} datasets ({available_count} available)") - console.print(f"[bold]By category:[/bold] ", end="") + console.print("[bold]By category:[/bold] ", end="") console.print(", ".join(f"{cat}: {count}" for cat, count in sorted(by_category.items()))) console.print() @@ -1021,7 +1020,7 @@ Features: if input_type == "registry": ds = DATASET_REGISTRY[input_value] - console.print(f"\n[bold]Dataset Info:[/bold]") + console.print("\n[bold]Dataset Info:[/bold]") console.print(f" Name: {ds.name}") console.print(f" Category: {ds.category}") console.print(f" Size: {ds.size_gb} GB") @@ -1036,12 +1035,12 @@ Features: else: console.print(f" Estimated triples: {ds.triples}") - console.print(f"\n[bold]Would perform:[/bold]") + console.print("\n[bold]Would perform:[/bold]") console.print(f" 1. {'Stream from ' + DATASET_REGISTRY[input_value].url if input_type == 'registry' else 'Read from ' + input_value}") console.print(f" 2. Parse RDF triples in {rdf_format} format") console.print(f" 3. Convert to Parquet shards ({args.rows_per_shard:,} rows each)") console.print(f" 4. Upload shards to {args.repo_id}") - console.print(f" 5. Create and upload README.md") + console.print(" 5. Create and upload README.md") console.print("\n[bold green]✓ Dry run complete. Use without --dry-run to execute.[/bold green]") return 0 diff --git a/scripts/upload_all_datasets.py b/scripts/upload_all_datasets.py index 8182dc4..87d429b 100755 --- a/scripts/upload_all_datasets.py +++ b/scripts/upload_all_datasets.py @@ -1287,6 +1287,8 @@ def main() -> int: print(f"`decompress` took {one_dataset_decompress_duration:.2f} seconds:") # Convert + assert rdf_file is not None + if not args.dry_run: rdf_file = get_most_recent_file(rdf_file.parent) if rdf_file is None: diff --git a/src/streaming/extract_stream.py b/src/streaming/extract_stream.py index ea1cee1..29bb8cf 100644 --- a/src/streaming/extract_stream.py +++ b/src/streaming/extract_stream.py @@ -6,9 +6,9 @@ import io import lzma import tarfile import tempfile +from collections.abc import Iterable, Iterator from dataclasses import dataclass from pathlib import Path -from typing import Iterable, Iterator from urllib.parse import urlparse diff --git a/src/streaming/http_stream.py b/src/streaming/http_stream.py index 24d3598..72bf5cb 100644 --- a/src/streaming/http_stream.py +++ b/src/streaming/http_stream.py @@ -1,8 +1,8 @@ from __future__ import annotations import time +from collections.abc import Iterator from dataclasses import dataclass -from typing import Iterator import httpx -- 2.52.0