diff --git a/scripts/convert_rdf_to_hf_dataset_streaming_parallel.py b/scripts/convert_rdf_to_hf_dataset_streaming_parallel.py index d0d8aee..4462702 100755 --- a/scripts/convert_rdf_to_hf_dataset_streaming_parallel.py +++ b/scripts/convert_rdf_to_hf_dataset_streaming_parallel.py @@ -1,7 +1,9 @@ #!/usr/bin/env python3 -"""Parallel streaming converter for RDF to HuggingFace dataset format with minimal memory usage. +"""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. +This version uses multiprocessing.Pool.imap for simpler and more reliable +parallelization. """ from __future__ import annotations @@ -180,7 +182,9 @@ 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}. Skipping line. Error: {e}" + ) continue return triples @@ -202,7 +206,7 @@ def batch_file_lines(file_path: Path, batch_size: int = 1000, format: str = "geo # Handle compressed files if file_path.suffix == ".gz": try: - file_obj = gzip.open(file_path, "rt", encoding="utf-8") + file_obj = gzip.open(file_path, "rt", encoding="utf-8") # noqa SIM115 # Test read to check if file is valid test_line = file_obj.readline() if not test_line and file_path.stat().st_size > 0: @@ -212,12 +216,14 @@ def batch_file_lines(file_path: Path, batch_size: int = 1000, format: str = "geo from rich.console import Console console = Console() - console.print(f"[red]Error: Compressed file is corrupted or incomplete: {e}[/red]") + console.print("[red]Error: Compressed file is corrupted or incomplete: " + f"{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") + file_obj = open(file_path, encoding="utf-8") # noqa SIM115 try: current_batch = [] @@ -284,7 +290,8 @@ def stream_geonames_parallel( num_workers = max(1, mp.cpu_count() - 1) console = Console() - console.print(f"[yellow]Using parallel GeoNames parser with {num_workers} workers[/yellow]") + console.print("[yellow]Using parallel GeoNames parser with " + f"{num_workers} workers[/yellow]") # Create batches of lines batches = batch_file_lines(file_path, batch_size=chunk_size, format="geonames") @@ -293,7 +300,11 @@ def stream_geonames_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): + for triples in pool.imap_unordered( + process_geonames_lines, + batches, + chunksize=1 + ): if triples: yield triples except Exception as e: @@ -318,7 +329,8 @@ def stream_ntriples_parallel( num_workers = max(1, mp.cpu_count() - 1) console = Console() - console.print(f"[yellow]Using parallel N-Triples parser with {num_workers} workers[/yellow]") + console.print("[yellow]Using parallel N-Triples parser with " + f"{num_workers} workers[/yellow]") # Create batches of lines batches = batch_file_lines(file_path, batch_size=chunk_size, format="ntriples") @@ -326,7 +338,11 @@ def stream_ntriples_parallel( # 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): + for triples in pool.imap_unordered( + process_ntriples_lines, + batches, + chunksize=1 + ): if triples: yield triples except Exception as e: @@ -334,7 +350,9 @@ def stream_ntriples_parallel( raise e -def stream_turtle_chunks(file_path: Path, chunk_size: int = 10000) -> Iterator[list[dict[str, str]]]: +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 @@ -357,13 +375,13 @@ def stream_turtle_chunks(file_path: Path, chunk_size: int = 10000) -> Iterator[l if is_bz2: import bz2 - file_obj = bz2.open(file_path, "rt", encoding="utf-8", errors="ignore") + file_obj = bz2.open(file_path, "rt", encoding="utf-8", errors="ignore") # noqa: SIM115 elif is_gzipped: import gzip - file_obj = gzip.open(file_path, "rt", encoding="utf-8", errors="ignore") + file_obj = gzip.open(file_path, "rt", encoding="utf-8", errors="ignore") # noqa: SIM115 else: - file_obj = open(file_path, encoding="utf-8", errors="ignore") + file_obj = open(file_path, encoding="utf-8", errors="ignore") # noqa: SIM115 try: current_chunk = [] @@ -383,7 +401,8 @@ def stream_turtle_chunks(file_path: Path, chunk_size: int = 10000) -> Iterator[l continue # Collect prefix declarations - 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: @@ -398,7 +417,11 @@ def stream_turtle_chunks(file_path: Path, chunk_size: int = 10000) -> Iterator[l if triple_count >= chunk_size: # Try to parse this chunk - 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() @@ -409,7 +432,9 @@ def stream_turtle_chunks(file_path: Path, chunk_size: int = 10000) -> Iterator[l for s, p, o in graph: if isinstance(o, Literal): object_type = "literal" - object_datatype = str(o.datatype) if o.datatype else None + 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" @@ -493,7 +518,10 @@ def stream_turtle_chunks(file_path: Path, chunk_size: int = 10000) -> Iterator[l def stream_rdf_chunks( - file_path: Path, format: str = "turtle", chunk_size: int = 10000, num_workers: int | None = None + 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. @@ -525,8 +553,10 @@ def stream_rdf_chunks( 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]") + 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) @@ -619,21 +649,7 @@ def convert_rdf_to_hf_streaming( # 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) @@ -647,11 +663,11 @@ def convert_rdf_to_hf_streaming( # generators, so using simple print statements def dataset_generator(): nonlocal triple_count - 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 chunk_count, chunk in enumerate( + stream_rdf_chunks(input_path, rdf_format, chunk_size, num_workers) + ): for triple in chunk: triple_count += 1 yield triple @@ -705,11 +721,13 @@ def convert_rdf_to_hf_streaming( elapsed_time = time.time() - start_time if push_to_hub: - console.print(f"[yellow]Uploading to HuggingFace Hub: {hub_repo_id}...[/yellow]") + console.print("[yellow]Uploading to HuggingFace Hub: " + f"{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]") + console.print("[bold green]✓ Successfully uploaded to " + f"{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]") @@ -737,17 +755,20 @@ def convert_rdf_to_hf_streaming( "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, + "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("\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]") + 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: @@ -762,7 +783,11 @@ def main(): ) 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( + "output", + nargs="?", + type=Path, + help="Output directory for HuggingFace dataset (optional if --push-to-hub)") parser.add_argument( "-f", "--format", @@ -796,7 +821,12 @@ def main(): type=str, help="HuggingFace Hub repository ID (required if --push-to-hub)", ) - parser.add_argument("-v", "--verbose", action="store_true", help="Show verbose output") + parser.add_argument( + "-v", + "--verbose", + action="store_true", + help="Show verbose output" + ) args = parser.parse_args() diff --git a/scripts/rdf_to_hf_incremental.py b/scripts/rdf_to_hf_incremental.py index 54e9026..b6c77d0 100644 --- a/scripts/rdf_to_hf_incremental.py +++ b/scripts/rdf_to_hf_incremental.py @@ -54,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)) +import contextlib + 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 ( @@ -322,10 +324,8 @@ def open_rdf_line_stream_from_url( ) return lines, f"{name}:{member}" finally: - try: + with contextlib.suppress(OSError): os.remove(tmp_path) - except OSError: - pass # Compressed files and plain text decompressed = iter_decompressed_bytes(name, byte_iter)