Refactor translation function to support paragraph translation

app.py

@@ -34,26 +34,36 @@ def load_tokenizer(src_lang, tgt_lang):
 def translate(text: str, src_lang: str, tgt_lang: str):
     tokenizer = load_tokenizer(src_lang, tgt_lang)
 
-    sentences = nltk.sent_tokenize(text)
-    translated_sentences = []
-
-    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([input_tokens]).to(device),
-            forced_bos_token_id=tokenizer.lang_code_to_id[code_mapping[tgt_lang]],
-            max_length=len(input_tokens) + 50,
-            num_return_sequences=1,
-        )
-        translated_chunk = tokenizer.decode(
-            translated_chunk[0], skip_special_tokens=True
-        )
-        translated_sentences.append(translated_chunk)
-
-    translated_text = " ".join(translated_sentences)
-    return translated_text
+    paragraphs = text.split("\n")
+    translated_paragraphs = []
+
+    for paragraph in paragraphs:
+        sentences = nltk.sent_tokenize(paragraph)
+        translated_sentences = []
+
+        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([input_tokens]).to(device),
+                forced_bos_token_id=tokenizer.lang_code_to_id[code_mapping[tgt_lang]],
+                max_length=len(input_tokens) + 50,
+                num_return_sequences=1,
+            )
+            translated_chunk = tokenizer.decode(
+                translated_chunk[0], skip_special_tokens=True
+            )
+            translated_sentences.append(translated_chunk)
+
+        translated_paragraph = " ".join(translated_sentences)
+        translated_paragraphs.append(translated_paragraph)
+
+    return "\n".join(translated_paragraphs)
 
 
 description = """