Spaces:
Runtime error
Runtime error
File size: 3,272 Bytes
ca4ed58 c3c40f4 ca4ed58 c3c40f4 ca4ed58 c3c40f4 08c0be3 c3c40f4 08c0be3 c3c40f4 8b33e3f c3c40f4 8b33e3f c3c40f4 ca4ed58 c3c40f4 ca4ed58 c3c40f4 ca4ed58 c3c40f4 ca4ed58 c3c40f4 7603254 d9da639 7603254 d9da639 7603254 c3c40f4 d9da639 c3c40f4 d9da639 c3c40f4 | 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 | from typing import List, Dict
from faster_whisper import WhisperModel
import os
import torch
from pyannote.audio import Pipeline
# Load environment variables
WHISPER_MODEL_SIZE = os.getenv("WHISPER_MODEL_SIZE", "base")
HF_TOKEN = os.getenv("HF_TOKEN") # For pyannote
# Initialize Whisper model
try:
_whisper = WhisperModel(WHISPER_MODEL_SIZE, device="cpu", compute_type="int8")
except Exception as e:
print(f"Error loading Whisper model: {e}")
_whisper = WhisperModel("base", device="cpu", compute_type="int8")
# Initialize pyannote diarization pipeline
if HF_TOKEN:
try:
diarization_pipeline = Pipeline.from_pretrained(
"pyannote/speaker-diarization-3.1", token=HF_TOKEN
)
# Move pipeline to CPU if no GPU is available
if not torch.cuda.is_available():
diarization_pipeline = diarization_pipeline.to(torch.device("cpu"))
except Exception as e:
print(f"Error loading pyannote pipeline: {e}")
diarization_pipeline = None
else:
print("HF_TOKEN not set, skipping diarization.")
diarization_pipeline = None
def transcribe(audio_path: str) -> List[Dict]:
"""
Transcribe an audio file using faster-whisper and combine with
pyannote.audio for speaker diarization.
"""
# 1. Transcribe with Whisper
segments, _ = _whisper.transcribe(
audio_path, language="en", vad_filter=True, beam_size=1
)
whisper_segments = []
for seg in segments:
whisper_segments.append(
{"start": float(seg.start), "end": float(seg.end), "text": seg.text.strip()}
)
if not diarization_pipeline:
# If diarization is not available, return with a single speaker
for seg in whisper_segments:
seg["speaker"] = "A"
return whisper_segments
# 2. Perform Diarization
try:
diarization = diarization_pipeline(audio_path)
# --- DEBUGGING LINES (optional) ---
print("\n--- Diarization Output ---")
print(f"Type of diarization object: {type(diarization)}")
print("Diarization object content:")
print(diarization)
print("--- End Diarization Output ---\n")
# ----------------------------------
except Exception as e:
print(f"Error during diarization: {e}")
for seg in whisper_segments:
seg["speaker"] = "A"
return whisper_segments
# 3. Assign Speaker to Segments
out_segments = []
# For pyannote 4.x, use diarization.speaker_diarization
annotation = diarization.speaker_diarization # This is an Annotation object
for seg in whisper_segments:
midpoint = seg["start"] + (seg["end"] - seg["start"]) / 2
speaker = "UNKNOWN"
for turn, _, speaker_label in annotation.itertracks(yield_label=True):
if turn.start <= midpoint <= turn.end:
speaker = speaker_label
break
out_segments.append(
{
"start": seg["start"],
"end": seg["end"],
"speaker": speaker,
"text": seg["text"],
}
)
if not out_segments:
return [{"start": 0.0, "end": 0.0, "speaker": "A", "text": ""}]
return out_segments
|