import gradio as gr
from pyannote.audio import Pipeline
from transformers import pipeline

asr = pipeline(
    "automatic-speech-recognition",
    model="facebook/wav2vec2-large-960h-lv60-self",
    feature_extractor="facebook/wav2vec2-large-960h-lv60-self",
    
)
pipeline1 = Pipeline.from_pretrained("pyannote/speaker-segmentation")

def diarization(file_input,mic_input,selection):
    mic_path = None if mic_input is None else mic_input.name
    audio = file_input if selection == "Upload" else mic_path
    if audio is None:
        return "Please check your inputs!", ""

    speaker_output = pipeline1(audio)
    text_output = asr(audio,return_timestamps="word")
    
    full_text = text_output['text'].lower()
    chunks = text_output['chunks']

    diarized_output = ""
    i = 0
    for turn, _, speaker in speaker_output.itertracks(yield_label=True):
        diarized = ""
        while i < len(chunks):
            time_index = chunks[i]['timestamp'][1]
            if time_index >= turn.start and time_index <= turn.end:
                diarized += chunks[i]['text'].lower() + ' '
            if time_index >= turn.end: break
            i += 1

        diarized_output += "{} said '{}' from {:.3f} to {:.3f}\n".format(speaker,diarized,turn.start,turn.end)
        
    return diarized_output, full_text

title = "Speech Recognition with Speaker Diarization"
description = "Speaker Diarization is the act of attributing parts of the audio recording to different speakers. This space aims to distinguish the speakers and apply speech-to-text from a given input audio file. Pre-trained models from Pyannote[1] for the Speaker Diarization and [2]."
article = "<p style='text-align: center'><a href='https://github.com/pyannote/pyannote-audio' target='_blank'>[1] Pyannote - Speaker Diarization model</a></p>"
inputs = [gr.inputs.Audio(source="upload", type="filepath", label="Upload your audio file here:", optional=True),
        gr.inputs.Audio(source="microphone", type="file",label="Or use your Microphone:", optional=True),
        gr.inputs.Radio(["Upload","Microphone"], type="value", label="Select which input:")]
outputs = [gr.outputs.Textbox(type="auto", label="Diarized Output"),
        gr.outputs.Textbox(type="auto",label="Full ASR Text for comparison")]
examples = [["test_audio1.wav",None,"Upload"],
            ["test_audio2.wav",None,"Upload"]]

app = gr.Interface(fn=diarization,
                inputs=inputs,
                outputs=outputs,
                examples=examples,
                title=title,
                description=description,
                article=article,
                allow_flagging=False)
app.launch()