From e6503490fdbe913e7c058c7f13c134cc2f7f1e30 Mon Sep 17 00:00:00 2001 From: entity-0-0-2 Date: Thu, 8 Oct 2026 21:03:21 -0700 Subject: [PATCH] fix(server): read long documents whole in GLiNER2 extraction, not only their first window gliner2 reads at most max_seq_length words (512) of a document, so entity, relation and structured extraction saw nothing after that point. #449 fixed classification in the same adapter and #552 made an overlong item fail with INPUT_TOO_LONG instead of coming back with a partial result, but a caller still could not extract from a long document at all. Entity, relation and structured extraction now read a text longer than one window as the same overlapping windows classification reads, and merge the windows' results. Entities and relation endpoints move back to document offsets and their text is re-read from the document, a span cut at a window edge gives way when the window next to it read that region whole, and the same span or relation keeps its highest score. Structured extraction takes the first window that found a string field and every window's list values. Input tokens are metered as the sum of the windows. A text that fits one window is read exactly as before, and a text taking more than 128 windows still comes back with the per-item INPUT_TOO_LONG error. --- .../sie_server/adapters/gliner2/adapter.py | 351 ++++++++++++++---- .../adapters/test_gliner2_complete_text.py | 51 ++- .../tests/adapters/test_gliner2_long_text.py | 279 +++++++++++++- 3 files changed, 580 insertions(+), 101 deletions(-) diff --git a/packages/sie_server/src/sie_server/adapters/gliner2/adapter.py b/packages/sie_server/src/sie_server/adapters/gliner2/adapter.py index 7c1f10dcc..a18a04bd6 100644 --- a/packages/sie_server/src/sie_server/adapters/gliner2/adapter.py +++ b/packages/sie_server/src/sie_server/adapters/gliner2/adapter.py @@ -18,9 +18,11 @@ from sie_server.adapters._types import ERR_REQUIRES_TEXT, ComputePrecision from sie_server.adapters._word_window import ( MAX_WORD_CHARS, + DocumentWindow, SubwordCounter, Window, WindowedSplitter, + merge_window_spans, plan_forwards, quadratic_attention, subword_budget, @@ -73,6 +75,10 @@ # A longer text is rejected with INVALID_INPUT rather than classified in part, # since a verdict on part of a text reads as a verdict on all of it. _MAX_CLASSIFY_WINDOWS = 128 +# Most windows one text is extracted in. A longer text is not read in part +# either: that item comes back with the per-item ``INPUT_TOO_LONG`` error and +# no extraction, and the other items of the request are unaffected. +_MAX_EXTRACT_WINDOWS = 128 # Characters of the rest of a text read to find its next window, doubled until # the window ends at least MAX_WORD_CHARS before them (or the text ends), so a # window is found without copying the rest of a long text each time. @@ -99,11 +105,14 @@ class GLiNER2Adapter(BaseAdapter): 2048), and each label, task name, field name or choice at most 128 characters; a longer prompt is rejected with ``INVALID_INPUT``. - gliner2 reads at most ``max_seq_length`` words of a document. Entity, - relation and structured extraction read that first window. Classification - reads all of a longer text, as overlapping windows of that many words, and - pools the windows' results into one (see ``pool_window_classifications``); - a text that fits one window is classified exactly as before. + gliner2 reads at most ``max_seq_length`` words of a document. Every task + reads all of a longer text, as overlapping windows of that many words: the + windows' classifications are pooled into one (see + ``pool_window_classifications``), and the windows' extractions are merged + into one result at document offsets. A text that fits one window is read + exactly as before, in the same call. A text that takes more than + ``_MAX_EXTRACT_WINDOWS`` windows of extraction comes back with a per-item + ``INPUT_TOO_LONG`` error rather than being read in part. Reference models: - fastino/gliner2-base-v1 @@ -277,9 +286,7 @@ def extract( """Extract entities, relations, classifications, or flat structured data.""" self._check_loaded() texts = [self._extract_text(item) for item in items] - # gliner2 reads a bounded window of words: pass it only the prefix holding them. reads = [self._read(text) for text in texts] - windows = [(model_text, subwords) for model_text, subwords, _ in reads] opts = options or {} effective_threshold = self._validate_threshold(opts.get("threshold", self._threshold)) classification_task = opts.get("classification_task", self._classification_task) @@ -302,12 +309,10 @@ def extract( check_label_chars("GLiNER2", "output_schema property names", output_schema["properties"]) check_label_chars("GLiNER2", "output_schema enum values", choices) prompt = self._prompt_tokens(specs, key=("json", tuple(specs)), extra=len(choices)) - rows = self._row_tokens(windows, prompt) with torch.inference_mode(): - raw_results, errors, input_token_counts = self._run_complete( + raw_results, errors, input_token_counts = self._run_whole( texts, reads, - rows, lambda batch: self._model.batch_extract_json( batch, structures, @@ -318,7 +323,9 @@ def extract( max_len=self._max_seq_length, ), input_token_counts=input_token_counts, + prompt=prompt, rows_per_pass=len(texts), + merge=lambda _, __, ___, results: self._merge_structured(results), ) return ExtractOutput( entities=[[] for _ in texts], @@ -347,12 +354,10 @@ def extract( prompt = self._prompt_tokens( normalized_labels, per_entry=_PROMPT_TOKENS_PER_RELATION, key=("relations", tuple(normalized_labels)) ) - rows = self._row_tokens(windows, prompt) with torch.inference_mode(): - raw_results, errors, input_token_counts = self._run_complete( + raw_results, errors, input_token_counts = self._run_whole( texts, reads, - rows, lambda batch: self._model.batch_extract_relations( batch, normalized_labels, @@ -363,7 +368,11 @@ def extract( max_len=self._max_seq_length, ), input_token_counts=input_token_counts, + prompt=prompt, rows_per_pass=len(texts), + merge=lambda index, text, plan, results: self._merge_relations( + text, plan, results, normalized_entities[index] + ), ) return ExtractOutput( entities=[ @@ -426,72 +435,247 @@ def extract_entities(batch: list[str]) -> list[Any]: check_label_chars("GLiNER2", "labels", normalized_labels) prompt = self._prompt_tokens(normalized_labels, key=("entities", tuple(normalized_labels))) with torch.inference_mode(): - raw_results, errors, input_token_counts = self._run_complete( + raw_results, errors, input_token_counts = self._run_whole( texts, reads, - self._row_tokens(windows, prompt), extract_entities, input_token_counts=input_token_counts, - rows_per_pass=1 if len(texts) == 1 else _PACKAGE_BATCH_SIZE, + prompt=prompt, + rows_per_pass=_PACKAGE_BATCH_SIZE, + merge=lambda _, text, plan, results: self._merge_entities(text, plan, results), ) all_entities = [self._flatten_entities(result, text=text) for text, result in zip(texts, raw_results)] return ExtractOutput(entities=all_entities, errors=errors, input_token_counts=input_token_counts) - def _run_complete( + def _run_whole( self, texts: list[str], reads: list[tuple[str, int | None, Window | None]], - rows: list[int] | None, run: Callable[[list[str]], list[Any]], *, + prompt: int | None, input_token_counts: list[int] | None, rows_per_pass: int, + merge: Callable[[int, str, tuple[list[int], list[int], list[tuple[str, int | None]]], list[Any]], Any], ) -> tuple[list[Any], list[ExtractItemError | None] | None, list[int] | None]: - """Run extraction only for items whose entire text fits the model's word/subword window. - - Classification has its own whole-document windowing. Entity, relation, - and structured extraction must fail an overlong item rather than return - a successful partial prediction. Inspect the last piece actually read: - a long word's ``cut`` can be beyond that piece, and lowercase offsets - need not index the original text. The package's synthetic sentence end - and trailing whitespace are not unread user content. + """Run extraction over every window of each item's text and merge the windows' results. + + Every window of every item runs in the same planned batch as one row + (see ``_run_planned``), and ``merge`` builds one item's result from its + windows' (``_merge_entities``, ``_merge_relations`` or + ``_merge_structured``). A text that fits the model's window is read as + that one window, exactly as before. A text taking more than + ``_MAX_EXTRACT_WINDOWS`` windows, or one whose windows' spans cannot be + moved back to it, comes back with the per-item ``INPUT_TOO_LONG`` + error, is not inferred and is metered at zero; the other items of the + request are unaffected. """ + plans: list[tuple[list[int], list[int], list[tuple[str, int | None]]] | None] = [] errors: list[ExtractItemError | None] = [] - for text, (_, _, window) in zip(texts, reads, strict=True): + for text, read in zip(texts, reads, strict=True): + plan = self._windows(text, read, limit=_MAX_EXTRACT_WINDOWS) + if plan is None: + errors.append( + ExtractItemError( + code="INPUT_TOO_LONG", + message=f"GLiNER2 reads a text in at most {_MAX_EXTRACT_WINDOWS} windows of " + f"{self._max_seq_length or _DEFAULT_MAX_WORDS} words. Split the text into smaller items.", + ) + ) + plans.append(None) + continue source = text.lower() if self._lower_text_first else text - incomplete = ( - window is not None - and window.cut is not None - and bool(window.words) - and _NON_SPACE.search(source, window.words[-1][2]) is not None - ) - errors.append( - ExtractItemError( - code="INPUT_TOO_LONG", - message="GLiNER2 extraction requires the entire text to fit one word/subword window. " - "Split the text into smaller items.", + if len(plan[2]) > 1 and len(source) != len(text): + errors.append( + ExtractItemError( + code="INPUT_TOO_LONG", + message="GLiNER2 lowercases this text before reading it, and lowercasing changes its " + "length, so the spans of its windows cannot be moved back to the text. " + "Split the text into smaller items.", + ) ) - if incomplete - else None - ) - accepted = [index for index, error in enumerate(errors) if error is None] + plans.append(None) + continue + errors.append(None) + plans.append(plan) + flat = [window for plan in plans if plan is not None for window in plan[2]] results: list[Any] = [{} for _ in texts] - if accepted: + if flat: raw = self._run_planned( - [reads[index][0] for index in accepted], - [rows[index] for index in accepted] if rows is not None else None, + [model_text for model_text, _ in flat], + self._row_tokens(flat, prompt), run, - rows_per_pass=rows_per_pass, + rows_per_pass=1 if len(flat) == 1 else rows_per_pass, ) - for index, result in zip(accepted, raw, strict=True): - results[index] = result - counts = ( - [count if error is None else 0 for count, error in zip(input_token_counts, errors, strict=True)] - if input_token_counts is not None - else ([0] * len(texts) if not accepted else None) - ) - return results, errors if len(accepted) != len(texts) else None, counts + if len(raw) != len(flat): + raise ValueError("GLiNER2 returned results for a different number of items") + offset = 0 + for index, plan in enumerate(plans): + if plan is None: + continue + window_results = raw[offset : offset + len(plan[2])] + offset += len(plan[2]) + results[index] = ( + merge(index, texts[index], plan, window_results) if len(plan[2]) > 1 else window_results[0] + ) + if any(plan is not None and len(plan[2]) > 1 for plan in plans): + counts = self._windowed_input_token_counts(texts, [plan[2] if plan is not None else [] for plan in plans]) + else: + counts = ( + [count if error is None else 0 for count, error in zip(input_token_counts, errors, strict=True)] + if input_token_counts is not None + else ([0] * len(texts) if all(plan is None for plan in plans) else None) + ) + return results, errors if any(error is not None for error in errors) else None, counts + + def _merge_entities( + self, + text: str, + plan: tuple[list[int], list[int], list[tuple[str, int | None]]], + results: list[Any], + ) -> dict[str, Any]: + """One text's entities from the entities found in each of its windows. + + Each window's spans move back to the text by the window's start, and + the windows' spans merge as the GLiNER adapters merge theirs (see + ``_word_window.merge_window_spans``): a span cut at a window's edge + gives way when the window next to it read that region whole and found + an overlapping span of the same label, the same span keeps its highest + score, and overlapping spans resolve highest score first. A span's text + is re-read from the text, so it keeps its case whatever the windows + were read from. + """ + starts, read_ends, windows = plan + found: list[list[dict[str, Any]]] = [] + for result in results: + spans: list[dict[str, Any]] = [] + for label, label_spans in (result.get("entities") or {}).items(): + if not isinstance(label_spans, list): + continue + for span in label_spans: + if isinstance(span, dict): + spans.append( + { + "start": span.get("start"), + "end": span.get("end"), + "label": label, + "score": span.get("confidence", 0.0), + "text": span.get("text", ""), + } + ) + found.append(spans) + document = [ + DocumentWindow(start, start + len(model_text), read_end, None) + for start, read_end, (model_text, _) in zip(starts, read_ends, windows, strict=True) + ] + merged = merge_window_spans(document, found, text, flat_ner=False, multi_label=True) + entities: dict[str, Any] = {} + for span in merged: + entities.setdefault(span["label"], []).append( + {"text": span["text"], "start": span["start"], "end": span["end"], "confidence": span["score"]} + ) + return {"entities": entities} + + def _merge_relations( + self, + text: str, + plan: tuple[list[int], list[int], list[tuple[str, int | None]]], + results: list[Any], + entities: list[Entity], + ) -> dict[str, Any]: + """One text's relations from the relations found in each of its windows. + + A relation found by several windows is kept once per (head, tail, + relation) at its highest score, and a relation's head and tail are + moved back to the text: their spans move by the window's start and the + text is re-read there, or an endpoint's text matches an entity's text + case-insensitively when it carries no span. A relation whose head or + tail is none of ``entities`` is dropped, as in one window. A relation + whose head and tail no single window holds together is not found. + """ + starts = plan[0] + originals = {entity["text"] for entity in entities} + by_lower: dict[str, str] = {} + for original_text in originals: + by_lower.setdefault(original_text.lower(), original_text) + + def original(endpoint: Any, offset: int) -> str | None: + if not isinstance(endpoint, dict): + return None + found = endpoint.get("text") + if not isinstance(found, str): + return None + if found in originals: + return found + start, end = endpoint.get("start"), endpoint.get("end") + if ( + isinstance(start, int) + and not isinstance(start, bool) + and isinstance(end, int) + and not isinstance(end, bool) + ): + moved_start, moved_end = offset + start, offset + end + if 0 <= moved_start < moved_end <= len(text) and text[moved_start:moved_end].lower() == found.lower(): + return text[moved_start:moved_end] + return by_lower.get(found.lower()) + + best: dict[tuple[str, str, str], tuple[float, dict[str, Any]]] = {} + for offset, result in zip(starts, results, strict=True): + by_type = result.get("relation_extraction") or {} + if not isinstance(by_type, dict): + continue + for relation_type, candidates in by_type.items(): + if not isinstance(relation_type, str) or not isinstance(candidates, list): + continue + for candidate in candidates: + if not isinstance(candidate, dict): + continue + head, tail = candidate.get("head"), candidate.get("tail") + resolved_head, resolved_tail = original(head, offset), original(tail, offset) + if resolved_head is None or resolved_tail is None: + continue + head_confidence = head.get("confidence", 0.0) if isinstance(head, dict) else 0.0 + tail_confidence = tail.get("confidence", 0.0) if isinstance(tail, dict) else 0.0 + score = min(head_confidence, tail_confidence) + key = (resolved_head, resolved_tail, relation_type) + if key not in best or score > best[key][0]: + best[key] = ( + score, + { + "head": {"text": resolved_head, "confidence": head_confidence}, + "tail": {"text": resolved_tail, "confidence": tail_confidence}, + }, + ) + extraction: dict[str, Any] = {} + for (_, _, relation_type), (_, candidate) in best.items(): + extraction.setdefault(relation_type, []).append(candidate) + return {"relation_extraction": extraction} + + def _merge_structured(self, results: list[Any]) -> dict[str, Any]: + """One text's structured data from the data extracted from each of its windows. + + A string field takes the first window that found a value for it, and a + list field takes every window's values in window order, once each. A + field no window found stays missing. + """ + merged: dict[str, Any] = {} + for result in results: + values = result.get(_STRUCTURE_NAME) or [] + if not isinstance(values, list): + continue + for record in values: + if not isinstance(record, dict): + continue + for field, value in record.items(): + if isinstance(value, list): + kept = merged.setdefault(field, []) + for one in value: + if one not in kept: + kept.append(one) + elif value is not None and value != "" and merged.get(field) is None: + merged[field] = value + return {_STRUCTURE_NAME: [merged]} def _classify( self, @@ -1059,43 +1243,64 @@ def _classification_windows( """The windows ``text`` is classified in, as ``(model text, subwords)``. ``first`` is ``_read(text)``, the window gliner2 reads of ``text``. A - text all of whose words it reads is one window, the text as before. - Otherwise each next window starts at the start of the last - ``_CLASSIFY_OVERLAP_WORDS`` words the one before it read (at most half - of them, so every window reads new words), and is read as gliner2 - reads a text starting there: at most ``max_seq_length`` words within - the subword budget. A window after the first is cut from the text as - the splitter reads it (lowercased first by gliner2 1.x, which - lowercases it again to the same text), so that the words' offsets - index it. + text all of whose words it reads is one window, the text as before; + a longer text is read as the windows of ``_windows``. Raises: InvalidInputError: The text takes more than ``_MAX_CLASSIFY_WINDOWS`` windows. """ + windows = self._windows(text, first, limit=_MAX_CLASSIFY_WINDOWS) + if windows is None: + raise InvalidInputError( + f"GLiNER2 classifies a text in at most {_MAX_CLASSIFY_WINDOWS} windows of " + f"{self._max_seq_length or _DEFAULT_MAX_WORDS} words; this text needs more. " + "Split it into several items." + ) + return windows[2] + + def _windows( + self, text: str, first: tuple[str, int | None, Window | None], *, limit: int + ) -> tuple[list[int], list[int], list[tuple[str, int | None]]] | None: + """``(starts, read ends, windows)``: every window ``text`` is read in. + + Each window is ``(model text, subwords)`` as ``_read`` returns it, so a + text all of whose words gliner2 reads is one window, the text as + before. ``first`` is ``_read(text)``. Otherwise each next window starts + at the start of the last ``_CLASSIFY_OVERLAP_WORDS`` words the one + before it read (at most half of them, so every window reads new + words), and is read as gliner2 reads a text starting there: at most + ``max_seq_length`` words within the subword budget. A window after the + first is cut from the text as the splitter reads it (lowercased first + by gliner2 1.x, which lowercases it again to the same text), so that + the words' offsets index it. ``starts[i]`` is where window ``i`` begins + in ``text`` and ``read ends[i]`` is where its last word ends there, so + a span found in window ``i`` moves back to the text by ``starts[i]``. + + Returns ``None`` when the text takes more than ``limit`` windows. + """ model_text, subwords, window = first windows = [(model_text, subwords)] + starts = [0] + read_ends = [window.words[-1][2] if window is not None and window.words else 0] if window is None or window.cut is None: - return windows + return starts, read_ends, windows source = text.lower() if self._lower_text_first else text start = 0 while True: words = window.words read_end = start + words[-1][2] if _NON_SPACE.search(source, read_end) is None: - # Only the "." gliner2 appends, or nothing, is left unread. - return windows - if len(windows) >= _MAX_CLASSIFY_WINDOWS: - raise InvalidInputError( - f"GLiNER2 classifies a text in at most {_MAX_CLASSIFY_WINDOWS} windows of " - f"{self._max_seq_length or _DEFAULT_MAX_WORDS} words; this text needs more. " - "Split it into several items." - ) + return starts, read_ends, windows + if len(windows) >= limit: + return None shared = min(_CLASSIFY_OVERLAP_WORDS, len(words) // 2) start = start + words[len(words) - shared][1] if shared else read_end model_text, subwords, window = self._read_from(source, start) windows.append((model_text, subwords)) + starts.append(start) + read_ends.append(start + (window.words[-1][2] if window is not None and window.words else 0)) if window is None or window.cut is None: - return windows + return starts, read_ends, windows def _read_from(self, source: str, start: int) -> tuple[str, int | None, Window | None]: """``_read`` of ``source[start:]``, from a slice of it long enough to give the same window. diff --git a/packages/sie_server/tests/adapters/test_gliner2_complete_text.py b/packages/sie_server/tests/adapters/test_gliner2_complete_text.py index 81d445561..05c9f4c1b 100644 --- a/packages/sie_server/tests/adapters/test_gliner2_complete_text.py +++ b/packages/sie_server/tests/adapters/test_gliner2_complete_text.py @@ -75,7 +75,10 @@ def _request(task: str, texts: list[str]) -> tuple[list[Item], dict[str, Any]]: @pytest.mark.parametrize("task", ["entities", "relations", "structured"]) -def test_overlong_items_fail_without_losing_short_items_or_their_positions(task: str) -> None: +def test_overlong_items_fail_without_losing_short_items_or_their_positions( + task: str, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr("sie_server.adapters.gliner2.adapter._MAX_EXTRACT_WINDOWS", 2) adapter, model = _adapter() long = "Alice Acme " + "word " * 20 items, kwargs = _request(task, ["Alice Acme", long, "Bob Acme", long]) @@ -117,7 +120,8 @@ def test_overlong_items_fail_without_losing_short_items_or_their_positions(task: @pytest.mark.parametrize("task", ["entities", "relations", "structured"]) -def test_an_all_overlong_batch_never_calls_the_model(task: str) -> None: +def test_an_all_overlong_batch_never_calls_the_model(task: str, monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr("sie_server.adapters.gliner2.adapter._MAX_EXTRACT_WINDOWS", 2) adapter, model = _adapter() items, kwargs = _request(task, ["Alice Acme " + "word " * 20] * 2) output = adapter.extract(items, **kwargs) @@ -132,34 +136,36 @@ def test_an_all_overlong_batch_never_calls_the_model(task: str) -> None: @pytest.mark.parametrize("lower", [False, True]) @pytest.mark.parametrize( - ("text", "max_words", "max_subwords", "too_long"), + ("text", "max_words", "max_subwords"), [ - ("one two three", 3, 1000, False), # only the package's synthetic dot is unread - ("one two three \t\n", 3, 1000, False), - ("one two three.", 3, 1000, True), # the caller's punctuation is real input - ("one two three four", 3, 1000, True), - ("\u0130stanbul Ankara Izmir", 2, 1000, True), # lowercase expands offsets - ("\u039f\u0394\u039f\u03a3'\u0391 next", 1, 1000, True), - ("x" * 513, 2, 1000, True), # cut is the original word end, beyond the last read piece - ("abc " * 7, 8, 3, True), # fewer than max_words, but the subword budget is full + ("one two three", 3, 1000), # only the package's synthetic dot is unread + ("one two three \t\n", 3, 1000), + ("one two three.", 3, 1000), + ("one two three four", 3, 1000), + ("\u0130stanbul Ankara Izmir", 2, 1000), # lowercase expands offsets + ("\u039f\u0394\u039f\u03a3'\u0391 next", 1, 1000), + ("x" * 513, 2, 1000), # a long word is read in pieces across windows + ("abc " * 7, 8, 3), # the subword budget bounds each window's words ], ) -def test_completeness_uses_the_actual_word_window( - text: str, max_words: int, max_subwords: int, too_long: bool, lower: bool +def test_a_long_text_is_read_in_windows_without_an_error( + text: str, max_words: int, max_subwords: int, lower: bool ) -> None: adapter, model = _adapter(lower=lower, max_words=max_words, max_subwords=max_subwords) model.extract_entities.side_effect = None model.extract_entities.return_value = {"entities": {}} output = adapter.extract([Item(text=text)], labels=["person"]) - if too_long: + if lower and text.startswith("\u0130"): + # Lowercasing the text changes its length, so its windows' spans cannot be moved back to it. assert output.errors is not None assert output.errors[0] is not None assert output.errors[0].code == "INPUT_TOO_LONG" assert output.input_token_counts == [0] model.extract_entities.assert_not_called() - else: - assert output.errors is None - model.extract_entities.assert_called_once() + model.batch_extract_entities.assert_not_called() + return + assert output.errors is None + assert model.extract_entities.called or model.batch_extract_entities.called def test_invalid_request_parameters_are_still_validated_before_item_length_errors() -> None: @@ -170,13 +176,14 @@ def test_invalid_request_parameters_are_still_validated_before_item_length_error @pytest.mark.parametrize("index", [0, 1]) -def test_queue_outcome_preserves_length_error_and_zero_billing(index: int) -> None: +def test_queue_outcome_preserves_length_error_and_zero_billing(index: int, monkeypatch: pytest.MonkeyPatch) -> None: from types import SimpleNamespace import msgpack from sie_server.ipc_types import ExtractBatchItem from sie_server.queue_executor import _extract_success_outcome + monkeypatch.setattr("sie_server.adapters.gliner2.adapter._MAX_EXTRACT_WINDOWS", 2) adapter, _ = _adapter() items = [Item(text="word " * 20), Item(text="Alice Acme")] output = adapter.extract(items, labels=["person"]) @@ -206,13 +213,14 @@ def test_queue_outcome_preserves_length_error_and_zero_billing(index: int) -> No assert outcome.units.input_tokens == 4 -def test_queue_length_error_keeps_zero_billing_when_sibling_metering_fails() -> None: +def test_queue_length_error_keeps_zero_billing_when_sibling_metering_fails(monkeypatch: pytest.MonkeyPatch) -> None: from types import SimpleNamespace import msgpack from sie_server.ipc_types import ExtractBatchItem from sie_server.queue_executor import _extract_success_outcome + monkeypatch.setattr("sie_server.adapters.gliner2.adapter._MAX_EXTRACT_WINDOWS", 2) adapter, _ = _adapter() adapter._doc_input_token_counts = lambda _texts: None items = [Item(text="word " * 20), Item(text="Alice Acme")] @@ -243,13 +251,16 @@ def test_queue_length_error_keeps_zero_billing_when_sibling_metering_fails() -> @pytest.mark.parametrize(("error_code", "expected_input_tokens"), [("INPUT_TOO_LONG", 0), ("INFERENCE_ERROR", 7)]) -def test_queue_length_error_overrides_reported_positive_billing(error_code: str, expected_input_tokens: int) -> None: +def test_queue_length_error_overrides_reported_positive_billing( + error_code: str, expected_input_tokens: int, monkeypatch: pytest.MonkeyPatch +) -> None: from types import SimpleNamespace import msgpack from sie_server.ipc_types import ExtractBatchItem from sie_server.queue_executor import _extract_success_outcome + monkeypatch.setattr("sie_server.adapters.gliner2.adapter._MAX_EXTRACT_WINDOWS", 2) adapter, _ = _adapter() items = [Item(text="word " * 20), Item(text="Alice Acme")] output = adapter.extract(items, labels=["person"]) diff --git a/packages/sie_server/tests/adapters/test_gliner2_long_text.py b/packages/sie_server/tests/adapters/test_gliner2_long_text.py index 2c1bf8dd6..ea45bdd49 100644 --- a/packages/sie_server/tests/adapters/test_gliner2_long_text.py +++ b/packages/sie_server/tests/adapters/test_gliner2_long_text.py @@ -21,7 +21,12 @@ import pytest from sie_server.adapters._word_window import ATTENTION_BUDGET, MAX_WORD_CHARS, WindowedSplitter, subword_budget -from sie_server.adapters.gliner2.adapter import _PROMPT_TOKENS_PER_RELATION, _SUBWORDS_PER_WORD, GLiNER2Adapter +from sie_server.adapters.gliner2.adapter import ( + _MAX_EXTRACT_WINDOWS, + _PROMPT_TOKENS_PER_RELATION, + _SUBWORDS_PER_WORD, + GLiNER2Adapter, +) from sie_server.adapters.gliner2.classification import GLiNER2ClassificationAdapter from sie_server.adapters.gliner2.words import PACKAGE_PATTERN, LinearWordSplitter from sie_server.types.inputs import InvalidInputError, Item @@ -795,21 +800,279 @@ def test_positive_label_is_validated() -> None: @pytest.mark.parametrize("task", ["entities", "relations", "structured"]) -def test_long_item_rejection_uses_the_pinned_package_processor(task: str) -> None: +def test_a_long_item_is_read_whole_with_the_pinned_package_processor(task: str) -> None: adapter, model = make_adapter() - text = " ".join(["word"] * 512 + ["Alice"]) + text = " ".join(["word"] * 512 + ["Alice", "met", "Bob"]) kwargs: dict[str, Any] = {"labels": ["person"]} item = Item(text=text) if task == "structured": kwargs = {"output_schema": {"type": "object", "properties": {"name": {"type": "string"}}}} elif task == "relations": kwargs = {"labels": ["knows"]} - item = Item(text=text, metadata={"entities": [{"text": "Alice", "start": len(text) - 5, "end": len(text)}]}) + item = Item( + text=text, + metadata={ + "entities": [ + {"text": "Alice", "start": text.index("Alice"), "end": text.index("Alice") + 5}, + {"text": "Bob", "start": text.index("Bob"), "end": text.index("Bob") + 3}, + ] + }, + ) output = adapter.extract([item], **kwargs) + assert output.errors is None + assert output.input_token_counts is not None + assert output.input_token_counts[0] > 0 + (batch,) = model.inputs + assert len(batch.text_tokens) == 2 + assert "alice" in batch.text_tokens[-1] + + +# --- Extraction of texts longer than one window ------------------------------------------------- + + +class SpanModel(PackageModel): + """gliner2's preprocessing, with a stand-in extractor: a window yields the person it reads. + + Records each window's text. A window holding the whole name yields it as + one span at the window's offsets; a window cut inside the name yields the + word of it that window holds whole, as a model cut at an edge would. + """ + + def __init__(self, processor: Any, target: str = "priya raman") -> None: + super().__init__(processor) + self.target = target + self.rows: list[str] = [] + + def batch_extract_entities( + self, texts: list[str], labels: list[str], *, max_len: int | None, **_: Any + ) -> list[Any]: + self._collate(texts, gliner2_engine.Schema().entities(labels), max_len) + results = [] + for text in texts: + self.rows.append(text) + lower = text.lower() + start = lower.find(self.target) + length = len(self.target) + if start < 0: + start = lower.find(self.target.split()[0]) + length = len(self.target.split()[0]) + if start < 0: + results.append({"entities": {}}) + continue + results.append( + { + "entities": { + "person": [ + { + "text": text[start : start + length], + "start": start, + "end": start + length, + "confidence": 0.9, + } + ] + } + } + ) + return results + + +def make_span(*, max_seq_length: int = 512, target: str = "priya raman") -> tuple[GLiNER2Adapter, SpanModel]: + adapter = GLiNER2Adapter("fake/gliner2", max_seq_length=max_seq_length) + model = SpanModel(make_processor(), target=target) + adapter._use_linear_word_splitter(model.processor) + adapter._model = model + return adapter, model + + +def document(words: int, planted_at: int | None = None) -> str: + """An ordinary document of ``words`` words, with the name planted ``planted_at`` words in. + + No punctuation, so that gliner2 reads exactly these words (and the "." it appends). + """ + out = benign_email(words).split() + if planted_at is not None: + out[planted_at : planted_at + 2] = ["Priya", "Raman"] + return " ".join(out) + + +def test_an_entity_past_the_first_window_is_extracted_at_its_document_offsets() -> None: + text = document(1500, planted_at=1420) + adapter, model = make_span() + assert "priya raman" not in adapter._model_text(text).lower() + + output = adapter.extract([Item(text=text)], labels=["person"]) + + start = text.index("Priya") + assert output.errors is None + assert output.entities == [ + [{"text": "Priya Raman", "label": "person", "score": 0.9, "start": start, "end": start + len("Priya Raman")}] + ] + # 1,500 words (and gliner2's ".") in windows of 512 that start 448 words apart, as classification's. + assert len(model.rows) == 4 + for before, after in zip(model.rows, model.rows[1:], strict=False): + assert before.split()[-64:] == after.split()[:64] + assert [len(row.split()) for row in model.rows] == [512, 512, 512, 156] + + +def test_an_entity_straddling_a_window_edge_is_found_once_whole() -> None: + # The name starts on the first window's last word and ends on the next one's first. + text = document(900, planted_at=511) + adapter, model = make_span() + + output = adapter.extract([Item(text=text)], labels=["person"]) + + start = text.index("Priya") + assert output.errors is None + assert output.entities == [ + [{"text": "Priya Raman", "label": "person", "score": 0.9, "start": start, "end": start + len("Priya Raman")}] + ] + assert "priya raman" not in model.rows[0].lower() + assert "priya" in model.rows[0].lower() + + +def test_a_batch_of_long_and_short_documents_keeps_its_positions() -> None: + adapter, model = make_span() + texts = [document(1500, planted_at=1420), "Priya Raman works at Novartis.", document(900, planted_at=800)] + + output = adapter.extract([Item(text=text) for text in texts], labels=["person"]) + + assert output.errors is None + assert [len(found) for found in output.entities] == [1, 1, 1] + assert output.entities[1][0]["text"] == "Priya Raman" + assert output.entities[1][0]["start"] == 0 + assert len(model.rows) == 4 + 1 + 2 + + +def test_extraction_metering_counts_every_window() -> None: + text = document(1500, planted_at=1420) + adapter, _ = make_span() + windows = adapter._windows(text, adapter._read(text), limit=_MAX_EXTRACT_WINDOWS) + assert windows is not None + per_window = adapter._doc_input_token_counts([model_text for model_text, _ in windows[2]]) + assert per_window is not None + + output = adapter.extract([Item(text=text), Item(text="Priya Raman works at Novartis.")], labels=["person"]) + + assert output.input_token_counts == [ + sum(per_window), + adapter._doc_input_token_counts(["Priya Raman works at Novartis."])[0], + ] + assert per_window == [512, 512, 512, per_window[-1]] + assert 0 < per_window[-1] < 512 + + +def test_a_text_longer_than_the_extraction_cap_errors_per_item(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr("sie_server.adapters.gliner2.adapter._MAX_EXTRACT_WINDOWS", 3) + adapter, model = make_span(max_seq_length=8) + # Windows of 8 words, 4 apart: three windows read 16 words (and gliner2's appended "."). + fits = document(15) + too_long = document(17) + + output = adapter.extract([Item(text=fits), Item(text=too_long)], labels=["person"]) + assert output.errors is not None - assert output.errors[0] is not None - assert output.errors[0].code == "INPUT_TOO_LONG" - assert output.input_token_counts == [0] - assert not model.inputs + assert output.errors[0] is None + assert output.errors[1] is not None + assert output.errors[1].code == "INPUT_TOO_LONG" + assert output.entities[1] == [] + assert output.input_token_counts is not None + assert output.input_token_counts[1] == 0 + assert output.input_token_counts[0] > 0 + assert len(model.rows) == 3 + + +class WorksForModel(SpanModel): + """A window holds the works-for relation when it reads both the person and the employer.""" + + def batch_extract_relations( + self, texts: list[str], labels: list[str], *, max_len: int | None, **_: Any + ) -> list[Any]: + self._collate(texts, gliner2_engine.Schema().relations(labels), max_len) + results = [] + for text in texts: + self.rows.append(text) + lower = text.lower() + person, employer = lower.find("priya raman"), lower.find("novartis") + if person < 0 or employer < 0: + results.append({"relation_extraction": {}}) + continue + results.append( + { + "relation_extraction": { + "works_for": [ + { + "head": { + "text": text[person : person + 11], + "confidence": 0.9, + "start": person, + "end": person + 11, + }, + "tail": { + "text": text[employer : employer + 8], + "confidence": 0.8, + "start": employer, + "end": employer + 8, + }, + } + ] + } + } + ) + return results + + +def employment(planted_at: int) -> str: + out = benign_email(1500).split() + out[planted_at : planted_at + 5] = ["Priya", "Raman", "works", "at", "Novartis"] + return " ".join(out) + + +def test_a_relation_past_the_first_window_keeps_the_document_case() -> None: + text = employment(1420) + adapter = GLiNER2Adapter("fake/gliner2", max_seq_length=512) + model = WorksForModel(make_processor(), target="priya raman works at novartis") + adapter._use_linear_word_splitter(model.processor) + adapter._model = model + entities = [ + {"text": "Priya Raman", "start": text.index("Priya Raman"), "end": text.index("Priya Raman") + 11}, + {"text": "Novartis", "start": text.index("Novartis"), "end": text.index("Novartis") + 8}, + ] + + output = adapter.extract([Item(text=text, metadata={"entities": entities})], labels=["works_for"]) + + assert output.errors is None + assert output.relations == [[{"head": "Priya Raman", "tail": "Novartis", "relation": "works_for", "score": 0.8}]] + + +class EmployerModel(PackageModel): + """A window names the employer when it reads the person.""" + + def batch_extract_json( + self, texts: list[str], structures: dict[str, list[str]], *, max_len: int | None, **_: Any + ) -> list[Any]: + self._collate(texts, gliner2_engine.Schema().structure("_sie_root").field("employer", dtype="str"), max_len) + results = [] + for text in texts: + if "priya raman" in text.lower(): + results.append({"_sie_root": [{"employer": "Novartis"}]}) + else: + results.append({"_sie_root": [{"employer": None}]}) + return results + + +def test_a_structured_field_past_the_first_window_is_extracted() -> None: + text = document(1500, planted_at=1420) + adapter = GLiNER2Adapter("fake/gliner2", max_seq_length=512) + model = EmployerModel(make_processor()) + adapter._use_linear_word_splitter(model.processor) + adapter._model = model + + output = adapter.extract( + [Item(text=text)], + output_schema={"type": "object", "properties": {"employer": {"type": "string"}}}, + ) + + assert output.errors is None + assert output.data == [{"employer": "Novartis"}]