Delete diarization_utils.py

diarization_utils.py

@@ -1,141 +0,0 @@
-import torch
-import numpy as np
-from torchaudio import functional as F
-from transformers.pipelines.audio_utils import ffmpeg_read
-from starlette.exceptions import HTTPException
-import sys
-
-# Code from insanely-fast-whisper:
-# https://github.com/Vaibhavs10/insanely-fast-whisper
-
-import logging
-logger = logging.getLogger(__name__)
-
-def preprocess_inputs(inputs, sampling_rate):
-    inputs = ffmpeg_read(inputs, sampling_rate)
-
-    if sampling_rate != 16000:
-        inputs = F.resample(
-            torch.from_numpy(inputs), sampling_rate, 16000
-        ).numpy()
-
-    if len(inputs.shape) != 1:
-        logger.error(f"Diarization pipeline expecs single channel audio, received {inputs.shape}")
-        raise HTTPException(
-            status_code=400,
-            detail=f"Diarization pipeline expecs single channel audio, received {inputs.shape}"
-        )
-
-    # diarization model expects float32 torch tensor of shape `(channels, seq_len)`
-    diarizer_inputs = torch.from_numpy(inputs).float()
-    diarizer_inputs = diarizer_inputs.unsqueeze(0)
-
-    return inputs, diarizer_inputs
-
-
-def diarize_audio(diarizer_inputs, diarization_pipeline, parameters):
-    diarization = diarization_pipeline(
-        {"waveform": diarizer_inputs, "sample_rate": parameters.sampling_rate},
-        num_speakers=parameters.num_speakers,
-        min_speakers=parameters.min_speakers,
-        max_speakers=parameters.max_speakers,
-    )
-
-    segments = []
-    for segment, track, label in diarization.itertracks(yield_label=True):
-        segments.append(
-            {
-                "segment": {"start": segment.start, "end": segment.end},
-                "track": track,
-                "label": label,
-            }
-        )
-
-    # diarizer output may contain consecutive segments from the same speaker (e.g. {(0 -> 1, speaker_1), (1 -> 1.5, speaker_1), ...})
-    # we combine these segments to give overall timestamps for each speaker's turn (e.g. {(0 -> 1.5, speaker_1), ...})
-    new_segments = []
-    prev_segment = cur_segment = segments[0]
-
-    for i in range(1, len(segments)):
-        cur_segment = segments[i]
-
-        # check if we have changed speaker ("label")
-        if cur_segment["label"] != prev_segment["label"] and i < len(segments):
-            # add the start/end times for the super-segment to the new list
-            new_segments.append(
-                {
-                    "segment": {
-                        "start": prev_segment["segment"]["start"],
-                        "end": cur_segment["segment"]["start"],
-                    },
-                    "speaker": prev_segment["label"],
-                }
-            )
-            prev_segment = segments[i]
-
-    # add the last segment(s) if there was no speaker change
-    new_segments.append(
-        {
-            "segment": {
-                "start": prev_segment["segment"]["start"],
-                "end": cur_segment["segment"]["end"],
-            },
-            "speaker": prev_segment["label"],
-        }
-    )
-
-    return new_segments
-
-
-def post_process_segments_and_transcripts(new_segments, transcript, group_by_speaker) -> list:
-    # get the end timestamps for each chunk from the ASR output
-    end_timestamps = np.array(
-        [chunk["timestamp"][-1] if chunk["timestamp"][-1] is not None else sys.float_info.max for chunk in transcript])
-    segmented_preds = []
-
-    # align the diarizer timestamps and the ASR timestamps
-    for segment in new_segments:
-        # get the diarizer end timestamp
-        end_time = segment["segment"]["end"]
-        # find the ASR end timestamp that is closest to the diarizer's end timestamp and cut the transcript to here
-        upto_idx = np.argmin(np.abs(end_timestamps - end_time))
-
-        if group_by_speaker:
-            segmented_preds.append(
-                {
-                    "speaker": segment["speaker"],
-                    "text": "".join(
-                        [chunk["text"] for chunk in transcript[: upto_idx + 1]]
-                    ),
-                    "timestamp": (
-                        transcript[0]["timestamp"][0],
-                        transcript[upto_idx]["timestamp"][1],
-                    ),
-                }
-            )
-        else:
-            for i in range(upto_idx + 1):
-                segmented_preds.append({"speaker": segment["speaker"], **transcript[i]})
-
-        # crop the transcripts and timestamp lists according to the latest timestamp (for faster argmin)
-        transcript = transcript[upto_idx + 1:]
-        end_timestamps = end_timestamps[upto_idx + 1:]
-
-        if len(end_timestamps) == 0:
-            break
-
-    return segmented_preds
-
-
-def diarize(diarization_pipeline, file, parameters, asr_outputs):
-    _, diarizer_inputs = preprocess_inputs(file, parameters.sampling_rate)
-
-    segments = diarize_audio(
-        diarizer_inputs, 
-        diarization_pipeline, 
-        parameters
-    )
-
-    return post_process_segments_and_transcripts(
-        segments, asr_outputs["chunks"], group_by_speaker=False
-    )