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"}]