added text Batches

app.py

@@ -1,4 +1,6 @@
 import os
+import nltk
+nltk.download('punkt')
 from doctr.io import DocumentFile
 from doctr.models import ocr_predictor
 import gradio as gr
@@ -99,7 +101,19 @@ demo_ocr = gr.Interface(
 
 # demo_ocr.launch(debug=True)
 
-
+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:
@@ -107,9 +121,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

requirements.txt

@@ -8,4 +8,5 @@ transformers
 fairseq2==0.1 
 pydub
 yt-dlp
-sentencepiece
+sentencepiece
+nltk