improve translation speed

translation.py

@@ -1,3 +1,5 @@
+# -*- coding: utf-8 -*-
+
 import os
 import re
 import sys
@@ -7,6 +9,9 @@ import torch
 import pysbd
 from transformers import NllbTokenizer, AutoModelForSeq2SeqLM
 import unicodedata
+import time
+
+
 
 #hy_segmenter = pysbd.Segmenter(language="hy", clean=False) not needed
 
@@ -117,8 +122,8 @@ class Translator:
         self.tokenizer = NllbTokenizer.from_pretrained(MODEL_NAME, token=HF_TOKEN)
         init_tokenizer(self.tokenizer)
 
-        self.hyw_splitter = pysbd.Segmenter(language="hy", clean=False)
-        self.eng_splitter = pysbd.Segmenter(language="en", clean=False)
+        self.hyw_splitter = pysbd.Segmenter(language="hy", clean=True)
+        self.eng_splitter = pysbd.Segmenter(language="en", clean=True)
         self.languages = LANGUAGES
     
 
@@ -138,6 +143,7 @@ class Translator:
         )
         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],
@@ -155,36 +161,51 @@ class Translator:
     def translate(self, text: str,
                   src_lang: str,
                   tgt_lang: str,
-                  max_length="auto",
+                  max_length=256,
                   num_beams=4,
                   by_sentence=True,
                   clean=True,
                   **kwargs):
         
         if by_sentence:
-            if src_lang =="eng_Latn":
-                sents, fillers = sentenize_with_fillers(text, self.eng_splitter, ignore_errors=True)
+            if src_lang == "eng_Latn":
+                sents = self.eng_splitter.segment(text)
             elif src_lang == "hyw_Armn":
-                sents, fillers = sentenize_with_fillers(text, self.hyw_splitter, ignore_errors=True)
+                sents = self.hyw_splitter.segment(text)
 
-        else:
-            sents = [text]
-            fillers = ["", ""]
 
         if clean:
             sents = [clean_text(sent, src_lang) for sent in sents]
 
-        results = []
-        for sent, sep in zip(sents, fillers):
-            results.append(sep)
-            results.append(self.translate_single(sent, src_lang, tgt_lang, max_length, num_beams, **kwargs))
-        
-        results.append(fillers[-1])
+
+        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.")
-    print(translator.translate("Hello, world!", "eng_Latn", "hyw_Armn"))
+
+    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)