diff --git a/packages/sie_server/models/google__madlad400-3b-mt.yaml b/packages/sie_server/models/google__madlad400-3b-mt.yaml index dc8f0cf86..fda799e93 100644 --- a/packages/sie_server/models/google__madlad400-3b-mt.yaml +++ b/packages/sie_server/models/google__madlad400-3b-mt.yaml @@ -36,6 +36,7 @@ profiles: compute_type: bfloat16 native_batch_max_size: 8 native_batch_max_tokens: 4096 + unescape_moses_entities: true runtime: first_chunk_timeout_s: 180 inter_chunk_timeout_s: 10 diff --git a/packages/sie_server/src/sie_server/adapters/ctranslate2/_moses_entities.py b/packages/sie_server/src/sie_server/adapters/ctranslate2/_moses_entities.py new file mode 100644 index 000000000..e0a89e70e --- /dev/null +++ b/packages/sie_server/src/sie_server/adapters/ctranslate2/_moses_entities.py @@ -0,0 +1,101 @@ +"""Undo the Moses tokenizer's entity escaping in translation output. + +The Moses tokenizer escapes ``& < > ' " [ ] |`` as entities and splits them +off the surrounding words with spaces. MADLAD-400 returns text in that form +for some target languages, so a French translation can come back as +``d ' Android`` instead of ``d'Android``. + +The repair is deliberately conservative: + +* An entity is only unescaped when the input text does not contain it + literally. A request that really translates ``&`` keeps it. +* Spacing is only repaired where the form is unambiguous: French elision, + when the prompt targets French (``<2fr>``, ``<2fr_CA>``), so + ``l ' offre`` becomes ``l'offre``; and balanced, space-delimited + ``"`` pairs (``" Bonjour "`` becomes ``"Bonjour"``). + Everything else is unescaped in place and keeps its spaces, so + ``R & D`` becomes ``R & D``. +""" + +from __future__ import annotations + +import re + +_ENTITIES = { + "&": "&", + "<": "<", + ">": ">", + "'": "'", + """: '"', + "[": "[", + "]": "]", + "|": "|", +} +# One pass over the text, so an escaped entity such as ``&lt;`` is +# unescaped once, to ``<``, and not again to ``<``. +_ENTITY = re.compile("|".join(re.escape(entity) for entity in _ENTITIES)) + +_APOS = "'" +_QUOT = """ + +# French elided forms followed by an escaped apostrophe and the next word: +# c', d', j', l', m', n', s', t', the qu' family (qu', jusqu', lorsqu', +# quelqu'...) and aujourd'hui. Only the elided word and the word after it +# decide the join, so a single-quoted span such as ``dit ' oui '`` +# is left spaced, while a quoted single letter that is itself an elided form +# (``' l ' est``) is joined to the word after it. The join only runs +# when the prompt targets French. +_ELISION = re.compile( + r"(?") + + +def unescape_moses_entities(text: str, source: str) -> str: + """Return ``text`` with Moses entities unescaped and their spacing repaired. + + Args: + text: Decoded model output. + source: The input text the output was generated from, target tag + included. Entities that occur literally in it are kept as they + are, and French elision is only repaired for a French target tag. + """ + escaped = frozenset(entity for entity in _ENTITIES if entity in text and entity not in source) + if not escaped: + return text + if _APOS in escaped and _FRENCH_TARGET.match(source): + text = _ELISION.sub(r"\1'", text) + if _QUOT in escaped: + text = _attach_quote_pairs(text) + return _ENTITY.sub(lambda match: _ENTITIES[match[0]] if match[0] in escaped else match[0], text) + + +def _attach_quote_pairs(text: str) -> str: + """Attach balanced ``"`` pairs to the words they enclose. + + The tokenizer turns ``"word"`` into ``" word "``. The output is + only respaced when every escaped quote can be paired in order, each pair + encloses text with exactly one space on either side, an opening quote + starts the text or follows whitespace and a closing quote does not precede + a word character. A literal ``"`` elsewhere in the output, an odd count or any other + shape leaves the spacing as it is. + """ + if '"' in text: + return text + parts = text.split(_QUOT) + if len(parts) % 2 == 0: + return text + for index in range(1, len(parts), 2): + before, inner, after = parts[index - 1], parts[index], parts[index + 1] + if (before or index > 1) and not before[-1:].isspace(): + return text + if len(inner) < 3 or inner[0] != " " or inner[-1] != " " or inner[1].isspace() or inner[-2].isspace(): + return text + if after and (after[0].isalnum() or after[0] == "_"): + return text + pieces = [parts[0]] + for index in range(1, len(parts), 2): + pieces.extend(('"', parts[index][1:-1], '"', parts[index + 1])) + return "".join(pieces) diff --git a/packages/sie_server/src/sie_server/adapters/ctranslate2/generation.py b/packages/sie_server/src/sie_server/adapters/ctranslate2/generation.py index 813a2e675..1e4fd6ea3 100644 --- a/packages/sie_server/src/sie_server/adapters/ctranslate2/generation.py +++ b/packages/sie_server/src/sie_server/adapters/ctranslate2/generation.py @@ -28,6 +28,7 @@ ) from sie_server.adapters._spec import AdapterSpec from sie_server.adapters._types import ERR_NOT_LOADED, ComputePrecision +from sie_server.adapters.ctranslate2._moses_entities import unescape_moses_entities as _unescape_moses_entities from sie_server.types.grammar import GrammarSpec from sie_server.types.inputs import ImageInput, VideoInput @@ -109,6 +110,7 @@ class _DecodeOptions: @dataclass(frozen=True, slots=True) class _PreparedGeneration: + prompt: str source_tokens: tuple[str, ...] prompt_tokens: int options: _DecodeOptions @@ -171,6 +173,7 @@ def __init__( native_batch_max_tokens: int = 8192, native_batch_wait_ms: float = 2.0, max_pending_requests: int = 64, + unescape_moses_entities: bool = False, ) -> None: artifact = Path(artifact_path) if not str(artifact_path): @@ -186,6 +189,11 @@ def __init__( self._native_batch_max_tokens = _positive_int("native_batch_max_tokens", native_batch_max_tokens) self._native_batch_wait_s = _non_negative_float("native_batch_wait_ms", native_batch_wait_ms) / 1000 self._max_pending_requests = _positive_int("max_pending_requests", max_pending_requests) + if not isinstance(unescape_moses_entities, bool): + raise ValueError("unescape_moses_entities must be a boolean") + # For models that emit Moses-escaped entities such as ``'`` in their + # output; see ``_moses_entities``. + self._unescape_moses_entities = unescape_moses_entities self._translator: Any | None = None self._tokenizer: Any | None = None @@ -433,6 +441,7 @@ def _prepare_generation(self, parameters: Mapping[str, Any]) -> _PreparedGenerat f"prompt token count {prompt_tokens} exceeds native_batch_max_tokens={self._native_batch_max_tokens}" ) return _PreparedGeneration( + prompt=prompt, source_tokens=source_tokens, prompt_tokens=prompt_tokens, options=_DecodeOptions( @@ -606,6 +615,8 @@ def callback(step: Any) -> bool: ended = bool(result_tokens and result_tokens[-1] in end_tokens) output_tokens = result_tokens[:-1] if ended else result_tokens text = self._decode_tokens(output_tokens) + if self._unescape_moses_entities: + text = _unescape_moses_entities(text, request.prepared.prompt) completion_tokens = len(output_tokens) finish_reason: Literal["stop", "length"] = "stop" if ended else "length" request.loop.call_soon_threadsafe(request.complete, text, completion_tokens, finish_reason) diff --git a/packages/sie_server/tests/adapters/test_ctranslate2_generation.py b/packages/sie_server/tests/adapters/test_ctranslate2_generation.py index 24a8aacd1..3682662ea 100644 --- a/packages/sie_server/tests/adapters/test_ctranslate2_generation.py +++ b/packages/sie_server/tests/adapters/test_ctranslate2_generation.py @@ -23,7 +23,24 @@ def __init__(self, *, clean_up_tokenization_spaces: bool = False) -> None: self.encode_calls = 0 self.decode_cleanup_calls: list[bool | None] = [] self.clean_up_tokenization_spaces = clean_up_tokenization_spaces - self._tokens = ["", "", "hello", "world", "bonjour", " monde", " !", "!"] + self._tokens = [ + "", + "", + "hello", + "world", + "bonjour", + " monde", + " !", + "!", + "&", + "l", + " '", + " ami", + " &", + " co", + "<2fr>", + "<2en>", + ] self.eos_token_id = 1 def encode(self, text: str, *, add_special_tokens: bool) -> list[int]: @@ -59,8 +76,15 @@ class _FakeResult: class _FakeTranslator: - def __init__(self, *, block: threading.Event | None = None, append_end_token: bool = False) -> None: + def __init__( + self, + *, + block: threading.Event | None = None, + append_end_token: bool = False, + outputs: list[str] | None = None, + ) -> None: self.calls: list[dict[str, Any]] = [] + self.outputs = outputs or ["bonjour", " monde"] self.block = block self.append_end_token = append_end_token self.started = threading.Event() @@ -73,7 +97,7 @@ def translate_batch(self, source: list[list[str]], **kwargs: Any) -> list[_FakeR if self.block is not None: self.block.wait(timeout=2) max_length = kwargs["max_decoding_length"] - outputs = ["bonjour", " monde"][:max_length] + outputs = self.outputs[:max_length] callback = kwargs["callback"] for batch_id in range(len(source)): for token in outputs: @@ -117,6 +141,7 @@ def _adapter( native_batch_max_tokens: int = 64, native_batch_wait_ms: float = 5, max_pending_requests: int = 64, + unescape_moses_entities: bool = False, ) -> tuple[CTranslate2GenerationAdapter, _FakeTranslator, _FakeCTranslate2, _FakeTokenizer]: fake_translator = translator or _FakeTranslator() fake_module = _FakeCTranslate2(fake_translator) @@ -129,6 +154,7 @@ def _adapter( native_batch_max_tokens=native_batch_max_tokens, native_batch_wait_ms=native_batch_wait_ms, max_pending_requests=max_pending_requests, + unescape_moses_entities=unescape_moses_entities, ) adapter._translator = fake_translator adapter._tokenizer = tokenizer @@ -332,6 +358,46 @@ def test_natural_end_token_at_max_length_reports_stop_without_counting_or_decodi assert chunks[-1].completion_tokens == 2 +_ESCAPED_OUTPUT = ["l", " '", " ami", " &", " co"] + + +@pytest.mark.parametrize( + ("unescape_moses_entities", "prompt", "expected"), + [ + (False, "<2fr> hello world", "l ' ami & co"), + (True, "<2fr> hello world", "l'ami & co"), + (True, "<2fr> hello & world", "l'ami & co"), + # Elision is only rejoined for a French target tag. + (True, "<2en> hello world", "l ' ami & co"), + ], +) +def test_output_unescapes_moses_entities_only_when_enabled( + tmp_path: Path, + unescape_moses_entities: bool, + prompt: str, + expected: str, +) -> None: + adapter, _, _, _ = _adapter( + tmp_path, + translator=_FakeTranslator(outputs=_ESCAPED_OUTPUT), + unescape_moses_entities=unescape_moses_entities, + ) + + chunks = asyncio.run(_collect(adapter.generate(prompt, max_new_tokens=8, seed=0))) + + assert "".join(chunk.text_delta for chunk in chunks) == expected + assert chunks[-1].completion_tokens == len(_ESCAPED_OUTPUT) + + +def test_unescape_moses_entities_option_must_be_a_boolean(tmp_path: Path) -> None: + with pytest.raises(ValueError, match="unescape_moses_entities must be a boolean"): + CTranslate2GenerationAdapter( + artifact_path=tmp_path, + ct2_compute_type="float16", + unescape_moses_entities="yes", # type: ignore[arg-type] + ) + + def test_scheduler_batches_only_exact_compatible_requests(tmp_path: Path) -> None: async def run() -> tuple[list[Any], list[Any], list[Any]]: adapter, translator, _, _ = _adapter(tmp_path, native_batch_wait_ms=20) diff --git a/packages/sie_server/tests/adapters/test_ctranslate2_moses_entities.py b/packages/sie_server/tests/adapters/test_ctranslate2_moses_entities.py new file mode 100644 index 000000000..96e9379e4 --- /dev/null +++ b/packages/sie_server/tests/adapters/test_ctranslate2_moses_entities.py @@ -0,0 +1,114 @@ +from __future__ import annotations + +import pytest +from sie_server.adapters.ctranslate2._moses_entities import unescape_moses_entities + +_SOURCE = "<2fr> A plain source sentence." + + +@pytest.mark.parametrize( + ("escaped", "expected"), + [ + ("C ' est l ' heure d ' aller à l ' école.", "C'est l'heure d'aller à l'école."), + ("Je n ' ai pas dit qu ' Anne s ' en va.", "Je n'ai pas dit qu'Anne s'en va."), + ("Il m ' a dit qu ' il t ' attend.", "Il m'a dit qu'il t'attend."), + ("J ' arrive jusqu ' ici.", "J'arrive jusqu'ici."), + ("Lorsqu ' il pleut, quelqu ' un reste.", "Lorsqu'il pleut, quelqu'un reste."), + ("Aujourd ' hui, l ' été commence.", "Aujourd'hui, l'été commence."), + ("L ' Hôtel ouvre et D ' ailleurs ferme.", "L'Hôtel ouvre et D'ailleurs ferme."), + ], +) +def test_french_elision_rejoins_the_apostrophe(escaped: str, expected: str) -> None: + assert unescape_moses_entities(escaped, _SOURCE) == expected + + +@pytest.mark.parametrize( + ("escaped", "expected"), + [ + # Not an elided form: unescaped, spacing kept. + ("Il dit ' oui ' .", "Il dit ' oui ' ."), + # Elided form followed by a digit rather than a word. + ("l ' 2", "l ' 2"), + # Elided letter inside a longer word. + ("le mot ' x", "le mot ' x"), + ], +) +def test_other_apostrophes_are_unescaped_without_respacing(escaped: str, expected: str) -> None: + assert unescape_moses_entities(escaped, _SOURCE) == expected + + +@pytest.mark.parametrize( + ("escaped", "expected"), + [ + ("Il a dit " bonjour " et il est parti.", 'Il a dit "bonjour" et il est parti.'), + ("" Oui, tout de suite ", dit-elle.", '"Oui, tout de suite", dit-elle.'), + ("" Un " puis " deux ".", '"Un" puis "deux".'), + ("Le projet " Nord " avance.", 'Le projet "Nord" avance.'), + ], +) +def test_balanced_quote_pairs_attach_to_the_quoted_text(escaped: str, expected: str) -> None: + assert unescape_moses_entities(escaped, _SOURCE) == expected + + +@pytest.mark.parametrize( + ("escaped", "expected"), + [ + # Odd count: which quote opens is unknown. + ("Fin du texte. "", 'Fin du texte. "'), + ("" a " b "", '" a " b "'), + # A literal quote elsewhere makes the pairing ambiguous. + ('"Oui" et " non "', '"Oui" et " non "'), + # A quote that touches a word is not a standalone token. + ("mot" x "", 'mot" x "'), + ("" x "mot", '" x "mot'), + # Not exactly one space inside the pair. + ("" x "", '" x "'), + ("" "", '" "'), + ], +) +def test_ambiguous_quotes_are_unescaped_without_respacing(escaped: str, expected: str) -> None: + assert unescape_moses_entities(escaped, _SOURCE) == expected + + +def test_remaining_moses_entities_are_unescaped_in_place() -> None: + escaped = "R & D, a < b > c, [ d ] et e | f" + + assert unescape_moses_entities(escaped, _SOURCE) == "R & D, a < b > c, [ d ] et e | f" + + +def test_unescaping_is_a_single_pass() -> None: + assert unescape_moses_entities("Écrivez &lt; et &apos;", _SOURCE) == "Écrivez < et '" + + +def test_entities_present_in_the_source_are_kept() -> None: + source = "<2fr> Tom & Jerry" + + assert unescape_moses_entities("Tom & Jerry", source) == "Tom & Jerry" + assert unescape_moses_entities("l ' ami de Tom & Jerry", source) == "l'ami de Tom & Jerry" + + +def test_elision_is_only_repaired_for_a_french_target() -> None: + escaped = "Rock ' n ' Roll" + + assert unescape_moses_entities(escaped, "<2en> Rock and roll") == "Rock ' n ' Roll" + assert unescape_moses_entities("l ' ami", "<2fr_CA> the friend") == "l'ami" + + +def test_a_quoted_elided_letter_is_joined_like_an_elision() -> None: + # Pinned on purpose: a quoted single letter that is an elided form is + # indistinguishable from an elision, so it joins the word after it. + escaped = "la lettre ' l ' est muette" + + assert unescape_moses_entities(escaped, _SOURCE) == "la lettre ' l'est muette" + + +def test_apostrophe_present_in_the_source_keeps_its_spacing() -> None: + source = "<2fr> Say ' hello" + + assert unescape_moses_entities("Dis l ' ami", source) == "Dis l ' ami" + + +def test_text_without_entities_is_returned_unchanged() -> None: + text = "Une phrase avec ' des \" signes & et < > [ ] |." + + assert unescape_moses_entities(text, _SOURCE) is text