Adding speaker segmentation

app.py

@@ -1,11 +1,13 @@
 import gradio as gr
 from transformers import pipeline, Wav2Vec2ProcessorWithLM
+from pyannote.audio import Pipeline
 from librosa import load, resample
 from rpunct import RestorePuncts
 
 asr_model = 'patrickvonplaten/wav2vec2-base-100h-with-lm'
 processor = Wav2Vec2ProcessorWithLM.from_pretrained(asr_model)
 asr = pipeline('automatic-speech-recognition', model=asr_model, tokenizer=processor.tokenizer, feature_extractor=processor.feature_extractor, decoder=processor.decoder)
+speaker_segmentation = Pipeline.from_pretrained("pyannote/speaker-segmentation")
 
 rpunct = RestorePuncts()
 
@@ -13,19 +15,36 @@ def transcribe(filepath):
 	speech, sampling_rate = load(filepath)
 	if sampling_rate != 16000:
 		speech = resample(speech, sampling_rate, 16000)
-	text = asr(speech)['text']
-	text = rpunct.punctuate(text.lower())
-	return text
+	speaker_output = speaker_segmentation(speech)
+	text = asr(speech, return_timestamps="word")
+
+	full_text = text['text'].lower()
+	chunks = text['chunks']
+
+	diarizaed_output = ""
+	i = 0
+	for turn, _, speaker in speaker_output.itertracks(yield_label=True):
+		diarized = ""
+		while i < len(chunks) and chunks[i]['timestamp'][1] <= turn.end:
+			diarized += chunks[i]['text'].lower() + ' '
+			i += 1
+
+		if diarized != "":
+			diarized = rpunct.punctuate(diarized)
+			diarized_output += "{}: ''{}'' from {:.3f}-{:.3f}\n".format(speaker,diarized,turn.start,turn.end)
+
+	return diarizaed_output, full_text
 
 mic = gr.inputs.Audio(source='microphone', type='filepath', label='Speech input', optional=False)
 
-transcript = gr.outputs.Textbox(type='auto', label='Transcription')
+diarized_transcript = gr.outputs.Textbox(type='auto', label='Diarized Output')
+full_transcript = gr.outputs.Textbox(type='auto', label='Full Transcript')
 
 iface = gr.Interface(
 	theme='huggingface',
 	description='Testing transcription',
 	fn=transcribe,
 	inputs=[mic],
-	outputs=[transcript]
+	outputs=[diarized_transcript, full_transcript]
 )
 iface.launch()

requirements.txt

@@ -4,3 +4,4 @@ librosa
 pyctcdecode
 pypi-kenlm
 git+https://github.com/anuragshas/rpunct.git
+https://github.com/pyannote/pyannote-audio/archive/develop.zip