Update modules/translation/translation_base.py

modules/translation/translation_base.py

@@ -136,6 +136,56 @@ class TranslationBase(ABC):
         finally:
             self.release_cuda_memory()
 
+    def translate_text(self,
+                       input_list_dict: list,
+                       model_size: str,
+                       src_lang: str,
+                       tgt_lang: str,
+                       max_length: int = 200,
+                       add_timestamp: bool = True,
+                       progress=gr.Progress()) -> list:
+        """
+        Translate text from source language to target language
+        Parameters
+        ----------
+        str_text: str
+            List[dict] to translate
+        model_size: str
+            Whisper model size from gr.Dropdown()
+        src_lang: str
+            Source language of the file to translate from gr.Dropdown()
+        tgt_lang: str
+            Target language of the file to translate from gr.Dropdown()
+        max_length: int
+            Max length per line to translate
+        add_timestamp: bool
+            Boolean value that determines whether to add a timestamp
+        progress: gr.Progress
+            Indicator to show progress directly in gradio.
+            I use a forked version of whisper for this. To see more info : https://github.com/jhj0517/jhj0517-whisper/tree/add-progress-callback
+        Returns
+        ----------
+        A List of
+        List[dict] with translation
+        """
+        try:
+            self.cache_parameters(model_size=model_size,src_lang=src_lang,tgt_lang=tgt_lang,max_length=max_length,add_timestamp=add_timestamp)
+            self.update_model(model_size=model_size,src_lang=src_lang,tgt_lang=tgt_lang,progress=progress)
+            
+            total_progress = len(input_list_dict)
+            for index, dic in enumerate(input_list_dict):
+                progress(index / total_progress, desc="Translating..")
+                translated_text = self.translate(dic["text"], max_length=max_length)
+                dic["text"] = translated_text
+
+            return input_list_dict
+
+        except Exception as e:
+            print(f"Error translating text: {e}")
+            raise
+        finally:
+            self.release_cuda_memory()
+    
     @staticmethod
     def get_device():
         if torch.cuda.is_available():