Source code for langcheck.metrics.de._translation

from math import floor

from nltk.tokenize import sent_tokenize
from transformers.pipelines import pipeline




class Translate:
    """Translation class based on HuggingFace's translation pipeline."""

    def __init__(self, model_name: str) -> None:
        """
        Initialize the Translation class with given parameters.

        Args:
            model_name: The name of the model to use for translation
        """
        self._translation_pipeline = pipeline(
            "translation",
            model=model_name,
            tokenizer=model_name,
            truncation=True,
        )
        self._max_length = self._translation_pipeline.model.config.max_length

    def _translate(self, texts: str) -> str:
        """Translate the texts using the translation pipeline.
        It splits the texts into blocks and translates each block separately,
        avoiding problems with long texts.
        Args:
            texts: The texts to translate
        Returns:
            The translated texts
        """
        tokenization = self._translation_pipeline.tokenizer(
            texts, return_tensors="pt"
        )  # type: ignore
        if tokenization.input_ids.shape[1] > (self._max_length / 2):
            # Split the text into blocks, if it is too long
            # starting from 2 * num_tokens / max_length to be sure
            # NB: this comes from a few 100 tests, but it is not a science
            blocks = floor(
                2 * tokenization.input_ids.shape[1] / self._max_length
            )
            sentences = sent_tokenize(texts)
            # Split sentences into a number of blocks, e.g., 2 blocks = 2 groups
            len_block = floor(len(sentences) / blocks) + 1
            sentences_list = []
            for i in range(blocks):
                sentences_list.append(
                    sentences[i * len_block : (i + 1) * len_block]
                )
            text_list = [" ".join(sent) for sent in sentences_list]
        else:
            text_list = [texts]
        translated_texts = []
        for text in text_list:
            text_en = [
                str(d["translation_text"])  # type: ignore
                for d in self._translation_pipeline(text)  # type: ignore
            ]
            translated_texts.append(" ".join(text_en))
        text_translated_final = " ".join(translated_texts)
        return text_translated_final

    def __call__(self, text: str) -> str:
        """Translate the text using the translation pipeline.
        Args:
            text: The text to translate
        Returns:
            The translated text
        """
        return self._translate(text)