| import os |
| import py3langid as langid |
|
|
|
|
| from .common import OfflineTranslator |
|
|
| ISO_639_1_TO_MBart50 = { |
|
|
| 'ar': 'ar_AR', |
| 'de': 'de_DE', |
| 'en': 'en_XX', |
| 'es': 'es_XX', |
| 'fr': 'fr_XX', |
| 'hi': 'hi_IN', |
| 'it': 'it_IT', |
| 'ja': 'ja_XX', |
| 'nl': 'nl_XX', |
| 'pl': 'pl_PL', |
| 'pt': 'pt_XX', |
| 'ru': 'ru_RU', |
| 'sw': 'sw_KE', |
| 'th': 'th_TH', |
| 'tr': 'tr_TR', |
| 'ur': 'ur_PK', |
| 'vi': 'vi_VN', |
| 'zh': 'zh_CN', |
| |
| |
| } |
|
|
| class MBart50Translator(OfflineTranslator): |
| |
| |
| _LANGUAGE_CODE_MAP = { |
| "ARA": "ar_AR", |
| "DEU": "de_DE", |
| "ENG": "en_XX", |
| "ESP": "es_XX", |
| "FRA": "fr_XX", |
| "HIN": "hi_IN", |
| "ITA": "it_IT", |
| "JPN": "ja_XX", |
| "NLD": "nl_XX", |
| "PLK": "pl_PL", |
| "PTB": "pt_XX", |
| "RUS": "ru_RU", |
| "SWA": "sw_KE", |
| "THA": "th_TH", |
| "TRK": "tr_TR", |
| "URD": "ur_PK", |
| "VIN": "vi_VN", |
| "CHS": "zh_CN", |
| } |
| |
| _MODEL_SUB_DIR = os.path.join(OfflineTranslator._MODEL_DIR, OfflineTranslator._MODEL_SUB_DIR, 'mbart50') |
| |
| _TRANSLATOR_MODEL = "facebook/mbart-large-50-many-to-many-mmt" |
|
|
|
|
|
|
| async def _load(self, from_lang: str, to_lang: str, device: str): |
| from transformers import ( |
| MBartForConditionalGeneration, |
| AutoTokenizer, |
| ) |
| if ':' not in device: |
| device += ':0' |
| self.device = device |
| self.model = MBartForConditionalGeneration.from_pretrained(self._TRANSLATOR_MODEL) |
| if self.device != 'cpu': |
| self.model.to(self.device) |
| self.model.eval() |
| self.tokenizer = AutoTokenizer.from_pretrained(self._TRANSLATOR_MODEL) |
|
|
| async def _unload(self): |
| del self.model |
| del self.tokenizer |
|
|
| async def _infer(self, from_lang: str, to_lang: str, queries: list[str]) -> list[str]: |
| if from_lang == 'auto': |
| detected_lang = langid.classify('\n'.join(queries))[0] |
| target_lang = self._map_detected_lang_to_translator(detected_lang) |
|
|
| if target_lang == None: |
| self.logger.warn('Could not detect language from over all sentence. Will try per sentence.') |
| else: |
| from_lang = target_lang |
|
|
| return [self._translate_sentence(from_lang, to_lang, query) for query in queries] |
|
|
| def _translate_sentence(self, from_lang: str, to_lang: str, query: str) -> str: |
|
|
| if not self.is_loaded(): |
| return '' |
|
|
| if from_lang == 'auto': |
| detected_lang = langid.classify(query)[0] |
| from_lang = self._map_detected_lang_to_translator(detected_lang) |
|
|
| if from_lang == None: |
| self.logger.warn(f'MBart50 Translation Failed. Could not detect language (Or language not supported for text: {query})') |
| return '' |
| |
| self.tokenizer.src_lang = from_lang |
| tokens = self.tokenizer(query, return_tensors="pt") |
| |
| if self.device != 'cpu': |
| tokens = tokens.to(self.device) |
| generated_tokens = self.model.generate(**tokens, forced_bos_token_id=self.tokenizer.lang_code_to_id[to_lang]) |
| result = self.tokenizer.batch_decode(generated_tokens, skip_special_tokens=True)[0] |
| return result |
|
|
| def _map_detected_lang_to_translator(self, lang): |
| if lang not in ISO_639_1_TO_MBart50: |
| return None |
|
|
| return ISO_639_1_TO_MBart50[lang] |
|
|
| async def _download(self): |
| import huggingface_hub |
| |
| huggingface_hub.snapshot_download(self._TRANSLATOR_MODEL, cache_dir=self._MODEL_SUB_DIR, ignore_patterns=["*.msgpack", "*.h5", '*.ot',".*", "*.safetensors"]) |
|
|
| def _check_downloaded(self) -> bool: |
| import huggingface_hub |
| return huggingface_hub.try_to_load_from_cache(self._TRANSLATOR_MODEL, 'pytorch_model.bin', cache_dir=self._MODEL_SUB_DIR) is not None |
