diff --git a/config.json b/config.json index d82d550..2aad271 100644 --- a/config.json +++ b/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 } } diff --git a/lib/Service.py b/lib/Service.py index 9cc4026..681e650 100644 --- a/lib/Service.py +++ b/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: