Upload 4 files

modules/translation/nllb_inference.py

@@ -3,10 +3,10 @@ import gradio as gr
 import os
 
 from modules.utils.paths import TRANSLATION_OUTPUT_DIR, NLLB_MODELS_DIR
-from modules.translation.translation_base import TranslationBase
+import modules.translation.translation_base as base
 
 
-class NLLBInference(TranslationBase):
+class NLLBInference(base.TranslationBase):
     def __init__(self,
                  model_dir: str = NLLB_MODELS_DIR,
                  output_dir: str = TRANSLATION_OUTPUT_DIR

modules/translation/translation_base.py

@@ -7,10 +7,12 @@ from typing import List
 from datetime import datetime
 
 import modules.translation.nllb_inference as nllb
+from modules.whisper.data_classes import *
 from modules.utils.subtitle_manager import *
 from modules.utils.files_manager import load_yaml, save_yaml
 from modules.utils.paths import DEFAULT_PARAMETERS_CONFIG_PATH, NLLB_MODELS_DIR, TRANSLATION_OUTPUT_DIR
 
+
 class TranslationBase(ABC):
     def __init__(self,
                  model_dir: str = NLLB_MODELS_DIR,
@@ -127,64 +129,6 @@ 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 file: {e}")
-            raise
-        finally:
-            self.release_cuda_memory()
-
     def offload(self):
         """Offload the model and free up the memory"""
         if self.model is not None: