separate transcription and diarization for longer records

handler.py

@@ -8,7 +8,7 @@ import torch
 #     stdout, stderr = process.communicate()
 
 import whisperx
-import os
+import os, gc
 
 import time
 import json
@@ -188,7 +188,7 @@ class EndpointHandler:
     def __init__(self, path=""):
         # load the model
         device, batch_size, compute_type, whisper_model = whisper_config()
-        # self.model = whisperx.load_model(whisper_model, device=device, compute_type=compute_type, language="fr")
+        self.model = whisperx.load_model(whisper_model, device=device, compute_type=compute_type, language="fr")
         # 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 ERIC
         # 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 OLD
         # logger.info(f"Model {whisper_model} initialized")
@@ -218,17 +218,23 @@ class EndpointHandler:
         logger.info(display_gpu_infos())
 
         # 1. process input
-        # for diarization without transcription, the transcription is given as input, so data is now a tuple (inputs, transcription)
-        inputs_encoded, transcription = data.pop("inputs", data)
-        # inputs_encoded = data.pop("inputs", data)
         parameters = data.pop("parameters", None)
         options = data.pop("options", None)
 
         # OPTIONS are given as parameters
-        info = True if options and "info" in options.keys() and options["info"] else False
-        alignment = True if options and "alignment" in options.keys() and options["alignment"] else False
-        diarization = False if options and "diarization" in options.keys() and not options["diarization"] else True
-        language = parameters["language"] if parameters and "language" in parameters.keys() else "fr"
+        info = options.get("info", False)
+        transcribe = options.get("transcription", False)
+        alignment = options.get("alignment", False)
+        diarization = options.get("diarization", False)
+        language = parameters.get("language", "fr")
+        min_speakers = parameters.get("min_speakers", 2)
+        max_speakers = parameters.get("max_speakers", 25)
+
+        # for diarization without transcription, the transcription is given as input, so data is now a tuple (inputs, transcription)
+        if transcribe:
+            (inputs_encoded,) = data.pop("inputs", data)
+        elif diarization:
+            inputs_encoded, transcription = data.pop("inputs", data)
 
         inputs = base64.b64decode(inputs_encoded)
         logger.info(f"inputs decoded.")
@@ -237,82 +243,79 @@ class EndpointHandler:
             w.write(inputs)
         logger.info(f"inputs saved.")
 
-        # audio_nparray = ffmpeg_load_audio("/tmp/myfile.tmp", sr=SAMPLE_RATE, mono=True, out_type=np.float32)
         audio_nparray = load_audio("/tmp/myfile.tmp", sr=SAMPLE_RATE)
         logger.info(f"inputs loaded as mono 16kHz.")
         # clean up
         os.remove("/tmp/myfile.tmp")
         logger.info(f"temp file removed.")
 
-        # audio_nparray = ffmpeg_read(inputs, SAMPLE_RATE)
-        # audio_tensor = torch.from_numpy(audio_nparray)
-        # logger.info(f"inputs loaded as mono 16kHz.")
-
-        # get the end time
         et = time.time()
-
-        # get the execution time
         elapsed_time = et - st
+
         logger.info(f"TIME for audio processing : {elapsed_time:.2f} seconds")
         if info:
             print(f"TIME for audio processing : {elapsed_time:.2f} seconds")
 
         # 2. transcribe
-        # logger.info("--------------- STARTING TRANSCRIPTION ------------------------")
-        # transcription = self.model.transcribe(audio_nparray, batch_size=batch_size, language=language)
-        # if info:
-        #     print(transcription["segments"][0:10000])  # before alignment
-        # logger.info(transcription["segments"][0:10000])
-
-        # try:
-        #     first_text = transcription["segments"][0]["text"]
-        # except:
-        #     logger.warning("No transcription")
-        #     return {"transcription": transcription["segments"]}
-
-        # # get the execution time
-        # et = time.time()
-        # elapsed_time = et - st
-        # st = time.time()
-        # logger.info(f"TIME for audio transcription : {elapsed_time:.2f} seconds")
-        # if info:
-        #     print(f"TIME for audio transcription : {elapsed_time:.2f} seconds")
-
-        # # 3. align
-        # if alignment:
-        #     logger.info("--------------- STARTING ALIGNMENT ------------------------")
-        #     model_a, metadata = whisperx.load_align_model(language_code=transcription["language"], device=device)
-        #     transcription = whisperx.align(
-        #         transcription["segments"], model_a, metadata, audio_nparray, device, return_char_alignments=False
-        #     )
-        #     if info:
-        #         print(transcription["segments"][0:10000])
-        #     logger.info(transcription["segments"][0:10000])
-
-        #     # get the execution time
-        #     et = time.time()
-        #     elapsed_time = et - st
-        #     st = time.time()
-        #     logger.info(f"TIME for alignment : {elapsed_time:.2f} seconds")
-        #     if info:
-        #         print(f"TIME for alignment : {elapsed_time:.2f} seconds")
+        if transcribe:
+            gc.collect()
+            torch.cuda.empty_cache()
+            logger.info("--------------- STARTING TRANSCRIPTION ------------------------")
+            transcription = self.model.transcribe(audio_nparray, batch_size=batch_size, language=language)
+            if info:
+                print(transcription["segments"][0:10_000])  # before alignment
+            else:
+                logger.info(transcription["segments"][0:1_000])
+
+            try:
+                first_text = transcription["segments"][0]["text"]
+            except:
+                logger.warning("No transcription")
+                return {"transcription": transcription["segments"]}
+
+            et = time.time()
+            elapsed_time = et - st
+            st = time.time()
+            logger.info(f"TIME for audio transcription : {elapsed_time:.2f} seconds")
+            if info:
+                print(f"TIME for audio transcription : {elapsed_time:.2f} seconds")
+
+        # 3. align
+        if alignment:
+            gc.collect()
+            torch.cuda.empty_cache()
+            logger.info("--------------- STARTING ALIGNMENT ------------------------")
+            model_a, metadata = whisperx.load_align_model(language_code=transcription["language"], device=device)
+            transcription = whisperx.align(
+                transcription["segments"], model_a, metadata, audio_nparray, device, return_char_alignments=False
+            )
+            del model_a
+            if info:
+                print(transcription["segments"][0:10000])
+            logger.info(transcription["segments"][0:10000])
+
+            et = time.time()
+            elapsed_time = et - st
+            st = time.time()
+            logger.info(f"TIME for alignment : {elapsed_time:.2f} seconds")
+            if info:
+                print(f"TIME for alignment : {elapsed_time:.2f} seconds")
 
         # 4. Assign speaker labels
         if diarization:
+            gc.collect()
+            torch.cuda.empty_cache()
             logger.info("--------------- STARTING DIARIZATION ------------------------")
+            if not transcription:
+                logger.warning("No transcription to diarize")
             # add min/max number of speakers if known
-            diarize_segments = self.diarize_model(audio_nparray)
+            diarize_segments = self.diarize_model(audio_nparray, min_speakers=min_speakers, max_speakers=max_speakers)
             if info:
                 print(diarize_segments)
             logger.info(diarize_segments)
-            # diarize_model(audio, min_speakers=min_speakers, max_speakers=max_speakers)
 
             transcription = whisperx.assign_word_speakers(diarize_segments, transcription)
-            # if info:
-            #     print(transcription["segments"][0:10000])
-            # logger.info(transcription["segments"][0:10000])  # segments are now assigned speaker IDs
 
-            # get the execution time
             et = time.time()
             elapsed_time = et - st
             st = time.time()