split text to batches

app.py

@@ -1,4 +1,6 @@
 import gradio as gr
+import nltk
+nltk.download('punkt')
 from lang_list import (
     LANGUAGE_NAME_TO_CODE,
     T2TT_TARGET_LANGUAGE_NAMES,
@@ -15,6 +17,19 @@ processor = AutoProcessor.from_pretrained("facebook/hf-seamless-m4t-medium")
 # translated_text_from_text = processor.decode(output_tokens[0].tolist(), skip_special_tokens=True)
 # print(translated_text_from_text)
 
+def split_text_into_batches(text, max_tokens_per_batch):
+    sentences = nltk.sent_tokenize(text)  # Tokenize text into sentences
+    batches = []
+    current_batch = ""
+    for sentence in sentences:
+        if len(current_batch) + len(sentence) + 1 <= max_tokens_per_batch:  # Add 1 for space
+            current_batch += sentence + " "  # Add sentence to current batch
+        else:
+            batches.append(current_batch.strip())  # Add current batch to batches list
+            current_batch = sentence + " "  # Start a new batch with the current sentence
+    if current_batch:
+        batches.append(current_batch.strip())  # Add the last batch
+    return batches
 
 def run_t2tt(file_uploader , input_text: str, source_language: str, target_language: str) -> (str, bytes):
     if file_uploader is not None:
@@ -22,9 +37,15 @@ def run_t2tt(file_uploader , input_text: str, source_language: str, target_langu
             input_text=file.read()
     source_language_code = LANGUAGE_NAME_TO_CODE[source_language]
     target_language_code = LANGUAGE_NAME_TO_CODE[target_language]
-    text_inputs = processor(text = input_text, src_lang=source_language_code , return_tensors="pt")
-    output_tokens = model.generate(**text_inputs, tgt_lang=target_language_code)
-    output = processor.decode(output_tokens[0].tolist(), skip_special_tokens=True)
+    max_tokens_per_batch= 256
+    batches = split_text_into_batches(input_text, max_tokens_per_batch)
+    translated_text = ""
+    for batch in batches:
+        text_inputs = processor(text=batch, src_lang=source_language_code, return_tensors="pt")
+        output_tokens = model.generate(**text_inputs, tgt_lang=target_language_code)
+        translated_batch = processor.decode(output_tokens[0].tolist(), skip_special_tokens=True)
+        translated_text += translated_batch + " "
+    output=translated_text.strip()
     _output_name = "result.txt"
     open(_output_name, 'w').write(output)
     return str(output), _output_name