PodcastAdDetectionModel/run_pipeline.py at main · jonasgunklach/PodcastAdDetectionModel · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
"""
Main pipeline script to run all stages sequentially.
Usage: python run_pipeline.py [--skip-download] [--skip-transcribe] [--skip-ads] [--skip-train]
"""

import argparse
import sys
from pathlib import Path

# Add current directory to path
sys.path.insert(0, str(Path(__file__).parent))

def run_download():
    """Run download stage."""
    print("\n" + "="*60)
    print("STAGE 1: Downloading Episodes")
    print("="*60)
    try:
        import download_episodes
        download_episodes.download_episodes()
        return True
    except Exception as e:
        print(f"Error in download stage: {e}")
        return False

def run_transcribe():
    """Run transcription stage."""
    print("\n" + "="*60)
    print("STAGE 2: Transcribing Audio")
    print("="*60)
    try:
        import transcribe_audio
        transcribe_audio.transcribe_episodes()
        return True
    except Exception as e:
        print(f"Error in transcription stage: {e}")
        return False

def run_detect_ads():
    """Run ad detection stage."""
    print("\n" + "="*60)
    print("STAGE 3: Detecting Ads")
    print("="*60)
    try:
        import detect_ads
        detect_ads.detect_ads()
        return True
    except Exception as e:
        print(f"Error in ad detection stage: {e}")
        return False

def run_train():
    """Run training stage."""
    print("\n" + "="*60)
    print("STAGE 4: Training Model")
    print("="*60)
    try:
        import train_model
        train_model.train_model()
        return True
    except Exception as e:
        print(f"Error in training stage: {e}")
        return False

def main():
    """Run the complete pipeline."""
    parser = argparse.ArgumentParser(description='Run podcast ad detection pipeline')
    parser.add_argument('--skip-download', action='store_true', help='Skip download stage')
    parser.add_argument('--skip-transcribe', action='store_true', help='Skip transcription stage')
    parser.add_argument('--skip-ads', action='store_true', help='Skip ad detection stage')
    parser.add_argument('--skip-train', action='store_true', help='Skip training stage')
    parser.add_argument('--stage', type=int, choices=[1,2,3,4], help='Run only specific stage')

    args = parser.parse_args()

    print("="*60)
    print("Podcast Ad Detection Pipeline")
    print("="*60)

    if args.stage:
        # Run only specific stage
        stages = {
            1: run_download,
            2: run_transcribe,
            3: run_detect_ads,
            4: run_train
        }
        stages[args.stage]()
    else:
        # Run all stages unless skipped
        success = True

        if not args.skip_download:
            success = run_download() and success
        else:
            print("\nSkipping download stage...")

        if not args.skip_transcribe:
            success = run_transcribe() and success
        else:
            print("\nSkipping transcription stage...")

        if not args.skip_ads:
            success = run_detect_ads() and success
        else:
            print("\nSkipping ad detection stage...")

        if not args.skip_train:
            success = run_train() and success
        else:
            print("\nSkipping training stage...")

        if success:
            print("\n" + "="*60)
            print("Pipeline completed successfully!")
            print("="*60)
        else:
            print("\n" + "="*60)
            print("Pipeline completed with errors. Check logs above.")
            print("="*60)
            sys.exit(1)

if __name__ == "__main__":
    main()