Update app.py

app.py

@@ -22,14 +22,38 @@ from pyannote.audio import Audio
 from pyannote.core import Segment
 import wave
 import contextlib
-from sklearn.cluster import KMeans
-
+from sklearn.cluster import AgglomerativeClustering
 import numpy as np
 import json
 from datetime import timedelta
 
+from transformers import T5ForConditionalGeneration, T5Tokenizer
+
 __FILES = set()
-      
+wispher_models = list(whisper._MODELS.keys())
+
+def correct_grammar(input_text,num_return_sequences=1):
+    torch_device = 'cuda' if torch.cuda.is_available() else 'cpu'
+    tokenizer = T5Tokenizer.from_pretrained('deep-learning-analytics/GrammarCorrector')
+    model = T5ForConditionalGeneration.from_pretrained('deep-learning-analytics/GrammarCorrector').to(torch_device)
+    batch = tokenizer([input_text],truncation=True,padding='max_length',max_length=len(input_text), return_tensors="pt").to(torch_device)
+    results = model.generate(**batch,max_length=len(input_text),num_beams=2, num_return_sequences=num_return_sequences, temperature=1.5)
+    generated_sequences = []
+    for generated_sequence_idx, generated_sequence in enumerate(results):
+        text = tokenizer.decode(generated_sequence, clean_up_tokenization_spaces=True, skip_special_tokens=True)
+        generated_sequences.append(text)
+    generated_text = "".join(generated_sequences)
+    _generated_text = ""
+    for idx, _sentence in enumerate(generated_text.split('.'), 0):
+        if not idx:
+            _generated_text+=_sentence+'.'
+        elif _sentence[:1]!=' ':
+            _generated_text+=' '+_sentence+'.'
+        elif _sentence[:1]=='':
+            pass
+        else:
+            _generated_text+=_sentence+'.'
+    return _generated_text
 
 def CreateFile(filename):
     __FILES.add(filename)
@@ -140,14 +164,16 @@ def Transcribe_V1(NumberOfSpeakers, SpeakerNames="", audio="temp_audio.wav"):
     return (t_text, ({ "data": [{"speaker": speaker, "text": text} for speaker, text in conversation]}))
 
 
-def Transcribe_V2(num_speakers, speaker_names, audio="temp_audio.wav"):
-    model = whisper.load_model("medium")
+def Transcribe_V2(model, num_speakers, speaker_names, audio="temp_audio.wav"):
+    model = whisper.load_model(model)
     # embedding_model = SpeechBrainPretrainedSpeakerEmbedding("speechbrain/spkrec-ecapa-voxceleb")
+    
     embedding_model = SpeechBrainPretrainedSpeakerEmbedding(
         "speechbrain/spkrec-ecapa-voxceleb",
         device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
     )
     SPEAKER_DICT = {}
+    default_speaker_names = ['A','B','C','D','E','F','G','H','I','J','K','L','M','N','O','P','Q','R','S','T','U','V','W','X','Y','Z']
     SPEAKERS = [speaker.strip() for speaker in speaker_names.split(',') if len(speaker)]
     def GetSpeaker(sp):
         speaker = sp
@@ -156,6 +182,10 @@ def Transcribe_V2(num_speakers, speaker_names, audio="temp_audio.wav"):
                 t = SPEAKERS.pop(0)
                 SPEAKER_DICT[sp] = t
                 speaker = SPEAKER_DICT[sp]
+            elif len(default_speaker_names):
+                t = default_speaker_names.pop(0)
+                SPEAKER_DICT[sp] = t
+                speaker = SPEAKER_DICT[sp]
         else:
             speaker = SPEAKER_DICT[sp]
         return speaker
@@ -168,6 +198,9 @@ def Transcribe_V2(num_speakers, speaker_names, audio="temp_audio.wav"):
             return s
         as_audio = AudioSegment.from_wav(audio)
         DEMO_FILE = {'uri': 'blabal', 'audio': audio}
+        hparams = pipeline.parameters(instantiated=True)
+        hparams["segmentation"]["min_duration_off"] -= 0.25
+        pipeline.instantiate(hparams)
         if num_speakers:
             dz = pipeline(DEMO_FILE, num_speakers=num_speakers)  
         else:
@@ -201,6 +234,8 @@ def Transcribe_V2(num_speakers, speaker_names, audio="temp_audio.wav"):
             #         conversation.append([GetSpeaker(segment["speaker"]), segment["text"][1:]]) # segment["speaker"] + ' ' + str(time(segment["start"])) + '\n\n'
             # conversation[-1][1] += segment["text"][1:]
         # return output
+        for idx in range(len(conversation)):
+            conversation[idx][3] = correct_grammar(conversation[idx][3])
         return ("".join([f"[{start}] - {speaker} \n{text}\n" for start, end, speaker, text in conversation])), ({ "data": [{"start": start, "end":end, "speaker": speaker, "text": text} for start, end, speaker, text in conversation]})
 
     def get_duration(path):
@@ -224,7 +259,7 @@ def Transcribe_V2(num_speakers, speaker_names, audio="temp_audio.wav"):
         return embedding_model(waveform[None])
 
     def add_speaker_labels(segments, embeddings, num_speakers):
-        clustering = KMeans(num_speakers).fit(embeddings)
+        clustering = AgglomerativeClustering(num_speakers).fit(embeddings)
         labels = clustering.labels_
         for i in range(len(segments)):
             segments[i]["speaker"] = 'SPEAKER ' + str(labels[i] + 1)
@@ -236,9 +271,9 @@ def Transcribe_V2(num_speakers, speaker_names, audio="temp_audio.wav"):
     if duration > 4 * 60 * 60:
         return "Audio duration too long"
 
-    print(json.dumps(diarization(audio)))
+    # print(json.dumps(diarization(audio)))
     result = model.transcribe(audio)
-    print(json.dumps(result))
+    # print(json.dumps(result))
 
     segments = result["segments"]
 
@@ -251,7 +286,7 @@ def Transcribe_V2(num_speakers, speaker_names, audio="temp_audio.wav"):
     return get_output(segments)
     # return output
 
-def AudioTranscribe(NumberOfSpeakers=None, SpeakerNames="", audio="", retries=5):
+def AudioTranscribe(NumberOfSpeakers=None, SpeakerNames="", audio="", retries=5, model='base'):
     print(f"{NumberOfSpeakers}, {SpeakerNames}, {retries}")
     if retries:
         # subprocess.call(['ffmpeg', '-i', audio,'temp_audio.wav'])
@@ -262,11 +297,11 @@ def AudioTranscribe(NumberOfSpeakers=None, SpeakerNames="", audio="", retries=5)
             return AudioTranscribe(NumberOfSpeakers, SpeakerNames, audio, retries-1)
         if not (os.path.isfile("temp_audio.wav")):
             return AudioTranscribe(NumberOfSpeakers, SpeakerNames, audio, retries-1)
-        return Transcribe_V2(NumberOfSpeakers, SpeakerNames)
+        return Transcribe_V2(model, NumberOfSpeakers, SpeakerNames)
     else:
         raise gr.Error("There is some issue ith Audio Transcriber. Please try again later!")
 
-def VideoTranscribe(NumberOfSpeakers=None, SpeakerNames="", video="", retries=5):
+def VideoTranscribe(NumberOfSpeakers=None, SpeakerNames="", video="", retries=5, model='base'):
     if retries:
         try:
             clip = mp.VideoFileClip(video)
@@ -278,12 +313,11 @@ def VideoTranscribe(NumberOfSpeakers=None, SpeakerNames="", video="", retries=5)
             return VideoTranscribe(NumberOfSpeakers, SpeakerNames, video, retries-1)
         if not (os.path.isfile("temp_audio.wav")):
             return VideoTranscribe(NumberOfSpeakers, SpeakerNames, video, retries-1)
-        return Transcribe_V2(NumberOfSpeakers, SpeakerNames)
+        return Transcribe_V2(model, NumberOfSpeakers, SpeakerNames)
     else:
         raise gr.Error("There is some issue ith Video Transcriber. Please try again later!")
-    return Transcribe_V2(NumberOfSpeakers, SpeakerNames)
 
-def YoutubeTranscribe(NumberOfSpeakers=None, SpeakerNames="", URL="", retries = 5):
+def YoutubeTranscribe(NumberOfSpeakers=None, SpeakerNames="", URL="", retries = 5, model='base'):
     if retries:
         if "youtu" not in URL.lower():
             raise gr.Error(f"{URL} is not a valid youtube URL.")
@@ -305,42 +339,28 @@ def YoutubeTranscribe(NumberOfSpeakers=None, SpeakerNames="", URL="", retries =
             stream = ffmpeg.input('temp_audio.m4a')
             stream = ffmpeg.output(stream, 'temp_audio.wav')
             RemoveFile("temp_audio.m4a")
-            return Transcribe_V2(NumberOfSpeakers, SpeakerNames)
+            return Transcribe_V2(model, NumberOfSpeakers, SpeakerNames)
     else:
         raise gr.Error(f"Unable to get video from {URL}")
  
-ut = gr.Interface(
-    fn=YoutubeTranscribe,
-    inputs=[gr.Number(label="Number of Speakers", placeholder="2"), gr.Textbox(label="Name of the Speakers (ordered by the time they speak and separated by comma)", placeholder="If Speaker 1 is first to speak followed by Speaker 2 then -> Speaker 1, Speaker 2"), gr.Textbox(label="Youtube Link", placeholder="https://www.youtube.com/watch?v=GECcjrYHH8w"),],
-    outputs=[gr.Textbox(label="Transcribed Text", lines=15), gr.JSON(label="Transcribed JSON")]
-)
-vt = gr.Interface(
-    fn=VideoTranscribe,
-    inputs=[gr.Number(label="Number of Speakers", placeholder="2"), gr.Textbox(label="Name of the Speakers (ordered by the time they speak and separated by comma)", placeholder="If Speaker 1 is first to speak followed by Speaker 2 then -> Speaker 1, Speaker 2"), 'video'],
-    outputs=[gr.Textbox(label="Transcribed Text", lines=15), gr.JSON(label="Transcribed JSON")]
-)
-at = gr.Interface(
-    fn=AudioTranscribe,
-    inputs=[gr.Number(label="Number of Speakers", placeholder="2"), gr.Textbox(label="Name of the Speakers (ordered by the time they speak and separated by comma)", placeholder="If Speaker 1 is first to speak followed by Speaker 2 then -> Speaker 1, Speaker 2"), 'audio'],
-    outputs=[gr.Textbox(label="Transcribed Text", lines=15), gr.JSON(label="Transcribed JSON")]
-)
 
-# demo = gr.TabbedInterface([ut, vt, at], ["Youtube URL", "Video", "Audio"])
-# demo.launch()
 with gr.Blocks() as yav_ui:
     with gr.Row():
         with gr.Column():
             with gr.Tab("Youtube", id=1):
+                ysz = gr.Dropdown(label="Model Size", choices=wispher_models , value='base')
                 yinput_nos = gr.Number(label="Number of Speakers", placeholder="2")
                 yinput_sn = gr.Textbox(label="Name of the Speakers (ordered by the time they speak and separated by comma)", placeholder="If Speaker 1 is first to speak followed by Speaker 2 then -> Speaker 1, Speaker 2")
                 yinput = gr.Textbox(label="Youtube Link", placeholder="https://www.youtube.com/watch?v=GECcjrYHH8w")
                 ybutton_transcribe = gr.Button("Transcribe", show_progress=True, scroll_to_output=True)
             with gr.Tab("Video", id=2):
+                vsz = gr.Dropdown(label="Model Size", choices=wispher_models, value='base')
                 vinput_nos = gr.Number(label="Number of Speakers", placeholder="2")
                 vinput_sn = gr.Textbox(label="Name of the Speakers (ordered by the time they speak and separated by comma)", placeholder="If Speaker 1 is first to speak followed by Speaker 2 then -> Speaker 1, Speaker 2")
                 vinput = gr.Video(label="Video")
                 vbutton_transcribe = gr.Button("Transcribe", show_progress=True, scroll_to_output=True)
             with gr.Tab("Audio", id=3):
+                asz = gr.Dropdown(label="Model Size", choices=wispher_models , value='base')
                 ainput_nos = gr.Number(label="Number of Speakers", placeholder="2")
                 ainput_sn = gr.Textbox(label="Name of the Speakers (ordered by the time they speak and separated by comma)", placeholder="If Speaker 1 is first to speak followed by Speaker 2 then -> Speaker 1, Speaker 2")
                 ainput = gr.Audio(label="Audio", type="filepath")
@@ -352,17 +372,17 @@ with gr.Blocks() as yav_ui:
                 output_json = gr.JSON(label="Transcribed JSON")
     ybutton_transcribe.click(
                 fn=YoutubeTranscribe,
-                inputs=[yinput_nos,yinput_sn,yinput],
+                inputs=[yinput_nos,yinput_sn,yinput, ysz],
                 outputs=[output_textbox,output_json]
             )
     abutton_transcribe.click(
                 fn=AudioTranscribe,
-                inputs=[ainput_nos,ainput_sn,ainput],
+                inputs=[ainput_nos,ainput_sn,ainput, asz],
                 outputs=[output_textbox,output_json]
             )
     vbutton_transcribe.click(
                 fn=VideoTranscribe,
-                inputs=[vinput_nos,vinput_sn,vinput],
+                inputs=[vinput_nos,vinput_sn,vinput, vsz],
                 outputs=[output_textbox,output_json]
             )
 yav_ui.launch(debug=True)