Cache parameters

modules/translation/deepl_api.py

@@ -4,8 +4,9 @@ import os
 from datetime import datetime
 import gradio as gr
 
-from modules.utils.paths import TRANSLATION_OUTPUT_DIR
+from modules.utils.paths import TRANSLATION_OUTPUT_DIR, DEFAULT_PARAMETERS_CONFIG_PATH
 from modules.utils.subtitle_manager import *
+from modules.utils.files_manager import load_yaml, save_yaml
 
 """
 This is written with reference to the DeepL API documentation.
@@ -125,6 +126,12 @@ class DeepLAPI:
         String to return to gr.Textbox()
         Files to return to gr.Files()
         """
+        self.cache_parameters(
+            api_key=auth_key,
+            is_pro=is_pro,
+            source_lang=source_lang,
+            target_lang=target_lang
+        )
 
         files_info = {}
         for fileobj in fileobjs:
@@ -199,4 +206,18 @@ class DeepLAPI:
         }
         response = requests.post(url, headers=headers, data=data).json()
         time.sleep(self.api_interval)
-        return response["translations"]
+        return response["translations"]
+
+    @staticmethod
+    def cache_parameters(api_key: str,
+                         is_pro: bool,
+                         source_lang: str,
+                         target_lang: str):
+        cached_params = load_yaml(DEFAULT_PARAMETERS_CONFIG_PATH)
+        cached_params["translation"]["deepl"] = {
+            "api_key": api_key,
+            "is_pro": is_pro,
+            "source_lang": source_lang,
+            "target_lang": target_lang
+        }
+        save_yaml(cached_params, DEFAULT_PARAMETERS_CONFIG_PATH)

modules/translation/translation_base.py

@@ -7,6 +7,8 @@ from datetime import datetime
 
 from modules.whisper.whisper_parameter 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
 
 
 class TranslationBase(ABC):
@@ -75,6 +77,11 @@ class TranslationBase(ABC):
         Files to return to gr.Files()
         """
         try:
+            self.cache_parameters(model_size=model_size,
+                                  src_lang=src_lang,
+                                  tgt_lang=tgt_lang,
+                                  max_length=max_length)
+
             self.update_model(model_size=model_size,
                               src_lang=src_lang,
                               tgt_lang=tgt_lang,
@@ -149,3 +156,17 @@ class TranslationBase(ABC):
         for file_path in file_paths:
             if file_path and os.path.exists(file_path):
                 os.remove(file_path)
+
+    @staticmethod
+    def cache_parameters(model_size: str,
+                         src_lang: str,
+                         tgt_lang: str,
+                         max_length: int):
+        cached_params = load_yaml(DEFAULT_PARAMETERS_CONFIG_PATH)
+        cached_params["translation"]["nllb"] = {
+            "model_size": model_size,
+            "src_lang": src_lang,
+            "tgt_lang": tgt_lang,
+            "max_length": max_length
+        }
+        save_yaml(cached_params, DEFAULT_PARAMETERS_CONFIG_PATH)