asrdiarization-handler / diarization_utils.py

Update diarization_utils.py

952ebd6 verified 3 months ago

No virus

5.01 kB

	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
	)