88 lines
3.2 KiB
Python
88 lines
3.2 KiB
Python
import warnings
|
|
from typing import List, Optional
|
|
|
|
import langdetect
|
|
from transformers import MarianMTModel, MarianTokenizer
|
|
|
|
from unstructured.nlp.tokenize import sent_tokenize
|
|
from unstructured.staging.huggingface import chunk_by_attention_window
|
|
|
|
|
|
def _get_opus_mt_model_name(source_lang: str, target_lang: str):
|
|
"""Constructs the name of the MarianMT machine translation model based on the
|
|
source and target language."""
|
|
return f"Helsinki-NLP/opus-mt-{source_lang}-{target_lang}"
|
|
|
|
|
|
def _validate_language_code(language_code: str):
|
|
if not isinstance(language_code, str) or len(language_code) != 2:
|
|
raise ValueError(
|
|
f"Invalid language code: {language_code}. Language codes must be two letter strings.",
|
|
)
|
|
|
|
|
|
def translate_text(text: str, source_lang: Optional[str] = None, target_lang: str = "en") -> str:
|
|
"""Translates the foreign language text. If the source language is not specified, the
|
|
function will attempt to detect it using langdetect.
|
|
|
|
Parameters
|
|
----------
|
|
text: str
|
|
The text to translate
|
|
target_lang: str
|
|
The two letter language code for the target langague. Defaults to "en".
|
|
source_lang: Optional[str]
|
|
The two letter language code for the language of the input text. If source_lang is
|
|
not provided, the function will try to detect it.
|
|
"""
|
|
if text.strip() == "":
|
|
return text
|
|
|
|
_source_lang: str = source_lang if source_lang is not None else langdetect.detect(text)
|
|
# NOTE(robinson) - Chinese gets detected with codes zh-cn, zh-tw, zh-hk for various
|
|
# Chinese variants. We normalizes these because there is a single model for Chinese
|
|
# machine translation
|
|
if _source_lang.startswith("zh"):
|
|
_source_lang = "zh"
|
|
|
|
_validate_language_code(target_lang)
|
|
_validate_language_code(_source_lang)
|
|
|
|
if target_lang == _source_lang:
|
|
return text
|
|
|
|
model_name = _get_opus_mt_model_name(_source_lang, target_lang)
|
|
print(f"Using model: {model_name}")
|
|
|
|
try:
|
|
tokenizer = MarianTokenizer.from_pretrained(model_name)
|
|
model = MarianMTModel.from_pretrained(model_name)
|
|
except OSError:
|
|
raise ValueError(
|
|
f"Transformers could not find the translation model {model_name}. "
|
|
"The requested source/target language combo is not supported.",
|
|
)
|
|
|
|
chunks: List[str] = chunk_by_attention_window(text, tokenizer, split_function=sent_tokenize)
|
|
|
|
translated_chunks: List[str] = []
|
|
for chunk in chunks:
|
|
translated_chunks.append(_translate_text(text, model, tokenizer))
|
|
|
|
return " ".join(translated_chunks)
|
|
|
|
|
|
def _translate_text(text, model, tokenizer):
|
|
"""Translates text using the specified model and tokenizer."""
|
|
# NOTE(robinson) - Suppresses the HuggingFace UserWarning resulting from the "max_length"
|
|
# key in the MarianMT config. The warning states that "max_length" will be deprecated
|
|
# in transformers v5
|
|
with warnings.catch_warnings():
|
|
warnings.simplefilter("ignore")
|
|
translated = model.generate(
|
|
**tokenizer([text], return_tensors="pt", padding=True, truncation=True),
|
|
)
|
|
return [tokenizer.decode(t, max_new_tokens=512, skip_special_tokens=True) for t in translated][
|
|
0
|
|
]
|