version 0.1

src/anki_hsk_creator/utility.py

@@ -0,0 +1,154 @@
+"""utility.py
+
+
+Static clasess and functions for general use
+"""
+
+# Standard Library
+from pathlib import Path
+
+# Pip
+import argostranslate.package
+import argostranslate.translate
+import torch
+from cedict_utils.cedict import CedictParser
+from chatterbox.mtl_tts import ChatterboxMultilingualTTS
+
+# Local
+from .constants import CCCEDICT_PATH, INPUT, LANGUAGES, OUTPUT, RESOURCES
+
+# Static Clases
+
+
+class TRANS:
+    """Static Class for Argos translate"""
+
+    UPDATED = False
+    PACKAGES = None
+
+    @staticmethod
+    def create_translator(from_code, to_code):
+        """Download and install Argos Translate package"""
+        if not TRANS.UPDATED:
+            argostranslate.package.update_package_index()
+            TRANS.PACKAGES = argostranslate.package.get_available_packages()
+            TRANS.UPDATED = True
+        package_to_install = next(
+            filter(
+                lambda x: x.from_code == from_code and x.to_code == to_code,
+                TRANS.PACKAGES,
+            )
+        )
+        argostranslate.package.install_from_path(package_to_install.download())
+
+
+class CCCEDICT:
+    """Static Class for the CCCEDIT dictionary"""
+
+    PARSER = None
+    ENTRIES = []
+    DICTIONARY_LIST = {}
+
+    @staticmethod
+    def create_cedict(language_id=LANGUAGES.EN):
+        """Creates a create_cedict dictionary object"""
+        if not CCCEDICT.PARSER:
+            CCCEDICT.PARSER = CedictParser()
+            CCCEDICT.PARSER.read_file(CCCEDICT_PATH)
+            CCCEDICT.ENTRIES = CCCEDICT.PARSER.parse()
+        if language_id not in CCCEDICT.DICTIONARY_LIST:
+            dictionary = {}
+            for entry in CCCEDICT.ENTRIES:
+                if language_id != LANGUAGES.EN:
+                    TRANS.create_translator(LANGUAGES.EN, language_id)
+                    entry = argostranslate.translate.translate(
+                        entry, LANGUAGES.EN, language_id
+                    )
+                if entry.simplified not in dictionary:
+                    dictionary[entry.simplified] = [entry]
+                else:
+                    dictionary[entry.simplified].append(entry)
+            CCCEDICT.DICTIONARY_LIST[language_id] = dictionary
+        else:
+            dictionary = CCCEDICT.DICTIONARY_LIST[language_id]
+        return dictionary
+
+
+class TTS:
+    """Static class for the the TTS engine"""
+
+    MODEL = None
+    DEVICE = None
+
+    @staticmethod
+    def create_tts():
+        """Creates a TTS engine"""
+        if TTS.DEVICE is None:
+            # Automatically detect the best available device
+            if torch.cuda.is_available():
+                TTS.DEVICE = "cuda"
+            elif torch.backends.mps.is_available():
+                TTS.DEVICE = "mps"
+            else:
+                TTS.DEVICE = "cpu"
+        if TTS.MODEL is None:
+            TTS.MODEL = ChatterboxMultilingualTTS.from_pretrained(
+                device=TTS.DEVICE, t3_model="v3"
+            )
+
+
+# Clases
+
+
+class ProcessFile:
+    """Class that represents a file to processs
+
+    diferent input files has direfent process_files depending on language
+    """
+
+    def __init__(self, input_file: Path, language_id: str = None):
+        self.input_file = input_file
+        self._language_id = language_id
+        # process file type
+        self.out_folder = OUTPUT / input_file.parent
+        self.out_folder.mkdir(parents=True, exist_ok=True)
+        resources = RESOURCES / input_file
+        self.resources = resources.parent / resources.stem
+        self.resources.mkdir(parents=True, exist_ok=True)
+
+    @property
+    def absolute_input_file(self):
+        """Absolute input file"""
+        return INPUT / self.input_file
+
+    @property
+    def language_id(self):
+        """language for this trasnlation process"""
+        return self._language_id
+
+    @language_id.setter
+    def language_id(self, value):
+        self._language_id = value
+
+    @property
+    def output_name(self):
+        """Posible name for the output file, still missing the filetype"""
+        if self.language_id is None:
+            raise ValueError("Not a valid language selected")
+        return self.input_file.parent / f"{self.input_file.stem}.{self.language_id})."
+
+
+class TranslationResult:
+    """Result of a translated process"""
+
+    def __init__(
+        self,
+        language_id: str,
+        translated: str,
+        line: str,
+        audio_path: Path,
+    ):
+        self.language_id = language_id
+        self.translated = translated
+        self.line = line
+        self.audio_path = audio_path