translation.py

# -*- coding: utf-8 -*-

import os
import re
import sys
import typing as tp

import torch
import pysbd
from transformers import NllbTokenizer, AutoModelForSeq2SeqLM
import unicodedata
import time



#hy_segmenter = pysbd.Segmenter(language="hy", clean=False) not needed

MODEL_NAME = "AriNubar/nllb-200-distilled-600m-en-hyw"

LANGUAGES = {
    "Արեւմտահայերէն | Western Armenian": "hyw_Armn",
    "Անգլերէն | English": "eng_Latn",
}

HF_TOKEN = os.environ.get("HF_TOKEN")

def get_non_printing_char_replacer(replace_by: str = " "):
    non_printable_map = {
        ord(c): replace_by
        for c in (chr(i) for i in range(sys.maxunicode + 1))
        # same as \p{C} in perl
        # see https://www.unicode.org/reports/tr44/#General_Category_Values
        if unicodedata.category(c) in {"C", "Cc", "Cf", "Cs", "Co", "Cn"}
    }

    def replace_non_printing_char(line) -> str:
        return line.translate(non_printable_map)

    return replace_non_printing_char

def clean_text(text: str, lang) -> str:
    HYW_CHARS_TO_NORMALIZE = {
        "«": '"',
        "»": '"',
        "“": '"',
        "”": '"',
        "’": "'",
        "‘": "'",
        "–": "-",
        "—": "-",
        "ՙ": "'",
        "՚": "'",
    }    

    DOUBLE_CHARS_TO_NORMALIZE = {
        "Կ՛": "Կ'",
        "կ՛": "կ'",        
        "Չ՛": "Չ'",
        "չ՛": "չ'",
        "Մ՛": "Մ'",
        "մ՛": "մ'",
    }
    replace_nonprint = get_non_printing_char_replacer()

    text = replace_nonprint(text)
    # print(text)
    text = text.replace("\t", " ").replace("\n", " ").replace("\r", " ").replace(r"[^\x00-\x7F]+", " ").replace(r"\s+", " ")
    text = text.strip()

    if lang == "hyw_Armn":
        text = text.translate(str.maketrans(HYW_CHARS_TO_NORMALIZE))
        for k, v in DOUBLE_CHARS_TO_NORMALIZE.items():
            text = text.replace(k, v)
    
    return text

def sentenize_with_fillers(text, splitter, fix_double_space=True, ignore_errors=False):
    if fix_double_space:
        text = re.sub(r"\s+", " ", text)
    text = text.strip()

    sentences = splitter.segment(text)

    fillers = []
    i = 0

    for sent in sentences:
        start_idx = text.find(sent, i)
        if ignore_errors and start_idx == -1:
            start_idx = i + 1
        assert start_idx != -1, f"Sent not found after index {i} in text: {text}"

        fillers.append(text[i:start_idx])
        i = start_idx + len(sent)
    
    fillers.append(text[i:])
    return sentences, fillers

def init_tokenizer(tokenizer, new_lang='hyw_Armn'):
    """ Add a new language token to the tokenizer vocabulary (this should be done each time after its initialization) """
    old_len = len(tokenizer) - int(new_lang in tokenizer.added_tokens_encoder)
    tokenizer.lang_code_to_id[new_lang] = old_len-1
    tokenizer.id_to_lang_code[old_len-1] = new_lang
    # always move "mask" to the last position
    tokenizer.fairseq_tokens_to_ids["<mask>"] = len(tokenizer.sp_model) + len(tokenizer.lang_code_to_id) + tokenizer.fairseq_offset

    tokenizer.fairseq_tokens_to_ids.update(tokenizer.lang_code_to_id)
    tokenizer.fairseq_ids_to_tokens = {v: k for k, v in tokenizer.fairseq_tokens_to_ids.items()}
    if new_lang not in tokenizer._additional_special_tokens:
        tokenizer._additional_special_tokens.append(new_lang)
    # clear the added token encoder; otherwise a new token may end up there by mistake
    tokenizer.added_tokens_encoder = {}
    tokenizer.added_tokens_decoder = {}

class Translator:
    def __init__(self) -> None:
        self.model = AutoModelForSeq2SeqLM.from_pretrained(MODEL_NAME, token=HF_TOKEN)

        if torch.cuda.is_available():
            self.model = self.model.cuda()
        
        self.tokenizer = NllbTokenizer.from_pretrained(MODEL_NAME, token=HF_TOKEN)
        init_tokenizer(self.tokenizer)

        self.hyw_splitter = pysbd.Segmenter(language="hy", clean=True)
        self.eng_splitter = pysbd.Segmenter(language="en", clean=True)
        self.languages = LANGUAGES
    

    def translate_single(
        self,
        text,
        src_lang,
        tgt_lang,
        max_length="auto",
        num_beams=4,
        n_out=None,
        **kwargs,
    ):
        self.tokenizer.src_lang = src_lang
        encoded = self.tokenizer(
            text, return_tensors="pt", truncation=True, max_length=256
        )
        if max_length == "auto":
            max_length = int(32 + 2.0 * encoded.input_ids.shape[1])
        
        generated_tokens = self.model.generate(
            **encoded.to(self.model.device),
            forced_bos_token_id=self.tokenizer.lang_code_to_id[tgt_lang],
            max_length=max_length,
            num_beams=num_beams,
            num_return_sequences=n_out or 1,
            **kwargs,
        )
        out = self.tokenizer.batch_decode(generated_tokens, skip_special_tokens=True)
        if isinstance(text, str) and n_out is None:
            return out[0]
        return out


    def translate(self, text: str,
                  src_lang: str,
                  tgt_lang: str,
                  max_length=256,
                  num_beams=4,
                  by_sentence=True,
                  clean=True,
                  **kwargs):
        
        if by_sentence:
            if src_lang == "eng_Latn":
                sents = self.eng_splitter.segment(text)
            elif src_lang == "hyw_Armn":
                sents = self.hyw_splitter.segment(text)


        if clean:
            sents = [clean_text(sent, src_lang) for sent in sents]


        if len(sents) > 1:
            results = self.translate_batch(sents, src_lang, tgt_lang, num_beams=num_beams, max_length=max_length, **kwargs)
        else:
            results = self.translate_single(sents, src_lang, tgt_lang, max_length=max_length, num_beams=num_beams, **kwargs)

        return " ".join(results)
    
    def translate_batch(self, texts, src_lang, tgt_lang, num_beams=4, max_length=256, **kwargs):
        self.tokenizer.src_lang = src_lang

        if torch.cuda.is_available():
            inputs = self.tokenizer(texts, return_tensors="pt", max_length=max_length, padding=True, truncation=True).input_ids.to("cuda")
            translated_tokens = self.model.generate(inputs, num_beams=num_beams, forced_bos_token_id=self.tokenizer.lang_code_to_id[tgt_lang])
        else:
            inputs = self.tokenizer(texts, return_tensors="pt", max_length=max_length, padding=True, truncation=True)
            translated_tokens = self.model.generate(**inputs, num_beams=num_beams, forced_bos_token_id=self.tokenizer.lang_code_to_id[tgt_lang])
        return self.tokenizer.batch_decode(translated_tokens, skip_special_tokens=True)

if __name__ == "__main__":

    print("Initializing translator...")
    translator = Translator()
    print("Translator initialized.")

    start_time = time.time()
    print(translator.translate("Hello world!", "eng_Latn", "hyw_Armn"))
    print("Time elapsed: ", time.time() - start_time)
    
    start_time = time.time()
    print(translator.translate("I am the greatest translator! Do not fuck with me!", "eng_Latn", "hyw_Armn"))
    print("Time elapsed: ", time.time() - start_time)