diff --git a/packages/sie_server/src/sie_server/adapters/gliformer/relation_decoding.py b/packages/sie_server/src/sie_server/adapters/gliformer/relation_decoding.py index 952b9eb0c..f4613a5b6 100644 --- a/packages/sie_server/src/sie_server/adapters/gliformer/relation_decoding.py +++ b/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]) diff --git a/packages/sie_server/src/sie_server/adapters/gliformer/structuring_decoding.py b/packages/sie_server/src/sie_server/adapters/gliformer/structuring_decoding.py index dfbe6cf07..34584aee9 100644 --- a/packages/sie_server/src/sie_server/adapters/gliformer/structuring_decoding.py +++ b/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]) diff --git a/packages/sie_server/tests/adapters/test_gliformer_span_decoding.py b/packages/sie_server/tests/adapters/test_gliformer_span_decoding.py index 7a3aa0bc2..b94f77427 100644 --- a/packages/sie_server/tests/adapters/test_gliformer_span_decoding.py +++ b/packages/sie_server/tests/adapters/test_gliformer_span_decoding.py @@ -504,6 +504,108 @@ def test_record_text_is_bounded_per_document() -> None: assert structuring_decoding._within_word_limit(slots, 50) == [(0, []), (1, [])] +def test_record_text_is_bounded_by_characters() -> None: + span = upstream.Span + # One word each. The long token is under the word cap and over the character cap. + long = span(start=0, end=0, entity_type="a", score=0.9) + short = span(start=1, end=1, entity_type="b", score=0.4) + slots = [(0, [long]), (1, [short])] + tokens = ["x" * 30, "ok"] + assert structuring_decoding.MAX_RECORD_CHARS == structuring_decoding.MAX_RECORD_WORDS * 32 + assert structuring_decoding._within_word_limit(slots, 10, tokens=tokens, max_chars=40) is slots + # Best first is the 30-character token. It does not fit in 20, so taking stops + # and the later short span is dropped too, as an over-word span would be. + assert structuring_decoding._within_word_limit(slots, 10, tokens=tokens, max_chars=20) == [(0, []), (1, [])] + # A normal short span still passes; the long token ranked after it does not. + passing = [ + (0, [span(start=1, end=1, entity_type="b", score=0.9)]), + (1, [span(start=0, end=0, entity_type="a", score=0.4)]), + ] + kept = structuring_decoding._within_word_limit(passing, 10, tokens=tokens, max_chars=10) + assert [(slot, [(item.start, item.end) for item in spans]) for slot, spans in kept] == [(0, [(1, 1)]), (1, [])] + # Decoded text joins tokens with spaces, so two short words can still exceed the character cap. + spaced = [(0, [span(start=0, end=1, entity_type="a", score=0.9)])] + assert structuring_decoding._within_word_limit(spaced, 10, tokens=["abcd", "ef"], max_chars=6) == [(0, [])] + assert structuring_decoding._within_word_limit(spaced, 10, tokens=["abcd", "ef"], max_chars=7) is spaced + + +def test_ordinary_words_reach_the_word_cap_before_the_character_cap() -> None: + """A document of short words reaches the word cap before the character cap. + + Eight letters plus the joining space average under 31 characters a word, + so that document stays under the character cap. Words averaging more than + 32 characters hit the character cap first. + """ + ordinary = 8 + words = structuring_decoding.MAX_RECORD_WORDS + # A full word cap of 8-letter words, plus the joining spaces, stays under the character cap. + joined = words * ordinary + (words - 1) + assert joined < structuring_decoding.MAX_RECORD_CHARS + assert ( + relation_decoding.MAX_RELATION_WORDS * ordinary + (relation_decoding.MAX_RELATION_WORDS - 1) + < relation_decoding.MAX_RELATION_CHARS + ) + # The same number of words, averaging more than 32 characters, exceeds the character cap first. + assert words * (32 + 1) > structuring_decoding.MAX_RECORD_CHARS + assert relation_decoding.MAX_RELATION_WORDS * (32 + 1) > relation_decoding.MAX_RELATION_CHARS + assert relation_decoding.MAX_RELATION_CHARS == relation_decoding.MAX_RELATION_WORDS * 32 + + +class _TracingTokens(list[str]): + def __init__(self, items: list[str]) -> None: + super().__init__(items) + self.reads: list[object] = [] + + def __getitem__(self, item: object) -> object: + self.reads.append(item) + return super().__getitem__(item) # type: ignore[index] + + +def test_word_cap_rejection_does_not_read_the_rejected_tail() -> None: + span = upstream.Span + tokens = _TracingTokens(["ok", "z" * 100_000]) + slots = [ + (0, [span(start=0, end=0, entity_type="a", score=0.9)]), + (1, [span(start=1, end=1, entity_type="b", score=0.1)]), + ] + kept = structuring_decoding._within_word_limit(slots, 1, tokens=tokens, max_chars=10**9) + assert [(slot, [(item.start, item.end) for item in spans]) for slot, spans in kept] == [(0, [(0, 0)]), (1, [])] + assert not any(isinstance(read, slice) and read.start == 1 for read in tokens.reads) + + +def _record_fields(records: list) -> list[dict]: + return [field for group in records[0] for record in group for field in record] + + +def test_decoded_record_text_is_bounded_by_characters() -> None: + decode = structuring_decoding.make_structuring_decode( + upstream.Span, + structuring_decoder.unflatten_by_batch_origin, + max_words=100, + max_chars=10, + ) + decoder = structuring_decoder.StructuringDecoder(SimpleNamespace()) + + def output(short_logit: float, long_logit: float) -> SimpleNamespace: + return SimpleNamespace( + structuring_logits=torch.tensor([[[short_logit, long_logit]]]), + structuring_field_logits=torch.tensor([[[6.0], [6.0]]]), + structuring_span_idx=torch.tensor([[[0, 0], [1, 1]]]), + structuring_span_mask=torch.tensor([[True, True]]), + structuring_batch_origin=torch.arange(1), + batch_size=1, + ) + + long_text = "z" * 30 + # The long token is one word, so the word cap would keep it. Ranked first, it + # misses the character cap and taking stops, dropping the short span too. + assert _record_fields(decode(decoder, output(2.0, 6.0), threshold=0.1, texts=[["ok", long_text]])) == [] + fields = _record_fields(decode(decoder, output(6.0, 2.0), threshold=0.1, texts=[["ok", long_text]])) + assert [field["text"] for field in fields] == ["ok"] + fields = _record_fields(decode(decoder, output(2.0, 6.0), threshold=0.1, texts=[["ok", "yo"]])) + assert [field["text"] for field in fields] == ["ok", "yo"] + + def test_fully_nested_records_stay_within_the_text_bound() -> None: spans = [(i, 8190 - i) for i in range(4095)] count = len(spans) @@ -518,7 +620,9 @@ def test_fully_nested_records_stay_within_the_text_bound() -> None: records = _STRUCTURING.decode(output, threshold=0.1, flat_ner=False, texts=[[f"w{i}" for i in range(8191)]]) fields = [field for group in records[0] for record in group for field in record] words = sum(field["end"] - field["start"] + 1 for field in fields) + chars = sum(len(field["text"]) for field in fields) assert 0 < words <= structuring_decoding.MAX_RECORD_WORDS + assert 0 < chars <= structuring_decoding.MAX_RECORD_CHARS def test_record_spans_within_an_allowance_keep_each_slots_best_spans() -> None: @@ -850,6 +954,107 @@ def test_relation_text_is_bounded_per_document() -> None: assert 0 < len(capped) < len(full) +def test_relation_text_is_bounded_by_characters() -> None: + span = upstream.Span + long = span(start=0, end=0, entity_type="a", score=1.0) + short = span(start=1, end=1, entity_type="b", score=1.0) + entities = [long, short] + texts = [["y" * 30, "ok"]] + assert relation_decoding.MAX_RELATION_CHARS == relation_decoding.MAX_RELATION_WORDS * 32 + + def rows(scores: list[float], *, max_chars: int, max_words: int = 10, row_texts: list[list[str]] = texts) -> list: + return relation_decoding._row_relations( + _RELATIONS, + torch.tensor([[score] for score in scores], dtype=torch.float32), + torch.tensor([[0, 1], [1, 1]]), + None, + 0.5, + {0: "links"}, + None, + entities, + row_texts, + 0, + allowance=None, + max_relations=10, + max_words=max_words, + max_chars=max_chars, + ) + + # 30 + 2 characters and 2 words, then 2 + 2. Under both caps. + both = rows([0.9, 0.6], max_chars=40) + assert [(item["head"]["text"], item["tail"]["text"]) for item in both] == [("y" * 30, "ok"), ("ok", "ok")] + # Best relation is two words (under the word cap) and 32 characters. It does not + # fit in 20, so taking stops and the later short relation is dropped too. + assert rows([0.9, 0.6], max_chars=20) == [] + # A normal short relation still passes; the long one ranked after it does not. + kept = rows([0.6, 0.9], max_chars=10) + assert [(item["head"]["text"], item["tail"]["text"]) for item in kept] == [("ok", "ok")] + # Joined tokens count the space: "abcd ef" is 7 characters and 2 words. + wide = [span(start=0, end=1, entity_type="a", score=1.0)] + spaced = relation_decoding._row_relations( + _RELATIONS, + torch.tensor([[0.9]], dtype=torch.float32), + torch.tensor([[0, 0]]), + None, + 0.5, + {0: "links"}, + None, + wide, + [["abcd", "ef"]], + 0, + allowance=None, + max_relations=10, + max_words=10, + max_chars=6, + ) + assert spaced == [] + fitting = relation_decoding._row_relations( + _RELATIONS, + torch.tensor([[0.9]], dtype=torch.float32), + torch.tensor([[0, 0]]), + None, + 0.5, + {0: "links"}, + None, + wide, + [["abcd", "ef"]], + 0, + allowance=None, + max_relations=10, + max_words=10, + max_chars=14, + ) + assert fitting[0]["head"]["text"] == "abcd ef" + assert len(fitting[0]["head"]["text"]) + len(fitting[0]["tail"]["text"]) == 14 + + +def test_relation_word_cap_does_not_read_the_rejected_endpoint() -> None: + span = upstream.Span + tokens = _TracingTokens(["ok", "z" * 100_000]) + entities = [ + span(start=0, end=0, entity_type="a", score=1.0), + span(start=1, end=1, entity_type="b", score=1.0), + ] + rows = relation_decoding._row_relations( + _RELATIONS, + torch.tensor([[0.9], [0.6]], dtype=torch.float32), + torch.tensor([[0, 0], [1, 1]]), + None, + 0.5, + {0: "links"}, + None, + entities, + [tokens], + 0, + allowance=None, + max_relations=10, + max_words=2, + max_chars=10**9, + ) + assert [(item["head"]["text"], item["tail"]["text"]) for item in rows] == [("ok", "ok")] + assert not any(isinstance(read, slice) and read.start == 1 for read in tokens.reads) + + def test_relations_within_an_allowance_keep_the_best_first_prefix() -> None: output = _relation_output(3, rows=1, pairs=60, types=6) kwargs = {"threshold": 0.1, "texts": [[f"w{i}" for i in range(30)]]}