fix(models): undo Moses entity escaping in MADLAD-400 output

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

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 &apos; 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 ``&amp;`` keeps it.
+* Spacing is only repaired where the form is unambiguous: French elision,
+  when the prompt targets French (``<2fr>``, ``<2fr_CA>``), so
+  ``l &apos; offre`` becomes ``l'offre``; and balanced, space-delimited
+  ``&quot;`` pairs (``&quot; Bonjour &quot;`` becomes ``"Bonjour"``).
+  Everything else is unescaped in place and keeps its spaces, so
+  ``R &amp; D`` becomes ``R & D``.
+"""
+
+from __future__ import annotations
+
+import re
+
+_ENTITIES = {
+    "&amp;": "&",
+    "&lt;": "<",
+    "&gt;": ">",
+    "&apos;": "'",
+    "&quot;": '"',
+    "&#91;": "[",
+    "&#93;": "]",
+    "&#124;": "|",
+}
+# One pass over the text, so an escaped entity such as ``&amp;lt;`` is
+# unescaped once, to ``&lt;``, and not again to ``<``.
+_ENTITY = re.compile("|".join(re.escape(entity) for entity in _ENTITIES))
+
+_APOS = "&apos;"
+_QUOT = "&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 &apos; oui &apos;``
+# is left spaced, while a quoted single letter that is itself an elided form
+# (``&apos; l &apos; est``) is joined to the word after it. The join only runs
+# when the prompt targets French.
+_ELISION = re.compile(
+    r"(?<!\w)([cdjlmnst]|[^\W\d_]*qu|aujourd) &apos; (?=[^\W\d_])",
+    re.IGNORECASE,
+)
+# MADLAD-400's target tag for French and its regional forms.
+_FRENCH_TARGET = re.compile(r"\s*<2fr(?:_[A-Za-z]+)?>")
+
+
+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 ``&quot;`` pairs to the words they enclose.
+
+    The tokenizer turns ``"word"`` into ``&quot; word &quot;``. 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)

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)

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 = ["<s>", "</s>", "hello", "world", "bonjour", " monde", " !", "!"]
+        self._tokens = [
+            "<s>",
+            "</s>",
+            "hello",
+            "world",
+            "bonjour",
+            " monde",
+            " !",
+            "!",
+            "&amp;",
+            "l",
+            " &apos;",
+            " ami",
+            " &amp;",
+            " 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", " &apos;", " ami", " &amp;", " co"]
+
+
+@pytest.mark.parametrize(
+    ("unescape_moses_entities", "prompt", "expected"),
+    [
+        (False, "<2fr> hello world", "l &apos; ami &amp; co"),
+        (True, "<2fr> hello world", "l'ami & co"),
+        (True, "<2fr> hello &amp; world", "l'ami &amp; 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)

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 &apos; est l &apos; heure d &apos; aller à l &apos; école.", "C'est l'heure d'aller à l'école."),
+        ("Je n &apos; ai pas dit qu &apos; Anne s &apos; en va.", "Je n'ai pas dit qu'Anne s'en va."),
+        ("Il m &apos; a dit qu &apos; il t &apos; attend.", "Il m'a dit qu'il t'attend."),
+        ("J &apos; arrive jusqu &apos; ici.", "J'arrive jusqu'ici."),
+        ("Lorsqu &apos; il pleut, quelqu &apos; un reste.", "Lorsqu'il pleut, quelqu'un reste."),
+        ("Aujourd &apos; hui, l &apos; été commence.", "Aujourd'hui, l'été commence."),
+        ("L &apos; Hôtel ouvre et D &apos; 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 &apos; oui &apos; .", "Il dit ' oui ' ."),
+        # Elided form followed by a digit rather than a word.
+        ("l &apos; 2", "l ' 2"),
+        # Elided letter inside a longer word.
+        ("le mot &apos; 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 &quot; bonjour &quot; et il est parti.", 'Il a dit "bonjour" et il est parti.'),
+        ("&quot; Oui, tout de suite &quot;, dit-elle.", '"Oui, tout de suite", dit-elle.'),
+        ("&quot; Un &quot; puis &quot; deux &quot;.", '"Un" puis "deux".'),
+        ("Le projet &quot; Nord &quot; 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. &quot;", 'Fin du texte. "'),
+        ("&quot; a &quot; b &quot;", '" a " b "'),
+        # A literal quote elsewhere makes the pairing ambiguous.
+        ('"Oui" et &quot; non &quot;', '"Oui" et " non "'),
+        # A quote that touches a word is not a standalone token.
+        ("mot&quot; x &quot;", 'mot" x "'),
+        ("&quot; x &quot;mot", '" x "mot'),
+        # Not exactly one space inside the pair.
+        ("&quot;  x &quot;", '"  x "'),
+        ("&quot; &quot;", '" "'),
+    ],
+)
+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 &amp; D, a &lt; b &gt; c, &#91; d &#93; et e &#124; 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 &amp;lt; et &amp;apos;", _SOURCE) == "Écrivez &lt; et &apos;"
+
+
+def test_entities_present_in_the_source_are_kept() -> None:
+    source = "<2fr> Tom &amp; Jerry"
+
+    assert unescape_moses_entities("Tom &amp; Jerry", source) == "Tom &amp; Jerry"
+    assert unescape_moses_entities("l &apos; ami de Tom &amp; Jerry", source) == "l'ami de Tom &amp; Jerry"
+
+
+def test_elision_is_only_repaired_for_a_french_target() -> None:
+    escaped = "Rock &apos; n &apos; Roll"
+
+    assert unescape_moses_entities(escaped, "<2en> Rock and roll") == "Rock ' n ' Roll"
+    assert unescape_moses_entities("l &apos; 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 &apos; l &apos; 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 &apos; hello"
+
+    assert unescape_moses_entities("Dis l &apos; ami", source) == "Dis l &apos; 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