Files
soprano-factory/custom/1_data_prep.py
T
2026-02-09 01:01:44 -05:00

97 lines
2.9 KiB
Python

import sys
import subprocess
import os
import json
import argparse
from pathlib import Path
# Windows UTF-8 Support
sys.stdout.reconfigure(encoding='utf-8')
def run_command(cmd, desc):
print(f"\n[Step 1] Starting: {desc}")
print(f"Command: {' '.join(cmd)}")
try:
subprocess.check_call(cmd)
print(f"[Step 1] Completed: {desc}\n")
except subprocess.CalledProcessError as e:
print(f"[Step 1] Error during {desc}: {e}")
sys.exit(1)
def verify_data_prep(output_dir):
print("[Step 1] Verifying outputs...")
train_json = output_dir / "train.json"
val_json = output_dir / "val.json"
if not train_json.exists():
print(f"[ERROR] train.json not found at {train_json}")
sys.exit(1)
if not val_json.exists():
print(f"[ERROR] val.json not found at {val_json}")
sys.exit(1)
try:
with open(train_json, "r", encoding="utf-8") as f:
train_data = json.load(f)
with open(val_json, "r", encoding="utf-8") as f:
val_data = json.load(f)
print(f"[SUCCESS] Data Prep Complete!")
print(f" - Train Samples: {len(train_data)}")
print(f" - Val Samples: {len(val_data)}")
if len(train_data) == 0:
print("[WARNING] Train dataset is empty! Check your input metadata.")
except Exception as e:
print(f"[ERROR] Failed to read JSON outputs: {e}")
sys.exit(1)
def main():
parser = argparse.ArgumentParser(description="Soprano Custom Pipeline - Step 1: Data Prep")
parser.add_argument("--test", action="store_true", help="Run in test mode (limit to 50 samples)")
args = parser.parse_args()
root_dir = Path(os.getcwd())
dataset_dir = root_dir / "mio_dataset"
custom_dir = root_dir / "custom"
encoder_path = root_dir / "weights" / "codec" / "encoder.pth"
# Ensure custom directory exists
custom_dir.mkdir(exist_ok=True)
# Check inputs
if not dataset_dir.exists():
print(f"[ERROR] Dataset directory not found: {dataset_dir}")
print("Please ensure 'mio_dataset' is in the project root.")
sys.exit(1)
if not encoder_path.exists():
print(f"[ERROR] Encoder checkpoint not found: {encoder_path}")
sys.exit(1)
# Run Generation
script_path = root_dir / "generate_dataset.py"
if not script_path.exists():
print(f"[ERROR] generate_dataset.py not found at {script_path}")
sys.exit(1)
cmd = [
sys.executable, str(script_path),
"--input-dir", str(dataset_dir),
"--output-dir", str(custom_dir),
"--encoder-ckpt", str(encoder_path)
]
if args.test:
cmd.extend(["--limit", "50"])
run_command(cmd, "Data Preparation (Audio -> Tokens)")
# Verify
verify_data_prep(custom_dir)
if __name__ == "__main__":
main()