fix: support large input texts with chunking

config.json

@@ -3,7 +3,8 @@
     "__comment::idle_polling_interval": "The interval in seconds to check for new messages when the app has no tasks",
     "__comment::tokenizer_file": "The tokenizer file name inside the model directory (loader.model_path)",
     "__comment::loader": "CTranslate2 loader options, see https://opennmt.net/CTranslate2/python/ctranslate2.Translator.html#ctranslate2.Translator.__init__. Use 'model_path' key for local paths or 'model_name' key for models hosted on Hugging Face. Both can't be used at the same time.",
-    "__comment::inference": "CTranslate2 inference options, see the kwargs in https://opennmt.net/CTranslate2/python/ctranslate2.Translator.html#ctranslate2.Translator.translate_batch",
+    "__comment::inference": "CTranslate2 inference options, see the kwargs in https://opennmt.net/CTranslate2/python/ctranslate2.Translator.html#ctranslate2.Translator.translate_iterable.",
+    "__comment::chunking": "Text chunking options, including the token threshold and chunk size.",
     "__comment::changes_to_the_config": "the program needs to be restarted if you change this file since it is stored in memory on startup",
     "log_level": 20,
     "idle_polling_interval": 5,
@@ -16,10 +17,14 @@
     },
     "inference": {
         "max_batch_size": 8192,
+        "max_decoding_length": 256,
         "sampling_temperature": 0.0001,
         "disable_unk": true,
-        "max_decoding_length": 10000,
         "repetition_penalty": 1.3,
         "patience": 1.5
+    },
+    "chunking": {
+        "chunk_threshold": 256,
+        "chunk_size": 80
     }
 }

lib/Service.py

@@ -7,6 +7,7 @@
 import json
 import logging
 import os
+import re
 from copy import deepcopy
 from time import perf_counter
 from typing import TypedDict
@@ -18,6 +19,25 @@
 
 logger = logging.getLogger(os.environ["APP_ID"] + __name__)
 
+# Languages that do not use spaces between words
+_NO_SPACE_LANGUAGES = {"zh", "yue", "ja", "th", "my", "km", "lo", "bo", "dz", "shn"}
+
+
+def _is_no_space_text(text: str) -> bool:
+    """Return True when the source text appears to use a no-space writing system.
+
+    Since the origin language is always "detect_language",
+    we inspect the text itself. A whitespace ratio below 5% is
+    treated as no-space text, which helps identify languages
+    such as Chinese, Japanese, Thai, and few more.
+    """
+    stripped = text.strip()
+    if not stripped:
+        return False
+    space_count = stripped.count(" ") + stripped.count("\t") + stripped.count("\n")
+    return (space_count / len(stripped)) < 0.05
+
+
 class ServiceException(Exception):
     pass
 
@@ -78,26 +98,112 @@ def load_model(self):
         except Exception as e:
             raise ServiceException("Error loading the translation model") from e
 
+    def _chunk_text(self, text: str, max_units: int, is_no_space: bool = False) -> list[str]:
+        """Split text into sentence-boundary chunks of a maximum size.
+
+        Space-delimited text is split by words, while no-space text is split by
+        characters. Sentence boundaries are preserved where possible, using
+        standard punctuation for space-delimited text and
+        CJK (Chinese, Japanese, and Korean-alike languages) punctuation for no-space text.
+        """
+        # Keep sentence punctuation attached to the preceding sentence.
+        # For no-space text (CJK etc.) use `\s*` because sentences run together without
+        # whitespace. For all other text use `\s+`.
+        sentences = (
+            re.split(r"(?<=[\u3002\uff01\uff1f])\s*", text)
+            if is_no_space
+            else re.split(r"(?<=[.!?])\s+", text)
+        )
+
+        chunks: list[str] = []
+        current_parts: list[str] = []
+        current_count = 0
+        sep = "" if is_no_space else " "
+
+        for sentence in sentences:
+            unit_count = len(sentence) if is_no_space else len(sentence.split())
+            if unit_count == 0:
+                continue
+
+            if current_count + unit_count > max_units and current_parts:
+                chunks.append(sep.join(current_parts))
+                current_parts = []
+                current_count = 0
+
+            if unit_count > max_units:
+                if is_no_space:
+                    for i in range(0, len(sentence), max_units):
+                        chunks.append(sentence[i:i + max_units])
+                else:
+                    words = sentence.split()
+                    for i in range(0, len(words), max_units):
+                        chunks.append(" ".join(words[i:i + max_units]))
+                continue
+
+            current_parts.append(sentence)
+            current_count += unit_count
+
+        if current_parts:
+            chunks.append(sep.join(current_parts))
+
+        return chunks if chunks else [text]
+
+    def _join_chunks(self, chunks: list[str], target_language: str) -> str:
+        """Join translated chunks respecting language-specific rules.
+
+        No-space languages are joined without a separator, while other languages
+        use a space. The translated chunks are already in the correct reading
+        order, so their order should not be reversed.
+        """
+        chunks = [c.strip() for c in chunks if c.strip()]
+        if not chunks:
+            return ""
+
+        target_base = target_language.split("_")[0].lower()
+
+        separator = "" if target_base in _NO_SPACE_LANGUAGES else " "
+        return separator.join(chunks)
+
     def translate(self, data: TranslateRequest) -> str:
         logger.debug(f"translating text to: {data['target_language']}")
 
         try:
             start = perf_counter()
-            input_tokens = self.tokenizer.Encode(
-                f"<2{data['target_language']}> {clean_text(data['input'])}",
-                out_type=str,
+            cleaned = clean_text(data["input"])
+            chunking = self.config.get("chunking", {})
+            chunk_threshold = chunking.get("chunk_threshold", 256)
+            chunk_size = chunking.get("chunk_size", 80)
+            is_no_space_source = _is_no_space_text(cleaned)
+            input_token_count = len(self.tokenizer.Encode(cleaned, out_type=str))
+            chunks = (
+                self._chunk_text(cleaned, chunk_size, is_no_space=is_no_space_source)
+                if input_token_count > chunk_threshold
+                else [cleaned]
             )
-            results = self.translator.translate_batch(
-                [input_tokens],
+
+            all_input_tokens = [
+                self.tokenizer.Encode(
+                    f"<2{data['target_language']}> {chunk}",
+                    out_type=str,
+                )
+                for chunk in chunks
+            ]
+
+            inference_config = dict(self.config["inference"])
+
+            results = list(self.translator.translate_iterable(
+                all_input_tokens,
                 batch_type="tokens",
-                **self.config["inference"],
-            )
+                **inference_config,
+            ))
 
-            if len(results) == 0 or len(results[0].hypotheses) == 0:
+            if len(results) != len(chunks) or any(len(r.hypotheses) == 0 for r in results):
                 raise ServiceException("Empty result returned from translator")
 
             # todo: handle multiple hypotheses
-            translation = self.tokenizer.Decode(results[0].hypotheses[0])
+            translated_chunks = [self.tokenizer.Decode(r.hypotheses[0]) for r in results]
+
+            translation = self._join_chunks(translated_chunks, data["target_language"])
             elapsed = perf_counter() - start
             logger.info(f"time taken: {elapsed:.2f}s")
         except Exception as e: