264 lines
8.1 KiB
Python
264 lines
8.1 KiB
Python
"""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 CedictEntry, CedictParser
|
|
from qwen_tts import Qwen3TTSModel
|
|
|
|
# Local
|
|
from .constants import CCCEDICT_PATH, GWEN_TTS, INPUT, LANGUAGES, OUTPUT, RESOURCES
|
|
|
|
# Static Clases
|
|
|
|
|
|
class TRANS:
|
|
"""Static Class for Argos translate"""
|
|
|
|
UPDATED = False
|
|
PACKAGES = []
|
|
|
|
@staticmethod
|
|
def create_translator(from_code, to_code):
|
|
"""Download and install Argos Translate package"""
|
|
print(f"Create translator from {from_code} to {to_code}")
|
|
if from_code == to_code:
|
|
return
|
|
if not TRANS.UPDATED:
|
|
argostranslate.package.update_package_index()
|
|
TRANS.PACKAGES = argostranslate.package.get_available_packages()
|
|
TRANS.UPDATED = True
|
|
packages = tuple(
|
|
filter(
|
|
lambda x: x.from_code == from_code or x.to_code == to_code,
|
|
TRANS.PACKAGES,
|
|
)
|
|
)
|
|
print(f"available packages {packages[:5]}")
|
|
packages_to_install = []
|
|
for in_package in packages:
|
|
if in_package.from_code == from_code:
|
|
for out_package in packages:
|
|
if out_package.to_code == to_code:
|
|
if in_package.to_code == out_package.from_code:
|
|
print(
|
|
f"Check in_package {in_package.from_code}"
|
|
f"{in_package.to_code}"
|
|
)
|
|
print(
|
|
f"Check out_package {out_package.from_code}"
|
|
f"{out_package.to_code}"
|
|
)
|
|
packages_to_install.append(in_package)
|
|
packages_to_install.append(out_package)
|
|
for package in packages_to_install:
|
|
print(f"instaling package {package}")
|
|
argostranslate.package.install_from_path(package.download())
|
|
|
|
|
|
class TranslatedEntry:
|
|
"""Holder class for CCCEDIT entry translated to `language_id`"""
|
|
|
|
def __init__(self, entry: CedictEntry, language_id: str):
|
|
self.entry = entry
|
|
self.language_id = language_id
|
|
self._translated_meanings = []
|
|
|
|
@property
|
|
def simplified(self):
|
|
"""Entry simplified"""
|
|
return self.entry.simplified
|
|
|
|
@property
|
|
def traditional(self):
|
|
"""Entry traditional"""
|
|
return self.entry.traditional
|
|
|
|
@property
|
|
def pinyin(self):
|
|
"""Entry piying"""
|
|
return self.entry.pinyin
|
|
|
|
@property
|
|
def meanings(self):
|
|
"""Entry translated meaning list"""
|
|
for meaning in self.entry.meanings:
|
|
if self.language_id != LANGUAGES.EN:
|
|
print(f"translating from {LANGUAGES.EN} to {self.language_id}")
|
|
print(f"-> {meaning}")
|
|
trans_meaning = argostranslate.translate.translate(
|
|
meaning, LANGUAGES.EN, self.language_id
|
|
)
|
|
else:
|
|
trans_meaning = meaning
|
|
self._translated_meanings.append(trans_meaning)
|
|
return self._translated_meanings
|
|
|
|
|
|
class CCCEDICT:
|
|
"""Static Class for the CCCEDIT dictionary"""
|
|
|
|
PARSER = None
|
|
ENTRIES = []
|
|
DICTIONARY_LIST = {}
|
|
|
|
@staticmethod
|
|
def create_cedict(
|
|
language_id: str = LANGUAGES.EN,
|
|
) -> dict[str, list[TranslatedEntry]]:
|
|
"""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:
|
|
TRANS.create_translator(LANGUAGES.EN, language_id)
|
|
dictionary = {}
|
|
for entry in CCCEDICT.ENTRIES:
|
|
trans_entry = TranslatedEntry(entry, language_id)
|
|
if entry.simplified not in dictionary:
|
|
dictionary[entry.simplified] = [trans_entry]
|
|
else:
|
|
dictionary[entry.simplified].append(trans_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
|
|
DEFAULTS = {
|
|
"language": "Chinese",
|
|
"speaker": "Uncle_Fu",
|
|
"instruct": "语速缓慢而审慎，每个音节的语调都拿捏得恰到好处，宛如教授在指导新生。",
|
|
}
|
|
|
|
@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:0"
|
|
elif torch.backends.mps.is_available():
|
|
TTS.DEVICE = "mps"
|
|
else:
|
|
TTS.DEVICE = "cpu"
|
|
if TTS.MODEL is None:
|
|
TTS.MODEL = Qwen3TTSModel.from_pretrained(
|
|
GWEN_TTS,
|
|
device_map=TTS.DEVICE,
|
|
dtype=torch.bfloat16,
|
|
attn_implementation="flash_attention_2",
|
|
)
|
|
|
|
@staticmethod
|
|
def generate(text: str):
|
|
"""Generates a Waw using the defaulst values"""
|
|
return TTS.MODEL.generate_custom_voice(text=text, **TTS.DEFAULTS)
|
|
|
|
|
|
# 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.out_folder / f"{self.input_file.stem}.{self.language_id}.temp"
|
|
|
|
@property
|
|
def dictionary_resource_file(self):
|
|
"""The path for the resource tsv for dictionary files"""
|
|
return self.resources / f"dictionary.{self.language_id}.tsv"
|
|
|
|
@property
|
|
def relative_dictionary_resource_file(self):
|
|
"""The path for the resource tsv for dictionary files"""
|
|
path = self.resources / f"dictionary.{self.language_id}.tsv"
|
|
return path.relative_to(RESOURCES)
|
|
|
|
@property
|
|
def available_dictionary_languages(self):
|
|
"""for a Dictionary file loads the avaliable proceced languages"""
|
|
return [lan.suffixes[0][1:] for lan in self.resources.glob("dictionary.*.tsv")]
|
|
|
|
|
|
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
|
|
|
|
|
|
class DictionaryResult:
|
|
"""Result of a dictionaty process"""
|
|
|
|
def __init__(
|
|
self,
|
|
language_id: str,
|
|
simplified: str,
|
|
traditional: str,
|
|
pinyin: str,
|
|
meaning: str,
|
|
audio_path: Path,
|
|
):
|
|
self.language_id = language_id
|
|
self.simplified = simplified
|
|
self.traditional = traditional
|
|
self.pinyin = pinyin
|
|
self.meaning = meaning
|
|
self.audio_path = audio_path
|