358 lines
11 KiB
Python
358 lines
11 KiB
Python
"""utility.py
|
|
|
|
|
|
Static clasess and functions for general use
|
|
"""
|
|
|
|
# Standard Library
|
|
import random
|
|
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,
|
|
DICTATION_TYPE,
|
|
DICT_TYPE,
|
|
GWEN_TTS,
|
|
INPUT,
|
|
LANGUAGES,
|
|
OUTPUT,
|
|
PHRASES_TYPE,
|
|
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:
|
|
if in_package.to_code == to_code:
|
|
# Single package between 2 languages
|
|
print(
|
|
f"Installing package {in_package.from_code}"
|
|
f"->{in_package.to_code}"
|
|
)
|
|
packages_to_install.append(in_package)
|
|
break
|
|
for out_package in packages:
|
|
if out_package.to_code == to_code:
|
|
if in_package.to_code == out_package.from_code:
|
|
print(
|
|
f"Installing in_package {in_package.from_code}"
|
|
f"->{in_package.to_code}"
|
|
)
|
|
print(
|
|
f"Installing out_package {out_package.from_code}"
|
|
f"->{out_package.to_code}"
|
|
)
|
|
packages_to_install.append(in_package)
|
|
packages_to_install.append(out_package)
|
|
break
|
|
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",
|
|
"instruct": "语速缓慢而审慎，每个音节的语调都拿捏得恰到好处，宛如教授在指导新生。",
|
|
"instruct": "Speak in a slow pace and enuntiate every word, as a teacher to a learning student",
|
|
}
|
|
VOICES = ["Vivian", "Serena", "Uncle_Fu", "Dylan", "Eric"]
|
|
|
|
@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"""
|
|
print(f"starting to generate {text}")
|
|
speaker = random.choice(TTS.VOICES)
|
|
audio = TTS.MODEL.generate_custom_voice(
|
|
text=text, **TTS.DEFAULTS, speaker=speaker
|
|
)
|
|
print(f"finish to generate {text}")
|
|
return audio
|
|
|
|
# Clases
|
|
|
|
|
|
class ProcessFolder:
|
|
"""Class that represents a folder to processs
|
|
|
|
diferent input files has direfent process_files depending on language
|
|
"""
|
|
|
|
def __init__(self, input_folder: Path, language_id: str = None):
|
|
self.input_folder = input_folder
|
|
self._language_id = language_id
|
|
self.absolute_input_folder = INPUT / self.input_folder
|
|
# process file type
|
|
self.out_folder = OUTPUT / input_folder
|
|
self.out_folder.mkdir(parents=True, exist_ok=True)
|
|
self.resources = RESOURCES / input_folder
|
|
self.resources.mkdir(parents=True, exist_ok=True)
|
|
|
|
@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_folder.stem}.{self.language_id}.temp"
|
|
|
|
@property
|
|
def input_files(self):
|
|
input_files = []
|
|
for file in self.absolute_input_folder.glob(f"*.txt"):
|
|
for file_type in (DICT_TYPE, PHRASES_TYPE, DICTATION_TYPE):
|
|
if file_type in file.suffixes:
|
|
input_files.append(
|
|
ProcessFile(
|
|
file.relative_to(INPUT), language_id=self.language_id
|
|
)
|
|
)
|
|
return input_files
|
|
|
|
@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
|
|
|
|
def __str__(self):
|
|
return (
|
|
f"Proccess {self.input_folder}"
|
|
f"(out: {self.out_folder}, res: {self.resources})"
|
|
)
|
|
|
|
|
|
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 file_type(self):
|
|
if DICTATION_TYPE in self.input_file.suffixes:
|
|
return DICTATION_TYPE
|
|
elif DICT_TYPE in self.input_file.suffixes:
|
|
return DICT_TYPE
|
|
elif PHRASES_TYPE in self.input_file.suffixes:
|
|
return PHRASES_TYPE
|
|
else:
|
|
raise ValueError("File type not recognized")
|
|
|
|
@property
|
|
def absolute_input_file(self):
|
|
"""Absolute input file"""
|
|
return INPUT / self.input_file
|
|
|
|
@property
|
|
def language_id(self):
|
|
"""language for this trasnlation proccess"""
|
|
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")]
|
|
|
|
def __str__(self):
|
|
return f"Proccess:{self.input_file}"
|
|
|
|
|
|
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,
|
|
n: str,
|
|
language_id: str,
|
|
simplified: str,
|
|
traditional: str,
|
|
pinyin: str,
|
|
meaning: str,
|
|
audio_path: Path,
|
|
):
|
|
self.n = n
|
|
self.language_id = language_id
|
|
self.simplified = simplified
|
|
self.traditional = traditional
|
|
self.pinyin = pinyin
|
|
self.meaning = meaning
|
|
self.audio_path = audio_path
|