use sentence split for translation

app.py

@@ -1,5 +1,6 @@
 import gradio as gr
 import random
+import spacy
 import torch
 from transformers import MT5Tokenizer, MT5ForConditionalGeneration
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
@@ -11,32 +12,46 @@ translator.eval()
 summarizer.eval()
 translator.to(device)
 summarizer.to(device)
-
+nlp = spacy.load("en_core_web_sm")
 
 def generate_output(
     task,
     text,
 ):
-    inputs = tokenizer(
-        [text],
-        padding="longest",
-        max_length=1024,
-        truncation=True,
-        return_tensors="pt",
-    ).to(device)
     if task == 'Translation':
-        outputs = translator.generate(
-            **inputs,
-            max_new_tokens=256,
-        )
+        sentences = [sent.text.strip() for sent in nlp(text).sents] # List[spacy.tokens.span.Span]
+        gen_texts = []
+        for sentence in sentences:
+            inputs = tokenizer(
+                [sentence],
+                padding="longest",
+                max_length=1024,
+                truncation=True,
+                return_tensors="pt",
+            ).to(device)
+            outputs = translator.generate(
+                **inputs,
+                max_new_tokens=256,
+            )
+            gen_text_ = tokenizer.decode(outputs[0], skip_special_tokens=True)
+            gen_texts.append(gen_text_)
+        return " ".join(gen_texts)
+
     elif task == 'Summarization':
+        inputs = tokenizer(
+            [text],
+            padding="longest",
+            max_length=1024,
+            truncation=True,
+            return_tensors="pt",
+        ).to(device)
         outputs = summarizer.generate(
             **inputs,
             max_new_tokens=256,
         )
+        gen_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
     else:
         raise ValueError("task undefined!")
-    gen_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
     return gen_text
 
 TASKS = ["Translation", "Summarization"]

requirements.txt

@@ -1,3 +1,4 @@
 torch>=1.10
 transformers>=4.11.3
 sentencepiece
+spacy