126 lines
4.2 KiB
Python
Executable File
126 lines
4.2 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 re
|
|
import subprocess
|
|
import sys
|
|
import time
|
|
from pathlib import Path
|
|
from typing import Any
|
|
|
|
from scripts.common import print_progress
|
|
|
|
|
|
def monitor_and_output_progress(proc: subprocess.Popen[Any]) -> int:
|
|
"""Monitor subprocess output and emit simple progress markers."""
|
|
|
|
last_progress = 0.0
|
|
last_emit_time = 0.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: int(m.group(1))),
|
|
]
|
|
|
|
# Start time for time-based progress
|
|
start_time = time.time()
|
|
|
|
assert proc.stdout is not None
|
|
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
|
|
|
|
last_progress = print_progress(
|
|
last_emit_time, last_progress, line, min_emit_interval, phase_patterns
|
|
)
|
|
|
|
# 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, int(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() -> int:
|
|
"""Run the converter with progress monitoring."""
|
|
# Determine which converter to run based on script name
|
|
script_path = Path(sys.argv[0])
|
|
remaining_args = []
|
|
|
|
if "fb15k" in script_path.name.lower():
|
|
target_script = "convert_fb15k237_to_hf.py"
|
|
elif "nell" in script_path.name.lower():
|
|
target_script = "convert_nell995_to_hf.py"
|
|
elif "conceptnet" in script_path.name.lower():
|
|
target_script = "convert_conceptnet_to_hf.py"
|
|
else:
|
|
# Default to passing through the first argument as the script name
|
|
if len(sys.argv) > 1:
|
|
target_script = sys.argv[1]
|
|
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
|
|
with subprocess.Popen(
|
|
cmd, stdout=subprocess.PIPE, stderr=subprocess.STDOUT, text=True, bufsize=1
|
|
) as proc:
|
|
return monitor_and_output_progress(proc)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
sys.exit(main())
|