Update modules/translation/translation_base.py

modules/translation/translation_base.py

@@ -1,4 +1,5 @@
 import os
+import re
 import torch
 import gradio as gr
 from abc import ABC, abstractmethod
@@ -178,9 +179,12 @@ class TranslationBase(ABC):
                 for index, dic in enumerate(input_list_dict):
                     progress(index / total_progress, desc="Translating..")
                     translated_text = self.translate(dic["text"], max_length=max_length)
-                    
-                    print(dic["text"])
-                    print(translated_text)
+
+                    # Add speaker ID to translated sentence when diarization is enabled
+                    if re.search(': ',dic['text']) != None:
+                        speaker_id = ((dic['text']).split(":", 1)[0]).strip()
+                        if speaker_id != "":
+                            translated_text = speaker_id + ": " + translated_text
                     
                     dic["text"] = translated_text