Add custom pipeline

pipeline.py

@@ -0,0 +1,43 @@
+from typing import Dict, Any
+from transformers import AutoModelForSeq2SeqLM, NllbTokenizer
+from translation import fix_tokenizer, TextPreprocessor, sentenize_with_fillers
+from sentence_splitter import SentenceSplitter
+import torch
+
+class PreTrainedPipeline():
+  def __init__(self, path=""):
+    self.model = AutoModelForSeq2SeqLM.from_pretrained(path)
+    if torch.cuda.is_available():
+      self.model = self.model.cuda()
+    self.tokenizer = NllbTokenizer.from_pretrained(path)
+    fix_tokenizer(self.tokenizer)
+    self.splitter = SentenceSplitter(language='es')
+    self.preprocessor = TextPreprocessor()
+
+  def __call__(self, data: Dict[str, Any]) -> Dict[str, Any]:
+    inputs = data.get("text", "")
+    src_lang = data.get("src_lang", "spa_Latn")
+    tgt_lang = data.get("tgt_lang", "agr_Latn")
+    preprocess = data.get("preprocess", True)
+    
+    sentences, fillers = sentenize_with_fillers(inputs, self.splitter)
+    if preprocess:
+      sentences = [self.preprocessor(sent) for sent in sentences]
+
+    translated_sentences  = []
+    for sentence in sentences:
+      self.tokenizer.src_lang = src_lang
+      encoded = self.tokenizer(sentence, return_tensors="pt")
+      generated_tokens = self.model.generate(
+          **encoded.to(self.model.device),
+          forced_bos_token_id=self.tokenizer.lang_code_to_id[tgt_lang]
+      )
+      translated_sentences.append(
+          self.tokenizer.batch_decode(generated_tokens, skip_special_tokens=True)[0]
+      )
+
+    output = "".join(
+        filler + sentence for filler, sentence in zip(fillers, translated_sentences)
+    ) + fillers[-1]
+    
+    return {"translation": output}

requirements.txt

@@ -0,0 +1,6 @@
+requests==2.27.*
+sentencepiece==0.1.*
+torch==1.11.*
+transformers==4.33.*
+sentence-splitter==1.4
+sacremoses== 0.0.45

translation.py

@@ -0,0 +1,93 @@
+import re
+import sys
+import typing as tp
+import unicodedata
+
+import torch
+from sacremoses import MosesPunctNormalizer
+from sentence_splitter import SentenceSplitter
+from transformers import AutoModelForSeq2SeqLM, NllbTokenizer
+
+L1 = "spa_Latn"
+L2 = "agr_Latn"
+LANGUAGES = {
+    "Spanish | spa": L1,
+    "Awajun | agr": L2,
+}
+
+def get_non_printing_char_replacer(replace_by: str = " ") -> tp.Callable[[str], 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
+
+class TextPreprocessor:
+    """
+    Mimic the text preprocessing made for the NLLB model.
+    This code is adapted from the Stopes repo of the NLLB team:
+    https://github.com/facebookresearch/stopes/blob/main/stopes/pipelines/monolingual/monolingual_line_processor.py#L214
+    """
+
+    def __init__(self, lang="en"):
+        self.mpn = MosesPunctNormalizer(lang=lang)
+        self.mpn.substitutions = [
+            (re.compile(r), sub) for r, sub in self.mpn.substitutions
+        ]
+        self.replace_nonprint = get_non_printing_char_replacer(" ")
+
+    def __call__(self, text: str) -> str:
+        clean = self.mpn.normalize(text)
+        clean = self.replace_nonprint(clean)
+        # replace 𝓕𝔯𝔞𝔫𝔠𝔢𝔰𝔠𝔞 by Francesca
+        clean = unicodedata.normalize("NFKC", clean)
+        return clean
+
+def fix_tokenizer(tokenizer, new_lang=L2):
+    """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 = {}
+
+def sentenize_with_fillers(text, splitter, fix_double_space=True, ignore_errors=False):
+    """Apply a sentence splitter and return the sentences and all separators before and after them"""
+    if fix_double_space:
+        text = re.sub(" +", " ", text)
+    sentences = splitter.split(text)
+    fillers = []
+    i = 0
+    for sentence in sentences:
+        start_idx = text.find(sentence, i)
+        if ignore_errors and start_idx == -1:
+            # print(f"sent not found after {i}: `{sentence}`")
+            start_idx = i + 1
+        assert start_idx != -1, f"sent not found after {i}: `{sentence}`"
+        fillers.append(text[i:start_idx])
+        i = start_idx + len(sentence)
+    fillers.append(text[i:])
+    return sentences, fillers