Add nltk dependency and update translate function to handle multiple sentences

app.py

@@ -4,6 +4,9 @@ from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
 from flores import code_mapping
 import platform
 import torch
+import nltk
+
+nltk.download("punkt")
 
 device = "cpu" if platform.system() == "Darwin" else "cuda"
 MODEL_NAME = "facebook/nllb-200-3.3B"
@@ -28,34 +31,29 @@ def load_tokenizer(src_lang, tgt_lang):
 
 
 @spaces.GPU
-def translate(
-    text: str,
-    src_lang: str,
-    tgt_lang: str,
-    window_size: int = 800,
-    overlap_size: int = 200,
-):
+def translate(text: str, src_lang: str, tgt_lang: str):
     tokenizer = load_tokenizer(src_lang, tgt_lang)
 
-    input_tokens = (
-        tokenizer(text, return_tensors="pt").input_ids[0].cpu().numpy().tolist()
-    )
-    translated_chunks = []
+    sentences = nltk.sent_tokenize(text)
+    translated_sentences = []
 
-    for i in range(0, len(input_tokens), window_size - overlap_size):
-        window = input_tokens[i : i + window_size]
+    for sentence in sentences:
+        input_tokens = (
+            tokenizer(sentence, return_tensors="pt").input_ids[0].cpu().numpy().tolist()
+        )
         translated_chunk = model.generate(
-            input_ids=torch.tensor([window]).to(device),
+            input_ids=torch.tensor([input_tokens]).to(device),
             forced_bos_token_id=tokenizer.lang_code_to_id[code_mapping[tgt_lang]],
-            max_length=window_size,
+            max_length=len(input_tokens) + 50,
             num_return_sequences=1,
         )
         translated_chunk = tokenizer.decode(
             translated_chunk[0], skip_special_tokens=True
         )
-        translated_chunks.append(translated_chunk)
+        translated_sentences.append(translated_chunk)
 
-    return " ".join(translated_chunks)
+    translated_text = " ".join(translated_sentences)
+    return translated_text
 
 
 description = """

requirements.txt

@@ -2,4 +2,5 @@
 transformers
 torch
 gradio
-spaces
+spaces
+nltk