Transcriptor/workers/diarize_worker.py

#!/home/superti/miniconda3/envs/diarization/bin/python
import warnings
warnings.filterwarnings("ignore")
import json
from pyannote.audio import Pipeline
import os
from torch import device as torch_device
import sys
import dotenv
dotenv.load_dotenv()

def diarize(file):
    pipeline = Pipeline.from_pretrained(
        "pyannote/speaker-diarization-3.1",
        token=os.environ["HF_TOKEN"],
    )
    pipeline.to(torch_device("cuda"))
    if not pipeline:
        raise RuntimeError("Pipeline not found")

    segments = pipeline(file)
    diariz = []
    for turn, speaker in segments.speaker_diarization:
        diariz.append({
            "start": turn.start,
            "end": turn.end,
            "speaker": speaker
        })

    print(json.dumps(diariz))

if __name__ == "__main__":
    diarize(sys.argv[1])