unstructured-io--unstructured/unstructured/cleaners/translate.py

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
    ]