fix(server): bound GLiFormer decoded text by characters

packages/sie_server/src/sie_server/adapters/gliformer/relation_decoding.py

@@ -10,14 +10,15 @@
 passing cells with array operations and returns the same relations, in the
 same order, with the same scores. It also bounds each document's relations:
 at most ``max_relations`` of them, whose heads and tails together span at
-most ``max_words`` words, and no more than the document's allowance affords.
+most ``max_words`` words and ``max_chars`` characters, and no more than the
+document's allowance affords.
 Past a bound, a document keeps the best-first prefix: the highest scores
 first, and among equal scores the package's own order.
 """
 
 from __future__ import annotations
 
-from collections.abc import Callable
+from collections.abc import Callable, Sequence
 from typing import Any
 
 import numpy as np
@@ -31,6 +32,12 @@
 # Words of head and tail text across one document's relations. Each relation
 # copies both endpoints' text, and endpoints can be long.
 MAX_RELATION_WORDS = 262144
+# Characters of that head and tail text (#378). Same rule as the record cap.
+# The constant is MAX_RELATION_WORDS * 32: 32 characters per counted word,
+# including the joining space. The bound is per document, so text averaging
+# under 31 characters a word reaches the word cap first, and text made of
+# very long tokens reaches this cap.
+MAX_RELATION_CHARS = MAX_RELATION_WORDS * 32
 # Allowance units per kept relation (building it) and per scored cell.
 RELATION_UNITS = 2
 _CELLS_PER_UNIT = 1024
@@ -42,6 +49,7 @@ def make_relation_decode(
     *,
     max_relations: int = MAX_RELATIONS,
     max_words: int = MAX_RELATION_WORDS,
+    max_chars: int = MAX_RELATION_CHARS,
 ) -> Callable[..., Any]:
     """A replacement for ``gliformer.tasks.joint_relex.decoder.JointRelexDecoder.decode``.
 
@@ -50,6 +58,7 @@ def make_relation_decode(
         unflatten_by_batch_origin: The package's helper of that name.
         max_relations: Most relations per document.
         max_words: Most head and tail words across a document's relations.
+        max_chars: Most characters of decoded head and tail text per document.
 
     Returns:
         The ``decode`` method.
@@ -128,6 +137,7 @@ def decode(
                     allowance=row_allowance(bn),
                     max_relations=max_relations,
                     max_words=max_words,
+                    max_chars=max_chars,
                 )
             )
 
@@ -156,6 +166,7 @@ def _row_relations(
     allowance: Allowance | None,
     max_relations: int,
     max_words: int,
+    max_chars: int = MAX_RELATION_CHARS,
 ) -> list[dict]:
     """One document row's relation triples, in the package's order (pair, then type).
 
@@ -195,13 +206,43 @@ def _row_relations(
     if pairs.size == 0:
         return []
 
-    # Endpoint widths, for the text bound; ids outside the decoded list
-    # resolve to empty endpoints, as in the package.
+    # Endpoint word widths are arithmetic on the span bounds. Character widths
+    # read token text, so they are filled for every endpoint only when the
+    # word total still fits. When the word cap already rejects the row, text
+    # is read only for the best-first prefix that fits in the word budget.
+    tokens = None if texts is None or source_batch_idx >= len(texts) else texts[source_batch_idx]
     widths = np.asarray([span.end - span.start + 1 for span in entities] + [0], dtype=np.int64)
-    head_words = widths[np.where((heads >= 0) & (heads < len(entities)), heads, len(entities))]
-    tail_words = widths[np.where((tails >= 0) & (tails < len(entities)), tails, len(entities))]
+    head_at = np.where((heads >= 0) & (heads < len(entities)), heads, len(entities))
+    tail_at = np.where((tails >= 0) & (tails < len(entities)), tails, len(entities))
+    word_cost = widths[head_at] + widths[tail_at]
     limit = max_relations if allowance is None else min(max_relations, allowance.affordable(RELATION_UNITS))
-    keep = _best_first(scores, head_words + tail_words, limit, max_words)
+    word_total = int(word_cost.sum()) if word_cost.size else 0
+    chars: np.ndarray | None = None
+    char_of: Callable[[int], int] | None = None
+    if pairs.size <= limit and word_total <= max_words:
+        char_widths = np.asarray(
+            [_span_chars(tokens, span.start, span.end) for span in entities] + [0],
+            dtype=np.int64,
+        )
+        chars = char_widths[head_at] + char_widths[tail_at]
+    else:
+        cache: dict[int, int] = {}
+
+        def char_of(index: int) -> int:
+            total = 0
+            for entity_id in (int(heads[index]), int(tails[index])):
+                cached = cache.get(entity_id)
+                if cached is None:
+                    if 0 <= entity_id < len(entities):
+                        entity = entities[entity_id]
+                        cached = _span_chars(tokens, entity.start, entity.end)
+                    else:
+                        cached = 0
+                    cache[entity_id] = cached
+                total += cached
+            return total
+
+    keep = _best_first(scores, word_cost, limit, max_words, chars, max_chars, char_of=char_of)
     if allowance is not None:
         allowance.spend(RELATION_UNITS * int(keep.size))
 
@@ -238,15 +279,68 @@ def endpoint(entity_id: int) -> dict[str, Any]:
     return triples
 
 
-def _best_first(scores: np.ndarray, words: np.ndarray, limit: int, max_words: int) -> np.ndarray:
+def _span_chars(tokens: Sequence[str] | None, start: int, end: int) -> int:
+    """Characters of one span's decoded text.
+
+    The package copies ``" ".join(tokens[start:end + 1])``. With no tokens it
+    copies nothing, so the character cap does not apply.
+    """
+    if not tokens or start >= len(tokens):
+        return 0
+    span = tokens[start : end + 1]
+    if not span:
+        return 0
+    return sum(len(token) for token in span) + len(span) - 1
+
+
+def _best_first(
+    scores: np.ndarray,
+    words: np.ndarray,
+    limit: int,
+    max_words: int,
+    chars: np.ndarray | None,
+    max_chars: int,
+    *,
+    char_of: Callable[[int], int] | None = None,
+) -> np.ndarray:
     """Positions of the kept relations, in their original order.
 
     Relations are taken best score first (stable), and taking stops at the
-    first one that would exceed ``limit`` relations or ``max_words`` words.
+    first one that would exceed ``limit`` relations, ``max_words`` words, or
+    ``max_chars`` characters of decoded text.
+
+    ``chars`` is the per-relation character cost when it is already known.
+    When the word total is over the cap, pass ``chars=None`` and ``char_of``
+    so token text is read only for the prefix that still fits in the word
+    budget.
     """
-    if scores.size <= limit and int(words.sum()) <= max_words:
-        return np.arange(scores.size)
+    count = int(scores.size)
+    if count == 0:
+        return np.arange(0)
+    word_total = int(words.sum())
+    if count <= limit and word_total <= max_words:
+        if chars is None:
+            if char_of is None:
+                msg = "char_of is required when chars is omitted"
+                raise ValueError(msg)
+            chars = np.fromiter((char_of(index) for index in range(count)), dtype=np.int64, count=count)
+        if int(chars.sum()) <= max_chars:
+            return np.arange(count)
     ranked = np.argsort(-scores, kind="stable")[:limit]
-    fits = np.cumsum(words[ranked]) <= max_words
-    taken = ranked[: int(np.argmin(fits)) if not fits.all() else ranked.size]
-    return np.sort(taken)
+    if ranked.size == 0:
+        return ranked
+    word_fits = np.cumsum(words[ranked]) <= max_words
+    word_count = int(np.argmin(word_fits)) if not bool(word_fits.all()) else int(ranked.size)
+    if word_count == 0:
+        return np.empty(0, dtype=ranked.dtype)
+    prefix = ranked[:word_count]
+    if chars is None:
+        if char_of is None:
+            msg = "char_of is required when chars is omitted"
+            raise ValueError(msg)
+        prefix_chars = np.fromiter((char_of(int(index)) for index in prefix), dtype=np.int64, count=word_count)
+    else:
+        prefix_chars = chars[prefix]
+    char_fits = np.cumsum(prefix_chars) <= max_chars
+    kept = int(np.argmin(char_fits)) if not bool(char_fits.all()) else word_count
+    return np.sort(prefix[:kept])

packages/sie_server/src/sie_server/adapters/gliformer/structuring_decoding.py

@@ -21,7 +21,7 @@
 
 from __future__ import annotations
 
-from collections.abc import Callable
+from collections.abc import Callable, Sequence
 from typing import Any
 
 import numpy as np
@@ -41,6 +41,11 @@
 # the text alone grows with spans x document length. Flat single-label
 # removal keeps disjoint spans per slot: at most 100 slots x 2048 words.
 MAX_RECORD_WORDS = 262144
+# Characters of that field text (#378). The constant is MAX_RECORD_WORDS * 32:
+# 32 characters per counted word, including the joining space. The bound is
+# per document, so text averaging under 31 characters a word reaches the word
+# cap first, and text made of very long tokens reaches this cap.
+MAX_RECORD_CHARS = MAX_RECORD_WORDS * 32
 
 
 def make_structuring_decode(
@@ -49,6 +54,7 @@ def make_structuring_decode(
     *,
     max_spans: int = MAX_RECORD_SPANS,
     max_words: int = MAX_RECORD_WORDS,
+    max_chars: int = MAX_RECORD_CHARS,
 ) -> Callable[..., Any]:
     """A replacement for ``gliformer.tasks.structuring.decoder.StructuringDecoder.decode``.
 
@@ -57,6 +63,7 @@ def make_structuring_decode(
         unflatten_by_batch_origin: The package's helper of that name.
         max_spans: Most field spans per document, across its record slots.
         max_words: Most words of field text per document, across its slots.
+        max_chars: Most characters of decoded field text per document.
 
     Returns:
         The ``decode`` method.
@@ -145,6 +152,8 @@ def decode(
             kept = _within_word_limit(
                 [(anchor_idx, self.greedy_search(spans, flat_ner, multi_label)) for anchor_idx, spans in slot_spans],
                 max_words,
+                tokens=None if texts is None or text_idx >= len(texts) else texts[text_idx],
+                max_chars=max_chars,
             )
             for anchor_idx, spans in kept:
                 fields = self._spans_to_fields(spans, texts, text_idx)
@@ -283,27 +292,59 @@ def slot_spans(position: int, cut: np.float32 | None) -> tuple[np.ndarray, np.nd
     return result
 
 
-def _within_word_limit(slots: list[tuple[int, list[Any]]], max_words: int) -> list[tuple[int, list[Any]]]:
-    """Drop the lowest-scoring kept fields until their text fits ``max_words``.
+def _span_chars(tokens: Sequence[str] | None, start: int, end: int) -> int:
+    """Characters of one span's decoded text.
+
+    The package copies ``" ".join(tokens[start:end + 1])``. With no tokens it
+    copies nothing, so the character cap does not apply.
+    """
+    if not tokens or start >= len(tokens):
+        return 0
+    span = tokens[start : end + 1]
+    if not span:
+        return 0
+    return sum(len(token) for token in span) + len(span) - 1
+
+
+def _within_word_limit(
+    slots: list[tuple[int, list[Any]]],
+    max_words: int,
+    *,
+    tokens: Sequence[str] | None = None,
+    max_chars: int = MAX_RECORD_CHARS,
+) -> list[tuple[int, list[Any]]]:
+    """Drop the lowest-scoring kept fields until their text fits both bounds.
 
     Fields are ranked best score first, and among equal scores by slot and
-    then position; the best ones whose words add up to at most ``max_words``
-    stay, in their original order.
+    then position. Taking stops at the first field that would exceed
+    ``max_words`` words or ``max_chars`` characters of decoded text. The ones
+    taken stay in their original order.
+
+    When the word total is already over ``max_words``, token text is read
+    only for the best-first prefix that still fits in that word budget. The
+    tail the word cap rejects is not walked.
     """
     words = sum(span.end - span.start + 1 for _, spans in slots for span in spans)
     if words <= max_words:
-        return slots
+        chars = sum(_span_chars(tokens, span.start, span.end) for _, spans in slots for span in spans)
+        if chars <= max_chars:
+            return slots
     ranked = sorted(
         ((slot, index, span) for slot, (_, spans) in enumerate(slots) for index, span in enumerate(spans)),
         key=lambda entry: -entry[2].score,
     )
     keep: set[tuple[int, int]] = set()
-    used = 0
+    used_words = 0
+    used_chars = 0
     for slot, index, span in ranked:
         width = span.end - span.start + 1
-        if used + width > max_words:
+        if used_words + width > max_words:
+            break
+        span_chars = _span_chars(tokens, span.start, span.end)
+        if used_chars + span_chars > max_chars:
             break
-        used += width
+        used_words += width
+        used_chars += span_chars
         keep.add((slot, index))
     return [
         (anchor_idx, [span for index, span in enumerate(spans) if (slot, index) in keep])