fix(server): limit the label prompt of GLiNER-family extract requests

packages/sie_server/src/sie_server/adapters/_prompt_limit.py

@@ -0,0 +1,131 @@
+"""Limits on the task prompt GLiNER-family adapters encode with each document.
+
+GLiNER, GLiNER2 and GLiREL encode a request's labels, relation types, class
+labels and schema fields (its task prompt) with every document. The prompt is
+not billed (input tokens count the document), so it is bounded instead, as
+GLiFormer and GLiNER2.5-Decide bound theirs. A request is rejected with
+``InvalidInputError`` (HTTP 400 ``INVALID_INPUT``) when
+
+* a label, relation type, class label, task name, field name or choice has
+  more than ``MAX_LABEL_CHARS`` characters (checked first, before anything is
+  tokenized), or
+* the prompt takes more than ``max_prompt_tokens`` tokens, counted with the
+  model's tokenizer.
+
+The defaults sit far above the label sets these models are used with: the
+largest label set among this repository's examples takes about 40 tokens, a
+60-type PII list about 230, and a structured-extraction schema of 50 described
+fields about 1,200, while ``DEFAULT_MAX_PROMPT_TOKENS`` holds several hundred
+entity types and ``DEFAULT_MAX_SCHEMA_PROMPT_TOKENS`` a schema of dozens of
+described fields with their choices.
+
+Descriptions are not limited one by one; the whole prompt is tokenized only
+after its characters are checked, so counting costs at most
+``MAX_PROMPT_CHARS_PER_TOKEN`` characters of tokenization per allowed token.
+"""
+
+from __future__ import annotations
+
+import hashlib
+from collections import OrderedDict
+from collections.abc import Callable, Hashable, Iterable
+from typing import Any
+
+from sie_server.types.inputs import InvalidInputError
+
+# Characters a label, relation type, class label, task name, field name or choice may have.
+MAX_LABEL_CHARS = 128
+# Tokens a request's labels, relation types and class labels may take.
+DEFAULT_MAX_PROMPT_TOKENS = 1024
+# Tokens a GLiNER2 request's labels or schema (field names, descriptions and choices) may take.
+DEFAULT_MAX_SCHEMA_PROMPT_TOKENS = 2048
+# No token of these tokenizers covers more characters than this.
+MAX_PROMPT_CHARS_PER_TOKEN = 32
+_PROMPT_CACHE_SIZE = 256
+
+
+def validate_max_prompt_tokens(value: object) -> int:
+    """A ``max_prompt_tokens`` adapter option, checked.
+
+    Raises:
+        ValueError: The value is not a positive integer.
+    """
+    if isinstance(value, bool) or not isinstance(value, int) or value < 1:
+        raise ValueError("max_prompt_tokens must be a positive integer")
+    return value
+
+
+def check_label_chars(model: str, kind: str, values: Iterable[str]) -> None:
+    """Reject a label (or relation type, class label, field name, choice) of more than ``MAX_LABEL_CHARS`` characters.
+
+    Raises:
+        InvalidInputError: A value is not a string or is too long.
+    """
+    for value in values:
+        if not isinstance(value, str):
+            raise InvalidInputError(f"{model} {kind} must be strings")
+        if len(value) > MAX_LABEL_CHARS:
+            raise InvalidInputError(f"{model} {kind} may have at most {MAX_LABEL_CHARS} characters each")
+
+
+class PromptLimit:
+    """Checks a request's task prompt against ``max_tokens``, remembering recent prompts' sizes."""
+
+    __slots__ = ("_counts", "max_tokens", "model")
+
+    def __init__(self, model: str, max_tokens: int) -> None:
+        self.model = model
+        self.max_tokens = validate_max_prompt_tokens(max_tokens)
+        self._counts: OrderedDict[bytes, int] = OrderedDict()
+
+    def check(self, texts: Iterable[str], count: Callable[[], int], key: Hashable) -> int:
+        """The prompt's tokens, from ``count()``, after checking ``texts`` (its strings) and the result.
+
+        ``key`` identifies the prompt (its strings, in order, and anything else
+        its size depends on); a prompt seen recently is not counted again.
+
+        Raises:
+            InvalidInputError: The prompt takes more than ``max_tokens`` tokens.
+        """
+        digest = hashlib.sha256(repr(key).encode("utf-8", "surrogatepass")).digest()
+        tokens = self._counts.get(digest)
+        if tokens is None:
+            chars = sum(len(text) for text in texts)
+            if chars > self.max_tokens * MAX_PROMPT_CHARS_PER_TOKEN:
+                raise InvalidInputError(self._message(None, chars))
+            tokens = int(count())
+            self._counts[digest] = tokens
+            if len(self._counts) > _PROMPT_CACHE_SIZE:
+                self._counts.popitem(last=False)
+        else:
+            self._counts.move_to_end(digest)
+        if tokens > self.max_tokens:
+            raise InvalidInputError(self._message(tokens, None))
+        return tokens
+
+    def _message(self, tokens: int | None, chars: int | None) -> str:
+        size = f"{tokens} tokens" if tokens is not None else f"{chars} characters"
+        return (
+            f"{self.model} labels, relation types, class labels and schema fields take {size}; "
+            f"a request may use at most {self.max_tokens} tokens for them"
+        )
+
+
+def gliner_prompt_counter(model: Any) -> Callable[[list[str], list[str]], int]:
+    """Tokens of the prompt a loaded ``gliner`` model builds for entity and relation types.
+
+    The prompt is built by the model's own processor (``prepare_inputs``, with
+    no document words) and tokenized as the processor tokenizes it.
+    """
+    processor = model.data_processor
+    tokenizer = processor.transformer_tokenizer
+
+    def count(entity_types: list[str], relation_types: list[str]) -> int:
+        kwargs = {"relations": relation_types} if relation_types else {}
+        (words,), _ = processor.prepare_inputs([[]], entity_types, **kwargs)
+        if not words:
+            return 0
+        encoding = tokenizer(list(words), is_split_into_words=True, add_special_tokens=False)
+        return len(encoding["input_ids"])
+
+    return count

packages/sie_server/src/sie_server/adapters/gliner/__init__.py

@@ -17,6 +17,11 @@
 Joint entity-relation ("relex") models also extract relations between the
 entities they find when a request names relation types in
 ``options["relation_labels"]``. Without it they return entities only.
+
+A request's labels and relation types, which GLiNER encodes with every
+document and does not bill, may have at most 128 characters each and take at
+most ``max_prompt_tokens`` tokens together (default 1024); a longer prompt is
+rejected with ``INVALID_INPUT``.
 """
 
 import math
@@ -27,6 +32,12 @@
 import torch
 
 from sie_server.adapters._base_adapter import BaseAdapter
+from sie_server.adapters._prompt_limit import (
+    DEFAULT_MAX_PROMPT_TOKENS,
+    PromptLimit,
+    check_label_chars,
+    gliner_prompt_counter,
+)
 from sie_server.adapters._spec import AdapterSpec
 from sie_server.adapters._types import ERR_REQUIRES_TEXT, ComputePrecision
 from sie_server.adapters._word_window import bound_gliner_words, plan_forwards
@@ -96,6 +107,7 @@ def __init__(
         multi_label: bool = False,
         merge_adjacent_entities: bool = False,
         relation_threshold: float | None = None,
+        max_prompt_tokens: int = DEFAULT_MAX_PROMPT_TOKENS,
         compute_precision: ComputePrecision = "float16",
         revision: str | None = None,
         **kwargs: Any,  # Accept extra args from loader (e.g., pooling)
@@ -112,6 +124,9 @@ def __init__(
             relation_threshold: Minimum relation score (0-1) for joint
                 entity-relation models. None uses the entity threshold, as the
                 gliner library does.
+            max_prompt_tokens: Most tokens a request's labels and relation
+                types may take in the prompt encoded with each document (see
+                ``_prompt_limit``).
             compute_precision: Compute precision for inference.
             revision: Optional HuggingFace revision/branch/commit SHA to pin when
                 loading model artifacts.
@@ -133,6 +148,9 @@ def __init__(
         self._extracts_relations = False
         # True when the encoder's attention memory grows with the square of a row (see ``_inference``).
         self._quadratic_attention = False
+        self._prompt_limit = PromptLimit("GLiNER", max_prompt_tokens)
+        # Tokens of the label prompt, as the loaded model builds it; None until loaded.
+        self._count_prompt: Any = None
 
     def load(self, device: str) -> None:
         """Load the model onto the specified device.
@@ -173,6 +191,7 @@ def load(self, device: str) -> None:
         # gliner's max_len counts words, whatever their subwords: read at most a
         # bounded number of subwords too, with a long word in pieces.
         self._quadratic_attention = bound_gliner_words(self._model)
+        self._count_prompt = gliner_prompt_counter(self._model)
 
     def extract(
         self,
@@ -227,6 +246,8 @@ def extract(
         if relation_labels and not self._extracts_relations:
             raise InvalidInputError(_ERR_NO_RELATIONS)
 
+        self._check_prompt(labels, relation_labels)
+
         # Extract texts from all items
         texts = [self._extract_text(item) for item in items]
         if any(not text.strip() for text in texts):
@@ -303,6 +324,27 @@ def extract(
 
         return ExtractOutput(entities=all_entities, relations=all_relations, input_token_counts=input_token_counts)
 
+    def _check_prompt(self, labels: list[str], relation_labels: list[str]) -> None:
+        """Reject a request whose labels and relation types take more than ``max_prompt_tokens``.
+
+        gliner encodes the label prompt with every document, and only the
+        document is billed.
+
+        Raises:
+            InvalidInputError: The prompt is too long, or a label is not a string.
+        """
+        check_label_chars("GLiNER", "labels", labels)
+        check_label_chars("GLiNER", "relation_labels", relation_labels)
+        entity_types = list(dict.fromkeys(labels))  # gliner drops repeated labels
+        count = self._count_prompt
+
+        def tokens() -> int:
+            return count(entity_types, relation_labels) if count is not None else 0
+
+        self._prompt_limit.check(
+            [*entity_types, *relation_labels], tokens, (tuple(entity_types), tuple(relation_labels))
+        )
+
     @staticmethod
     def _validate_relation_labels(value: Any, labels: list[str]) -> list[str]:
         """Return the requested relation types (empty when none were asked for)."""

packages/sie_server/src/sie_server/adapters/gliner2/adapter.py

@@ -13,6 +13,7 @@
 from huggingface_hub import snapshot_download
 
 from sie_server.adapters._base_adapter import BaseAdapter
+from sie_server.adapters._prompt_limit import DEFAULT_MAX_SCHEMA_PROMPT_TOKENS, PromptLimit, check_label_chars
 from sie_server.adapters._spec import AdapterSpec
 from sie_server.adapters._types import ERR_REQUIRES_TEXT, ComputePrecision
 from sie_server.adapters._word_window import (
@@ -76,6 +77,12 @@ class GLiNER2Adapter(BaseAdapter):
     - Batch methods cover entities, relations, structured data, and classification
     - Classification uses ``classify_text()`` / ``batch_classify_text()``
 
+    A request's labels, class labels, relation types or schema fields (field
+    names, descriptions and choices), which gliner2 encodes with every document
+    and does not bill, may take at most ``max_prompt_tokens`` tokens (default
+    2048), and each label, task name, field name or choice at most 128
+    characters; a longer prompt is rejected with ``INVALID_INPUT``.
+
     Reference models:
     - fastino/gliner2-base-v1
     - fastino/gliner2-large-v1
@@ -98,6 +105,7 @@ def __init__(
         default_labels: list[str] | None = None,
         multi_label: bool = False,
         max_seq_length: int | None = None,
+        max_prompt_tokens: int = DEFAULT_MAX_SCHEMA_PROMPT_TOKENS,
         compute_precision: ComputePrecision = "float16",
         revision: str | None = None,
         **kwargs: Any,
@@ -115,6 +123,9 @@ def __init__(
             multi_label: Whether the configured classification task may return
                 multiple labels.
             max_seq_length: Maximum document and schema input length.
+            max_prompt_tokens: Most tokens a request's labels, class labels,
+                relation types or schema fields may take in the task prompt
+                encoded with each document (see ``_prompt_limit``).
             compute_precision: Compute precision for inference.
             revision: Optional HuggingFace revision/branch/commit SHA to pin when
                 loading model artifacts.
@@ -127,6 +138,7 @@ def __init__(
         self._default_labels = self._validate_labels(default_labels) if default_labels is not None else None
         self._multi_label = multi_label
         self._max_seq_length = max_seq_length
+        self._prompt_limit = PromptLimit("GLiNER2", max_prompt_tokens)
         self._compute_precision = compute_precision
         self._revision = revision
 
@@ -254,8 +266,14 @@ def extract(
                 raise ValueError("GLiNER2 structured extraction does not accept classification_task")
             structures = self._json_schema_to_structures(output_schema)
             specs = [spec for fields in structures.values() for spec in fields]
-            # A field's choices are read twice: in its structure and in a prefix before the document.
-            rows = self._row_tokens(windows, specs + specs)
+            # A field's choices are listed in its structure, and each again in a prefix before the document.
+            choices = [
+                choice for definition in output_schema["properties"].values() for choice in definition.get("enum") or []
+            ]
+            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 = self._run_planned(
                     model_texts,
@@ -286,7 +304,11 @@ def extract(
             normalized_entities = [
                 self._normalize_input_entities(item, entities or []) for item, entities in zip(items, relation_entities)
             ]
-            rows = self._row_tokens(windows, normalized_labels, per_entry=_PROMPT_TOKENS_PER_RELATION)
+            check_label_chars("GLiNER2", "labels", normalized_labels)
+            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 = self._run_planned(
                     model_texts,
@@ -314,14 +336,20 @@ def extract(
         if classification_task is not None:
             if not isinstance(classification_task, str) or not classification_task.strip():
                 raise ValueError("GLiNER2 classification_task must be a non-empty string")
+            check_label_chars("GLiNER2", "classification_task", [classification_task])
+            check_label_chars("GLiNER2", "labels", normalized_labels)
+            prompt = self._prompt_tokens(
+                [classification_task, *normalized_labels],
+                key=("classification", classification_task, tuple(normalized_labels)),
+            )
             return self._classify(
                 model_texts,
                 normalized_labels,
                 task=classification_task,
                 multi_label=multi_label,
                 threshold=effective_threshold,
                 input_token_counts=input_token_counts,
-                rows=self._row_tokens(windows, [classification_task, *normalized_labels]),
+                rows=self._row_tokens(windows, prompt),
             )
 
         def extract_entities(batch: list[str]) -> list[Any]:
@@ -345,10 +373,12 @@ def extract_entities(batch: list[str]) -> list[Any]:
                 max_len=self._max_seq_length,
             )
 
+        check_label_chars("GLiNER2", "labels", normalized_labels)
+        prompt = self._prompt_tokens(normalized_labels, key=("entities", tuple(normalized_labels)))
         with torch.inference_mode():
             raw_results = self._run_planned(
                 model_texts,
-                self._row_tokens(windows, normalized_labels),
+                self._row_tokens(windows, prompt),
                 extract_entities,
                 rows_per_pass=1 if len(texts) == 1 else _PACKAGE_BATCH_SIZE,
             )
@@ -412,22 +442,43 @@ def classify(batch: list[str]) -> list[Any]:
             input_token_counts=input_token_counts,
         )
 
-    def _row_tokens(
+    def _prompt_tokens(
         self,
-        windows: list[tuple[str, int | None]],
-        prompt_entries: Iterable[str],
+        entries: list[str],
         *,
+        key: tuple[Any, ...],
         per_entry: int = _PROMPT_TOKENS_PER_ENTRY,
-    ) -> list[int] | None:
+        extra: int = 0,
+    ) -> int | None:
+        """Estimated tokens of the task prompt gliner2 builds from ``entries``, checked against the limit.
+
+        Each label, class label, relation type and schema field is counted
+        with the tokens gliner2 adds around it, plus ``extra`` (a token per
+        field choice, which gliner2 lists again before the document). None when words are
+        not counted (no bounded splitter is installed); the prompt's
+        characters are still checked.
+
+        Raises:
+            InvalidInputError: The prompt takes more than ``max_prompt_tokens``.
+        """
+        count = self._count_subwords
+        strings = [entry for entry in entries if isinstance(entry, str)]
+
+        def tokens() -> int:
+            if count is None:
+                return 0
+            return sum(count(strings)) + per_entry * len(strings) + extra + _ROW_OVERHEAD_TOKENS
+
+        prompt = self._prompt_limit.check(strings, tokens, (per_entry, extra, *key))
+        return prompt if count is not None else None
+
+    def _row_tokens(self, windows: list[tuple[str, int | None]], prompt: int | None) -> list[int] | None:
         """Estimated tokens of each item's encoder row: the task prompt, then the words it reads.
 
         None when the words were not counted (no bounded splitter is installed).
         """
-        count = self._count_subwords
-        if count is None or any(subwords is None for _, subwords in windows):
+        if prompt is None or any(subwords is None for _, subwords in windows):
             return None
-        entries = [entry for entry in prompt_entries if isinstance(entry, str)]
-        prompt = sum(count(entries)) + per_entry * len(entries) + _ROW_OVERHEAD_TOKENS
         return [prompt + (subwords or 0) for _, subwords in windows]
 
     def _run_planned(