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

GLiNER, GLiNER2 and GLiREL now validate the label prompt before any model work. Each label, relation type, class label, task name, schema field and choice is limited to 128 characters, and the whole prompt to a token limit counted with the model's tokenizer: 1024 for GLiNER and GLiREL, 2048 for GLiNER2. Requests over a limit get 400 INVALID_INPUT. The limit is configurable per model with max_prompt_tokens. GLiREL accepts at most 256 supplied entities per item, and rejects malformed input with INVALID_INPUT before any item runs.

Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>

Author: svonava

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(