run_pipeline.py

"""
============================================
Run Full Pipeline
============================================
Runs all steps sequentially.
Step 5 (training) requires GPU.

Author: Soham Dahivalkar
============================================
"""

import subprocess
import sys

STEPS = [
    ("step1_collect_injections.py",  "Collect Injection Prompts"),
    ("step2_collect_safe_prompts.py","Collect Safe Prompts"),
    ("step3_build_dataset.py",       "Build Balanced Dataset"),
    ("step4_upload_dataset.py",      "Upload Dataset to HuggingFace"),
    ("step5_train_model.py",         "Fine-Tune Model (GPU Required)"),
    ("step6_upload_model.py",        "Upload Model to HuggingFace"),
    ("step7_test_model.py",          "Test Published Model"),
]


def run_step(script: str, label: str) -> bool:
    print(f"\n{'=' * 60}")
    print(f"  Running: {label}")
    print(f"  Script:  {script}")
    print(f"{'=' * 60}")

    result = subprocess.run(
        [sys.executable, script],
        capture_output=False,
    )

    if result.returncode != 0:
        print(f"\n[ERROR] {script} failed with exit code {result.returncode}")
        return False

    print(f"\n[OK] {label} — Complete!")
    return True


if __name__ == "__main__":
    print("=" * 60)
    print("PROMPTSHIELD — Full Pipeline")
    print("=" * 60)

    # Check if running specific steps
    if len(sys.argv) > 1:
        step_range = sys.argv[1]  # e.g., "1-4" or "5-7"
        try:
            start, end = map(int, step_range.split("-"))
            steps_to_run = STEPS[start - 1: end]
        except Exception:
            print(f"Usage: python run_pipeline.py [start-end]")
            print(f"Example: python run_pipeline.py 1-4  (data only)")
            print(f"         python run_pipeline.py 5-7  (train + upload)")
            sys.exit(1)
    else:
        steps_to_run = STEPS

    failed = []
    for script, label in steps_to_run:
        ok = run_step(script, label)
        if not ok:
            failed.append(script)
            print(f"\n[ABORT] Pipeline stopped at: {script}")
            break

    print(f"\n{'=' * 60}")
    if not failed:
        print("✅ All steps completed successfully!")
    else:
        print(f"❌ Pipeline failed at: {failed}")
    print(f"{'=' * 60}")