update

handler.py

@@ -192,9 +192,6 @@ class EndpointHandler():
             Return:
                 A :obj:`dict`:. base64 encoded image
             """
-            for x in data.keys():
-                logger.info(f"key: {x}, value: {data[x]}    ")
-                print(f"key: {x}, value: {data[x]}    ")
             
             logger.info("--------------- CUDA ------------------------")
             logger.info(display_gpu_infos())
@@ -202,6 +199,16 @@ class EndpointHandler():
             # 1. process input
             inputs_encoded = data.pop("inputs", data)
             parameters = data.pop("parameters", None)
+            options = data.pop("options", None)
+            
+            # OPTIONS
+            info = False
+            if options and "info" in options.keys():
+                info = True
+
+            alignment = False
+            if options and "alignment" in options.keys():
+                info = True
         
             language = "fr"
             if parameters and "language" in parameters.keys():
@@ -220,35 +227,38 @@ class EndpointHandler():
             # audio_nparray = ffmpeg_read(inputs, SAMPLE_RATE)
             # audio_tensor= torch.from_numpy(audio_nparray)
             
-            results = []
-
             # 2. transcribe
             device, batch_size, compute_type, whisper_model = whisper_config()
             logger.info("--------------- STARTING TRANSCRIPTION ------------------------")
             transcription = self.model.transcribe(audio_nparray, batch_size=batch_size,language=language)
-            results.append({"transcription": transcription["segments"]})
-            
+            if info:
+                print(transcription["segments"]) # before alignment
             logger.info(transcription["segments"])
             
             # 3. align
-            logger.info("--------------- STARTING ALIGNMENT ------------------------")
-            # model_a, metadata = whisperx.load_align_model(
-            #     language_code=result["language"], device=device)
-            # transcription = whisperx.align(
-            #     result["segments"], model_a, metadata, audio_nparray, device, return_char_alignments=False)
-            # results.append({"aligned_transcription": transcription["segments"]})
-            # print(transcription["segments"])
+            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"])
+                logger.info(transcription["segments"])
             
             # 4. Assign speaker labels
             logger.info("--------------- STARTING DIARIZATION ------------------------")
             # add min/max number of speakers if known
-            #diarize_segments = self.diarize_model(audio_nparray)
-            #logger.info(diarize_segments)
+            diarize_segments = self.diarize_model(audio_nparray)
+            if info:
+                print(diarize_segments)
+            logger.info(diarize_segments)
             # diarize_model(audio, min_speakers=min_speakers, max_speakers=max_speakers)
 
-            #diarized_transcription = whisperx.assign_word_speakers(diarize_segments, transcription)
-            #logger.info(diarized_transcription["segments"])  # segments are now assigned speaker IDs
-            #results.append({"diarized_transcription": diarized_transcription["segments"]})
+            transcription = whisperx.assign_word_speakers(diarize_segments, transcription)
+            if info:
+                print(transcription["segments"])
+            logger.info(transcription["segments"])  # segments are now assigned speaker IDs
 
             if torch.cuda.is_available():
                 logger.info("--------------- GPU ------------------------")
@@ -259,7 +269,7 @@ class EndpointHandler():
             
             # results_json = json.dumps(results)
             # return {"results": results_json}
-            return results
+            return {"transcription": transcription["segments"]}