Files
dataset-uploader/scripts/specialized_converter_progress.py

136 lines
4.7 KiB
Python
Executable File

#!/usr/bin/env python3
"""
Wrapper for specialized converters (FB15k-237, NELL-995, ConceptNet) that adds progress output
for subprocess monitoring in parallel processing.
"""
import subprocess
import sys
import re
from pathlib import Path
import time
def monitor_and_output_progress(proc, converter_name="converter"):
"""Monitor subprocess output and emit simple progress markers."""
last_progress = 0
last_emit_time = 0
MIN_EMIT_INTERVAL = 0.5 # Minimum time between progress emissions
# Track phases
phase_patterns = [
# Common patterns across converters
(r"Converting.*to HuggingFace", lambda m: 5),
(r"Counting", lambda m: 10),
(r"Found \d+ (assertions|triples|lines)", lambda m: 15),
(r"Parsing", lambda m: 20),
(r"Reading.*file", lambda m: 25),
(r"Processing.*split", lambda m: 30),
# Progress indicators
(r"train.*processed|train.*parsed", lambda m: 40),
(r"valid.*processed|valid.*parsed", lambda m: 50),
(r"test.*processed|test.*parsed", lambda m: 60),
# Dataset creation
(r"Creating.*dataset|Building.*dataset", lambda m: 85),
(r"Saving.*dataset", lambda m: 95),
(r"Dataset saved|Successfully|Complete", lambda m: 100),
# Generic progress
(r"✓.*Parsed (\d+)", lambda m: min(80, 20 + int(m.group(1)) / 1000)),
(r"(\d+)%", lambda m: float(m.group(1))),
]
# Start time for time-based progress
start_time = time.time()
for line in iter(proc.stdout.readline, ""):
if not line:
break
# Pass through the original output
print(line, end="", flush=True)
# Look for explicit PROGRESS markers first
progress_match = re.search(r"PROGRESS:\s*(\d+)", line)
if progress_match:
progress = float(progress_match.group(1))
current_time = time.time()
# Only emit if progress increased
if progress > last_progress:
last_progress = progress
last_emit_time = current_time
# Already in output, don't double-emit
continue
# Try to infer progress from patterns
for pattern, extractor in phase_patterns:
match = re.search(pattern, line, re.IGNORECASE)
if match:
try:
progress = extractor(match)
current_time = time.time()
if progress > last_progress and (current_time - last_emit_time) > MIN_EMIT_INTERVAL:
last_progress = progress
last_emit_time = current_time
print(f"\nPROGRESS: {progress:.0f}", flush=True)
break
except:
pass
# Time-based progress as fallback
elapsed = time.time() - start_time
if elapsed > 5 and last_progress < 50: # After 5 seconds, show some progress
time_progress = min(80, elapsed * 2) # Rough estimate
if time_progress > last_progress + 10:
last_progress = time_progress
print(f"\nPROGRESS: {time_progress:.0f}", flush=True)
# Ensure we hit 100% on successful completion
if proc.wait() == 0 and last_progress < 100:
print("\nPROGRESS: 100", flush=True)
return proc.returncode
def main():
"""Run the converter with progress monitoring."""
# Determine which converter to run based on script name
script_path = Path(sys.argv[0])
if "fb15k" in script_path.name.lower():
target_script = "convert_fb15k237_to_hf.py"
converter_name = "FB15k-237"
elif "nell" in script_path.name.lower():
target_script = "convert_nell995_to_hf.py"
converter_name = "NELL-995"
elif "conceptnet" in script_path.name.lower():
target_script = "convert_conceptnet_to_hf.py"
converter_name = "ConceptNet"
else:
# Default to passing through the first argument as the script name
if len(sys.argv) > 1:
target_script = sys.argv[1]
converter_name = Path(target_script).stem
remaining_args = sys.argv[2:]
else:
print("Error: No converter specified", file=sys.stderr)
return 1
# Build command
cmd = [sys.executable, "-u", str(Path(__file__).parent / target_script)]
if "remaining_args" in locals():
cmd.extend(remaining_args)
else:
cmd.extend(sys.argv[1:])
# Run with monitoring
proc = subprocess.Popen(cmd, stdout=subprocess.PIPE, stderr=subprocess.STDOUT, text=True, bufsize=1)
return monitor_and_output_progress(proc, converter_name)
if __name__ == "__main__":
sys.exit(main())