Create asr_diarization.py

pipeline/asr_diarization.py

@@ -0,0 +1,25 @@
+import whisperx
+
+_asr_model = None
+_alignment_model = None
+_metadata = None
+
+def load_asr(device="cpu", model_size="medium"):
+    global _asr_model
+    if _asr_model is None:
+        _asr_model = whisperx.load_model(model_size, device=device)
+    return _asr_model
+
+def transcribe_and_align(audio_path: str, device="cpu"):
+    model = load_asr(device=device)
+    result = model.transcribe(audio_path, batch_size=16)
+
+    global _alignment_model, _metadata
+    if _alignment_model is None or _metadata is None:
+        _alignment_model, _metadata = whisperx.load_align_model(
+            language_code=result["language"], device=device
+        )
+    result_aligned = whisperx.align(
+        result["segments"], _alignment_model, _metadata, audio_path
+    )
+    return result, result_aligned