add profile_diarization.py

profile_diarization.py

@@ -0,0 +1,62 @@
+#!/usr/bin/env python3
+"""
+Memory profiling script for diarization
+"""
+
+import sys
+import os
+sys.path.insert(0, os.path.join(os.path.dirname(__file__), 'src'))
+
+import numpy as np
+import soundfile as sf
+from memory_profiler import profile
+from diarization import init_speaker_embedding_extractor, perform_speaker_diarization_on_utterances
+from asr import transcribe_file
+
+def main():
+    # Load audio
+    audio_path = "example.mp3"
+    print(f"Loading audio from {audio_path}")
+    audio, sample_rate = sf.read(audio_path)
+    print(f"Audio loaded: {len(audio)} samples at {sample_rate}Hz")
+
+    # Resample to 16kHz if needed
+    if sample_rate != 16000:
+        from scipy.signal import resample
+        audio = resample(audio, int(len(audio) * 16000 / sample_rate))
+        sample_rate = 16000
+
+    # Ensure mono
+    if len(audio.shape) > 1:
+        audio = audio.mean(axis=1)
+
+    print(f"Audio prepared: {len(audio)} samples at {sample_rate}Hz")
+
+    # Perform ASR to get utterances
+    print("Starting ASR...")
+    utterances = []
+    for current, all_utterances in transcribe_file(
+        audio_path=audio_path,
+        vad_threshold=0.5,
+        model_name="moonshine/tiny",
+        backend="moonshine"
+    ):
+        utterances = all_utterances
+    print(f"ASR completed: {len(utterances)} utterances")
+
+    # Initialize embedding extractor
+    print("Initializing embedding extractor...")
+    embedding_extractor, config_dict = init_speaker_embedding_extractor()
+    if embedding_extractor is None:
+        print("Failed to initialize embedding extractor")
+        return
+
+    # Perform diarization with profiling
+    print("Starting diarization...")
+    diarization_result = perform_speaker_diarization_on_utterances(
+        audio, sample_rate, utterances, embedding_extractor, config_dict
+    )
+    print(f"Diarization completed: {len(diarization_result)} results")
+
+if __name__ == "__main__":
+    main()