From 45e24307c21765ea400910d9e060ea21dfdebc8b Mon Sep 17 00:00:00 2001 From: Ingvar Date: Tue, 8 Sep 2026 11:47:38 +0300 Subject: [PATCH 1/2] fix inference issues --- docs/convert_to_onnx.md | 9 +- docs/instalation.md | 7 +- docs/intro.md | 13 +- docs/offline.md | 35 ++++ docs/serving.md | 6 +- docs/streaming.md | 7 + docs/training.md | 9 +- docs/usage.md | 48 ++++- gliner/decoding/decoder.py | 18 +- gliner/model.py | 33 ++- gliner/modeling/cache.py | 2 + tests/test_empty_ner.py | 162 +++++++++++++++ tests/test_local_files_only.py | 228 +++++++++++++++++++-- tests/test_streaming_label_descriptions.py | 122 +++++++++++ tests/test_token_padding.py | 74 +++++++ tests/test_trainer_checkpoints.py | 128 ++++++++++++ tests/test_training_args.py | 48 +++++ 17 files changed, 900 insertions(+), 49 deletions(-) create mode 100644 docs/offline.md create mode 100644 tests/test_empty_ner.py create mode 100644 tests/test_streaming_label_descriptions.py create mode 100644 tests/test_token_padding.py create mode 100644 tests/test_trainer_checkpoints.py create mode 100644 tests/test_training_args.py diff --git a/docs/convert_to_onnx.md b/docs/convert_to_onnx.md index b816e7e2..dda88503 100644 --- a/docs/convert_to_onnx.md +++ b/docs/convert_to_onnx.md @@ -16,16 +16,16 @@ automatically; keep an OpenVINO `.xml` file beside its matching `.bin` file. ## Installation -ONNX Runtime CPU support is included in the default installation: +For ONNX Runtime CPU support, install the optional `onnx` extra: ```bash -pip install gliner +pip install "gliner[onnx]" ``` Install the Python `onnx` package when exporting a PyTorch checkpoint to ONNX: ```bash -pip install gliner onnx +pip install "gliner[onnx]" onnx ``` For ONNX Runtime with CUDA execution providers, install the GPU extra: @@ -35,6 +35,9 @@ pip install "gliner[gpu]" ``` Add `onnx` to the same command if this environment will also export models. +Choose only one of `gliner[onnx]` and `gliner[gpu]`: both runtime packages +provide the same Python module and must not be installed together. The default +`gliner` installation does not install either ONNX Runtime package. For direct OpenVINO export and inference, install the OpenVINO extra. Direct OpenVINO conversion does not create an intermediate ONNX model: diff --git a/docs/instalation.md b/docs/instalation.md index 98ba6e89..a36bd7c5 100644 --- a/docs/instalation.md +++ b/docs/instalation.md @@ -9,12 +9,15 @@ To begin using the GLiNER model, you can install the GLiNER Python library throu pip install gliner ``` :::tip[ONNX Runtime] -CPU inference with ONNX Runtime is included in the default installation. Install the separate `onnx` package to export models, or use the GPU feature for CUDA execution providers. +ONNX Runtime is optional. Choose the `onnx` extra for CPU inference or the `gpu` extra for CUDA execution providers. Install only one: the CPU and GPU packages share the same Python module. Install the separate `onnx` package to export models. ::: ```bash -pip install onnx # required for ONNX export +pip install "gliner[onnx]" # CPU-backed ONNX Runtime +# or pip install "gliner[gpu]" # GPU-backed ONNX Runtime + +pip install onnx # additionally required for ONNX export ``` :::tip[OpenVINO] diff --git a/docs/intro.md b/docs/intro.md index 925d5b4c..f3028c20 100644 --- a/docs/intro.md +++ b/docs/intro.md @@ -85,13 +85,16 @@ To begin using the GLiNER model, you can install the GLiNER Python library throu !pip install gliner ``` -CPU ONNX Runtime support is included in the default installation. Install -`onnx` when exporting models, or install GLiNER with the GPU feature for CUDA -execution providers: +ONNX Runtime is optional. Choose the `onnx` extra for CPU inference or the +`gpu` extra for CUDA execution providers. Install only one: the CPU and GPU +packages share the same Python module. Also install `onnx` when exporting models: ```bash -!pip install onnx -!pip install "gliner[gpu]" +!pip install "gliner[onnx]" # CPU +# or +!pip install "gliner[gpu]" # NVIDIA GPU + +!pip install onnx # additionally required for ONNX export ``` Install OpenVINO support to export models directly to OpenVINO IR or run them diff --git a/docs/offline.md b/docs/offline.md new file mode 100644 index 00000000..2fa6be5d --- /dev/null +++ b/docs/offline.md @@ -0,0 +1,35 @@ +# Loading models offline + +Prepare a complete model directory on a machine with internet access: + +```python +from gliner import GLiNER + +model = GLiNER.from_pretrained("urchade/gliner_multi-v2.1") +model.save_pretrained("gliner-offline", safe_serialization=True) +``` + +Copy the entire `gliner-offline` directory to the offline machine, then load it: + +```python +from gliner import GLiNER + +model = GLiNER.from_pretrained("gliner-offline", local_files_only=True) +``` + +`save_pretrained` includes model weights, the resolved backbone configuration, +and tokenizers. Models with a separate label encoder or generative decoder also +save their auxiliary tokenizer in `labels_tokenizer/` or `decoder_tokenizer/`. +Keep these subdirectories with the model when copying it. + +Older Hub checkpoints may contain a GLiNER configuration that only names the +backbone, such as `microsoft/mdeberta-v3-base`, without embedding its configuration. +Downloading that checkpoint's files alone may therefore be insufficient. Loading +and saving it with the code above resolves and packages those dependencies. + +`local_files_only=True` restricts loading to local files and the Hugging Face +cache; it does not download missing dependencies. An incomplete legacy checkpoint +still requires its missing backbone configuration or tokenizer to be cached or +provided locally. If the backbone files are in a separate local directory, +`model_name` in `gliner_config.json` can point to that directory. Missing files +raise an error without attempting a network connection. diff --git a/docs/serving.md b/docs/serving.md index 16ef029c..16098705 100644 --- a/docs/serving.md +++ b/docs/serving.md @@ -80,7 +80,11 @@ outputs = client.predict( ``` Each dictionary key is returned as the entity label; its value is the prompt encoded -for that text. A single dictionary is shared across every input text. +for that text. A single dictionary is shared across every input text. When supplying +a list of dictionaries, its length and order must match the input texts, and each +dictionary contains all labels for its corresponding text. For a single text, put +all labels in one dictionary. See [Label descriptions](usage.md#label-descriptions) +for examples of shared and per-text label sets. Network or server errors surface as `gliner.serve.client.GLiNERClientError`. diff --git a/docs/streaming.md b/docs/streaming.md index 5af8da12..5eb4184d 100644 --- a/docs/streaming.md +++ b/docs/streaming.md @@ -45,6 +45,13 @@ finally: model.clear_session(session_id) ``` +For `model.inference(..., session_id=[...])`, labels can also be a single +`{"label name": "description"}` dictionary shared by all input chunks, or a list +of dictionaries with one complete label set per input chunk/session. See +[Label descriptions](usage.md#label-descriptions) for the input formats. Keep each +session's label names, descriptions, and order consistent across calls; changing +them requires `recompute=True` or clearing that session. + Each `snapshot` is the complete set of entities currently active for the accumulated session text. It is not a list of only the entities detected in the latest chunk. An entity has the same shape as an ordinary GLiNER prediction: diff --git a/docs/training.md b/docs/training.md index 7695016e..cde2f7c9 100644 --- a/docs/training.md +++ b/docs/training.md @@ -89,6 +89,13 @@ train_data = [ ### Advanced Format: Explicit Labels and Negatives +Examples with `"ner": []` are supported. If no candidate entity types are +available, the processor uses one empty-string label (`""`) with all-zero +targets. For uni-encoders, this produces an entity marker without label text. +Bi-encoders use the empty-string label only when the shared batch label pool is +empty. To train a negative example against specific entity types, supply +`ner_labels` or `ner_negatives` as described below. + You can improve training by explicitly defining: 1. **Positive labels** (`ner_labels`): Entity types present in this example, recomened to specify for fixed label set use cases. 2. **Negative labels** (`ner_negatives`): Entity types to use as negative examples @@ -377,4 +384,4 @@ python train.py --config config.yaml 7. **Save multiple checkpoints**: Set `save_total_limit` > 1 to keep best models 8. **Monitor training**: Use TensorBoard or W&B to track loss and metrics 9. **Start small**: Test your pipeline on a subset before full training -10. **Validate data format**: Ensure indices are correct and within bounds \ No newline at end of file +10. **Validate data format**: Ensure indices are correct and within bounds diff --git a/docs/usage.md b/docs/usage.md index 174b2fe7..57f3f5bf 100644 --- a/docs/usage.md +++ b/docs/usage.md @@ -178,29 +178,61 @@ for i, entities in enumerate(all_entities): ### Label descriptions -Models trained to use descriptive labels can receive a dictionary. Dictionary keys are -returned in predictions, while values are encoded as the label prompts: +Models trained to use descriptive labels can receive a dictionary mapping label names +to descriptions. Dictionary keys are returned in predictions, while values are encoded +as the label prompts. For a single text, put all labels in **one dictionary**: ```python +text = "Microsoft was founded by Bill Gates and Paul Allen." labels = { "person": "A human individual, including fictional characters", "organization": "A company, institution, agency, or other group of people", } entities = model.predict_entities(text, labels) + +for entity in entities: + print(entity["text"], "=>", entity["label"]) +``` + +For a batch that shares the same labels and descriptions, pass that dictionary directly: + +```python +texts = ["Alice works at Microsoft", "Bob works at Google"] +all_entities = model.inference(texts, labels) ``` -For batched inference, provide one dictionary or list of labels per text: +For **different labels or descriptions per input text**, pass a list containing one +dictionary per text. Each dictionary contains the complete label set for its text: ```python +texts = ["Alice works at Microsoft", "Paris is sunny"] label_sets = [ - {"person": "A human individual"}, - {"location": "A geographical place"}, + { + "person": "A human individual", + "organization": "A company or institution", + }, # Labels for texts[0]. + { + "location": "A geographical place", + }, # Labels for texts[1]. ] -entities = model.inference(["Alice arrived", "Paris is sunny"], label_sets) +all_entities = model.inference(texts, label_sets, batch_size=2) + +for text, entities in zip(texts, all_entities): + print(text) + for entity in entities: + print(entity["text"], "=>", entity["label"]) ``` -Descriptions must be unique within each label set. They are not supported with -precomputed prompt embeddings. +The outer list must have the same length and order as `texts`, including entries for +empty texts. A list such as `[{"person": "A human individual"}, {"location": "A geographical place"}]` +means two texts with one label each; it does not mean two labels for one text. To use +both labels with `predict_entities(text, ...)`, combine them into a single dictionary: +`{"person": "A human individual", "location": "A geographical place"}`. + +Without descriptions, use a shared list of label strings or one list of strings per +text. Description dictionaries must have string keys and values, and descriptions +must be unique within each dictionary. They are not supported with precomputed +prompt embeddings. ## Using Different Model Architectures diff --git a/gliner/decoding/decoder.py b/gliner/decoding/decoder.py index b7287d0e..2345dce7 100644 --- a/gliner/decoding/decoder.py +++ b/gliner/decoding/decoder.py @@ -1807,16 +1807,22 @@ class IDs to class names. multi_label_values = _expand_batch_param(multi_label, batch_size, "multi_label") spans = [] - for i, _ in enumerate(tokens): + for i, tokens_i in enumerate(tokens): id_to_class_i = id_to_class_per_item[i] input_spans_i = set(input_spans[i]) if input_spans is not None else None threshold_i = thresholds[i] + # Batch padding can exceed the threshold, so exclude it before + # matching boundaries or selecting overlapping spans. + token_count = len(tokens_i) + start_i = scores_start[i, :token_count] + end_i = scores_end[i, :token_count] + inside_i = scores_inside[i, :token_count] span_scores = self._calculate_span_score( - self._get_indices_above_threshold(scores_start[i], threshold_i), - self._get_indices_above_threshold(scores_end[i], threshold_i), - torch.sigmoid(scores_inside[i]), - torch.sigmoid(scores_start[i]), - torch.sigmoid(scores_end[i]), + self._get_indices_above_threshold(start_i, threshold_i), + self._get_indices_above_threshold(end_i, threshold_i), + torch.sigmoid(inside_i), + torch.sigmoid(start_i), + torch.sigmoid(end_i), id_to_class_i, threshold_i, input_spans_i=input_spans_i, diff --git a/gliner/model.py b/gliner/model.py index b9bd6011..acdd2540 100644 --- a/gliner/model.py +++ b/gliner/model.py @@ -4142,6 +4142,7 @@ def _state_from_output( *, batch_index=0, past_key_values=None, + label_names=None, ): current_attention = ( batch["label_attention_mask"] if "label_attention_mask" in batch else batch["attention_mask"] @@ -4172,6 +4173,7 @@ def _state_from_output( next_position_id=cached_length, session_id=session_id, labels=tuple(labels), + label_names=tuple(label_names) if label_names is not None else None, text=text, tokens=list(tokens), char_starts=list(char_starts), @@ -4196,7 +4198,8 @@ def _decode_session_output( if not state.span_logits: return [] - id_to_classes = {index + 1: label for index, label in enumerate(state.labels)} + label_names = state.label_names if state.label_names is not None else state.labels + id_to_classes = {index + 1: label for index, label in enumerate(label_names)} ordered_scores = sorted(state.span_logits.items()) span_idx = torch.tensor( [[boundary for boundary, _ in ordered_scores]], @@ -4234,11 +4237,16 @@ def _decode_session_output( entities.append(entity) return entities - def _prepare_session_item(self, text, labels, session_id, recompute): + def _prepare_session_item(self, text, labels, session_id, recompute, label_names=None): """Prepare semantic and token metadata for one streaming append.""" + label_names = tuple(labels if label_names is None else label_names) state = self._session_cache.get(session_id) - if state is not None and state.labels != tuple(labels) and not recompute: - raise ValueError(f"Labels for session {session_id!r} changed. Pass recompute=True or clear the session.") + if state is not None and not recompute: + previous_names = state.label_names if state.label_names is not None else state.labels + if state.labels != tuple(labels) or previous_names != label_names: + raise ValueError( + f"Labels for session {session_id!r} changed. Pass recompute=True or clear the session." + ) current_tokens, current_starts, current_ends = self.prepare_inputs([text]) current_tokens = current_tokens[0] @@ -4248,6 +4256,7 @@ def _prepare_session_item(self, text, labels, session_id, recompute): return { "session_id": session_id, "labels": tuple(labels), + "label_names": label_names, "state": state, "current_tokens": current_tokens, "current_starts": current_starts, @@ -4303,6 +4312,7 @@ def _run_full_session_batch(self, items): item["combined_ends"], batch_index=row, past_key_values=row_cache, + label_names=item["label_names"], ) ) for state in states: @@ -4772,6 +4782,7 @@ def _run_incremental_session_batch(self, items): next_position_id=cached.next_position + current_token_count, session_id=item["session_id"], labels=item["labels"], + label_names=item["label_names"], text=item["combined_text"], tokens=item["combined_tokens"], char_starts=item["combined_starts"], @@ -4807,6 +4818,7 @@ def _run_session_items_batched(self, requests, *, return_exceptions=False): request["labels"], request["session_id"], request.get("recompute", False), + label_names=request.get("label_names"), ) except Exception as error: results[index] = error @@ -5006,7 +5018,7 @@ def _run_session_item( def inference( self, texts: str | List[str], - labels: List[str], + labels: EntityLabels, flat_ner: bool = True, threshold: float = 0.5, multi_label: bool = False, @@ -5028,6 +5040,9 @@ def inference( takes precedence and refreshes every span over the combined session. Returned entities are a complete session snapshot, including any revisions or removals caused by the newly appended context. + Labels support shared or per-text label-to-description mappings, just + as in stateless inference. Changing names or descriptions in an existing + session requires ``recompute=True`` or clearing the session. """ if session_id is None: return super().inference( @@ -5066,8 +5081,9 @@ def inference( if not isinstance(recompute, bool): raise TypeError("recompute must be a boolean") - normalized_labels = list(dict.fromkeys(labels)) - if not normalized_labels: + normalized_labels, label_names = _normalize_labels(labels, len(texts), list(range(len(texts)))) + per_text_labels = bool(normalized_labels) and isinstance(normalized_labels[0], list) + if not normalized_labels or (per_text_labels and any(not label_set for label_set in normalized_labels)): raise ValueError("At least one label is required") if not isinstance(batch_size, int) or batch_size < 1: raise ValueError("batch_size must be a positive integer") @@ -5082,7 +5098,8 @@ def inference( requests.append( { "text": text, - "labels": normalized_labels, + "labels": normalized_labels[index] if per_text_labels else normalized_labels, + "label_names": label_names[index] if per_text_labels else label_names, "session_id": current_session, "recompute": recompute, "threshold": threshold, diff --git a/gliner/modeling/cache.py b/gliner/modeling/cache.py index 2493c2c5..caae3392 100644 --- a/gliner/modeling/cache.py +++ b/gliner/modeling/cache.py @@ -240,6 +240,8 @@ class CacheState: char_ends: list[int] = field(default_factory=list) span_logits: dict[tuple[int, int], torch.Tensor] = field(default_factory=dict) metadata: dict[str, Any] = field(default_factory=dict) + # Public names may differ from the model-facing prompts stored in labels. + label_names: tuple[str, ...] | None = None @property def next_position(self) -> int: diff --git a/tests/test_empty_ner.py b/tests/test_empty_ner.py new file mode 100644 index 00000000..e8350d9c --- /dev/null +++ b/tests/test_empty_ner.py @@ -0,0 +1,162 @@ +"""Regression coverage for training on empty annotations (issue #320).""" + +from copy import deepcopy + +import torch +import pytest +from tokenizers import Tokenizer +from transformers import BertConfig, PreTrainedTokenizerFast +from tokenizers.models import WordLevel +from tokenizers.processors import BertProcessing +from tokenizers.pre_tokenizers import WhitespaceSplit + +from gliner.config import ( + BiEncoderSpanConfig, + BiEncoderTokenConfig, + UniEncoderSpanConfig, + UniEncoderTokenConfig, + UniEncoderSpanRelexConfig, +) +from gliner.modeling.base import ( + BiEncoderSpanModel, + BiEncoderTokenModel, + UniEncoderSpanModel, + UniEncoderTokenModel, + UniEncoderSpanRelexModel, +) +from gliner.data_processing.collator import ( + SpanDataCollator, + TokenDataCollator, + RelationExtractionSpanDataCollator, +) +from gliner.data_processing.processor import ( + BiEncoderSpanProcessor, + BiEncoderTokenProcessor, + UniEncoderSpanProcessor, + UniEncoderTokenProcessor, + RelationExtractionSpanProcessor, +) + + +@pytest.fixture( + params=[ + (UniEncoderSpanConfig, UniEncoderSpanProcessor, SpanDataCollator, UniEncoderSpanModel), + (UniEncoderTokenConfig, UniEncoderTokenProcessor, TokenDataCollator, UniEncoderTokenModel), + (BiEncoderSpanConfig, BiEncoderSpanProcessor, SpanDataCollator, BiEncoderSpanModel), + (BiEncoderTokenConfig, BiEncoderTokenProcessor, TokenDataCollator, BiEncoderTokenModel), + ( + UniEncoderSpanRelexConfig, + RelationExtractionSpanProcessor, + RelationExtractionSpanDataCollator, + UniEncoderSpanRelexModel, + ), + ], + ids=["uni_span", "uni_token", "bi_span", "bi_token", "relex_span"], +) +def setup(request): + config_cls, processor_cls, collator_cls, model_cls = request.param + vocabulary = [ + "[PAD]", + "[UNK]", + "[CLS]", + "[SEP]", + "<>", + "<>", + "<>", + "The", + "weather", + "is", + "nice", + ".", + "Alice", + "person", + ] + vocab = {word: index for index, word in enumerate(vocabulary)} + backend = Tokenizer(WordLevel(vocab, unk_token="[UNK]")) + backend.pre_tokenizer = WhitespaceSplit() + backend.post_processor = BertProcessing(("[SEP]", vocab["[SEP]"]), ("[CLS]", vocab["[CLS]"])) + tokenizer = PreTrainedTokenizerFast( + tokenizer_object=backend, + unk_token="[UNK]", + pad_token="[PAD]", + cls_token="[CLS]", + sep_token="[SEP]", + model_max_length=64, + additional_special_tokens=["<>", "<>", "<>"], + ) + encoder_config = BertConfig( + vocab_size=len(tokenizer), + hidden_size=16, + num_hidden_layers=1, + num_attention_heads=2, + intermediate_size=32, + ).to_dict() + kwargs = {} + if processor_cls in (BiEncoderSpanProcessor, BiEncoderTokenProcessor): + kwargs = {"labels_encoder": "local-bert", "labels_encoder_config": deepcopy(encoder_config)} + config = config_cls( + model_name="local-bert", + encoder_config=encoder_config, + hidden_size=16, + vocab_size=len(tokenizer), + class_token_index=vocab["<>"], + rel_token_index=vocab["<>"], + num_rnn_layers=0, + max_width=2, + max_neg_type_ratio=0, + dropout=0.0, + augment_data_prob=0.0, + **kwargs, + ) + args = [config, tokenizer, None] + if kwargs: + args.append(tokenizer) + processor = processor_cls(*args) + return config, processor, collator_cls(config, data_processor=processor), model_cls, tokenizer + + +@pytest.mark.parametrize("case", ["single", "all_empty", "mixed", "explicit_negatives", "empty_labels"]) +def test_empty_examples_produce_training_gradients(setup, case): + config, processor, collator, model_cls, tokenizer = setup + examples = [{"tokenized_text": ["The", "weather", "is", "nice", "."], "ner": [], "relations": []}] + if case == "all_empty": + examples.append({"tokenized_text": ["nice", "."], "ner": [], "relations": []}) + elif case == "mixed": + examples.append({"tokenized_text": ["Alice"], "ner": [[0, 0, "person"]], "relations": []}) + elif case == "explicit_negatives": + examples[0]["ner_negatives"] = ["person"] + elif case == "empty_labels": + examples[0]["ner_labels"] = [] + original = deepcopy(examples) + raw = processor.collate_raw_batch(examples) + candidates = raw["classes_to_id"][0] + bi_encoder = isinstance(processor, (BiEncoderSpanProcessor, BiEncoderTokenProcessor)) + expected = "person" if case == "explicit_negatives" or (case == "mixed" and bi_encoder) else "" + assert candidates == {expected: 1} + if isinstance(processor, RelationExtractionSpanProcessor): + assert all(mapping == {} for mapping in raw["rel_class_to_ids"]) + + batch = collator(examples) + assert torch.count_nonzero(batch["labels"][0]) == 0 + if not bi_encoder and expected == "": + tokens = tokenizer.convert_ids_to_tokens(batch["input_ids"][0].tolist()) + assert tokens[:3] == ["[CLS]", "<>", "<>"] + # The empty label must not shift or discard any text words. + assert batch["words_mask"][0][batch["words_mask"][0] > 0].tolist() == [1, 2, 3, 4, 5] + + model = model_cls(config, from_pretrained=False) + output = model(**batch) + assert torch.isfinite(output.loss) and output.loss > 0 + output.logits.retain_grad() + output.loss.backward() + assert output.logits.grad[0].abs().sum() > 0 + gradients = [parameter.grad for parameter in model.parameters() if parameter.grad is not None] + assert gradients and all(torch.isfinite(gradient).all() for gradient in gradients) + assert examples == original + + +def test_explicit_empty_inference_labels_remain_empty(setup): + _, processor, _, _, _ = setup + example = {"tokenized_text": ["Alice"], "ner": None, "relations": []} + raw = processor.collate_raw_batch([example], entity_types=[[]]) + assert raw["classes_to_id"] == [{}] diff --git a/tests/test_local_files_only.py b/tests/test_local_files_only.py index e43295d2..77a378e4 100644 --- a/tests/test_local_files_only.py +++ b/tests/test_local_files_only.py @@ -1,10 +1,31 @@ """Tests that local_files_only is correctly forwarded through the loading path.""" + +import json +import socket from unittest.mock import MagicMock, patch +import torch import pytest - -from gliner.model import BaseGLiNER -from gliner.config import GLiNERConfig +from tokenizers import Tokenizer +from transformers import BertConfig, GPT2Config, PreTrainedTokenizerFast +from tokenizers.models import WordLevel +from tokenizers.pre_tokenizers import Whitespace + +from gliner import GLiNER +from gliner.model import ( + BaseGLiNER, + BiEncoderSpanGLiNER, + StreamingSpanGLiNER, + UniEncoderSpanGLiNER, + UniEncoderSpanDecoderGLiNER, +) +from gliner.config import ( + GLiNERConfig, + BiEncoderSpanConfig, + StreamingSpanConfig, + UniEncoderSpanConfig, + UniEncoderSpanDecoderConfig, +) @pytest.fixture @@ -26,32 +47,209 @@ class TestLoadTokenizerLocalFilesOnly: def test_with_tokenizer_config_present(self, tmp_path, config, mock_tokenizer): (tmp_path / "tokenizer_config.json").write_text("{}") - with patch("gliner.model.AutoTokenizer.from_pretrained", return_value=mock_tokenizer) as mock_ft: - with patch.object(BaseGLiNER, "_set_tokenizer_spec_tokens", return_value=mock_tokenizer): - BaseGLiNER._load_tokenizer(config, tmp_path, local_files_only=True) + with ( + patch("gliner.model.AutoTokenizer.from_pretrained", return_value=mock_tokenizer) as mock_ft, + patch.object(BaseGLiNER, "_set_tokenizer_spec_tokens", return_value=mock_tokenizer), + ): + BaseGLiNER._load_tokenizer(config, tmp_path, local_files_only=True) mock_ft.assert_called_once_with(tmp_path, cache_dir=None, local_files_only=True) def test_with_tokenizer_config_absent_uses_model_name(self, tmp_path, config, mock_tokenizer): - with patch("gliner.model.AutoTokenizer.from_pretrained", return_value=mock_tokenizer) as mock_ft: - with patch.object(BaseGLiNER, "_set_tokenizer_spec_tokens", return_value=mock_tokenizer): - BaseGLiNER._load_tokenizer(config, tmp_path, local_files_only=True) + with ( + patch("gliner.model.AutoTokenizer.from_pretrained", return_value=mock_tokenizer) as mock_ft, + patch.object(BaseGLiNER, "_set_tokenizer_spec_tokens", return_value=mock_tokenizer), + ): + BaseGLiNER._load_tokenizer(config, tmp_path, local_files_only=True) mock_ft.assert_called_once_with(config.model_name, cache_dir=None, local_files_only=True) def test_default_is_false(self, tmp_path, config, mock_tokenizer): - with patch("gliner.model.AutoTokenizer.from_pretrained", return_value=mock_tokenizer) as mock_ft: - with patch.object(BaseGLiNER, "_set_tokenizer_spec_tokens", return_value=mock_tokenizer): - BaseGLiNER._load_tokenizer(config, tmp_path) + with ( + patch("gliner.model.AutoTokenizer.from_pretrained", return_value=mock_tokenizer) as mock_ft, + patch.object(BaseGLiNER, "_set_tokenizer_spec_tokens", return_value=mock_tokenizer), + ): + BaseGLiNER._load_tokenizer(config, tmp_path) _, kwargs = mock_ft.call_args assert kwargs.get("local_files_only") is False def test_local_false_does_not_block_network(self, tmp_path, config, mock_tokenizer): """Sanity check: local_files_only=False still passes the flag through.""" - with patch("gliner.model.AutoTokenizer.from_pretrained", return_value=mock_tokenizer) as mock_ft: - with patch.object(BaseGLiNER, "_set_tokenizer_spec_tokens", return_value=mock_tokenizer): - BaseGLiNER._load_tokenizer(config, tmp_path, local_files_only=False) + with ( + patch("gliner.model.AutoTokenizer.from_pretrained", return_value=mock_tokenizer) as mock_ft, + patch.object(BaseGLiNER, "_set_tokenizer_spec_tokens", return_value=mock_tokenizer), + ): + BaseGLiNER._load_tokenizer(config, tmp_path, local_files_only=False) _, kwargs = mock_ft.call_args assert kwargs.get("local_files_only") is False + + +@pytest.fixture +def offline_checkpoint(tmp_path, request): + backend = Tokenizer(WordLevel({"[UNK]": 0, "[PAD]": 1, "Alice": 2, "person": 3}, unk_token="[UNK]")) + backend.pre_tokenizer = Whitespace() + tokenizer = PreTrainedTokenizerFast(tokenizer_object=backend, unk_token="[UNK]", pad_token="[PAD]") + encoder_config = BertConfig( + vocab_size=len(tokenizer), + hidden_size=16, + num_hidden_layers=1, + num_attention_heads=2, + intermediate_size=32, + ).to_dict() + config_kwargs = { + "model_name": "offline-test/missing-backbone", + "encoder_config": encoder_config, + "hidden_size": 16, + "max_width": 2, + "num_rnn_layers": 0, + "dropout": 0.0, + } + kind = getattr(request, "param", "uni") + if kind == "bi": + model_class = BiEncoderSpanGLiNER + config = BiEncoderSpanConfig( + **config_kwargs, + labels_encoder="offline-test/missing-labels", + labels_encoder_config=encoder_config.copy(), + ) + elif kind == "decoder": + model_class = UniEncoderSpanDecoderGLiNER + config = UniEncoderSpanDecoderConfig( + **config_kwargs, + labels_decoder="offline-test/missing-decoder", + labels_decoder_config=GPT2Config( + vocab_size=len(tokenizer), + n_embd=16, + n_layer=1, + n_head=2, + ).to_dict(), + ) + elif kind == "streaming": + model_class = StreamingSpanGLiNER + config = StreamingSpanConfig( + **config_kwargs, + decoder_config=GPT2Config(vocab_size=len(tokenizer), n_embd=16, n_layer=1, n_head=2).to_dict(), + labels_encoder_config={"model_type": "rnn"}, + ) + else: + model_class = UniEncoderSpanGLiNER + config = UniEncoderSpanConfig(**config_kwargs) + model = model_class(config, tokenizer=tokenizer, labels_tokenizer=tokenizer, decoder_tokenizer=tokenizer) + BaseGLiNER._resize_token_embeddings(model, config, tokenizer) + model.eval() + checkpoint = tmp_path / "checkpoint" + model.save_pretrained(checkpoint, safe_serialization=True) + return model, checkpoint + + +@pytest.fixture +def no_network(monkeypatch): + def reject_network(*args, **kwargs): + pytest.fail("Offline loading attempted a network connection") + + monkeypatch.setattr(socket.socket, "connect", reject_network) + + +@pytest.mark.parametrize("offline_checkpoint", ["uni", "bi", "decoder", "streaming"], indirect=True) +@pytest.mark.parametrize("low_cpu_mem_usage", [False, True]) +def test_saved_checkpoint_loads_with_empty_cache(offline_checkpoint, tmp_path, no_network, low_cpu_mem_usage): + original, checkpoint = offline_checkpoint + loaded = GLiNER.from_pretrained( + checkpoint, + local_files_only=True, + cache_dir=tmp_path / "empty-cache", + low_cpu_mem_usage=low_cpu_mem_usage, + strict=True, + ) + for name, expected in original.model.state_dict().items(): + torch.testing.assert_close(loaded.model.state_dict()[name], expected) + for name in ("transformer_tokenizer", "labels_tokenizer", "decoder_tokenizer"): + expected = getattr(original.data_processor, name, None) + if expected is not None: + assert getattr(loaded.data_processor, name).get_vocab() == expected.get_vocab() + + +@pytest.mark.parametrize("low_cpu_mem_usage", [False, True]) +def test_legacy_checkpoint_missing_backbone_fails_without_network( + offline_checkpoint, + tmp_path, + no_network, + low_cpu_mem_usage, +): + _, checkpoint = offline_checkpoint + config_file = checkpoint / "gliner_config.json" + config = json.loads(config_file.read_text()) + config.pop("encoder_config") + config_file.write_text(json.dumps(config)) + with pytest.raises(OSError): + GLiNER.from_pretrained( + checkpoint, + local_files_only=True, + cache_dir=tmp_path / "empty-cache", + low_cpu_mem_usage=low_cpu_mem_usage, + ) + + +@pytest.mark.parametrize("load_tokenizer", [True, False]) +def test_missing_tokenizer_fails_without_network(offline_checkpoint, tmp_path, no_network, load_tokenizer): + _, checkpoint = offline_checkpoint + (checkpoint / "tokenizer_config.json").unlink() + with pytest.raises(OSError): + GLiNER.from_pretrained( + checkpoint, + local_files_only=True, + cache_dir=tmp_path / "empty-cache", + load_tokenizer=load_tokenizer, + ) + + +@pytest.mark.parametrize( + "offline_checkpoint,missing_field", + [("bi", "labels_encoder_config"), ("decoder", "labels_decoder_config")], + indirect=["offline_checkpoint"], +) +def test_missing_auxiliary_config_fails_without_network(offline_checkpoint, missing_field, tmp_path, no_network): + _, checkpoint = offline_checkpoint + config_file = checkpoint / "gliner_config.json" + config = json.loads(config_file.read_text()) + config.pop(missing_field) + config_file.write_text(json.dumps(config)) + with pytest.raises(OSError): + GLiNER.from_pretrained(checkpoint, local_files_only=True, cache_dir=tmp_path / "empty-cache") + + +@pytest.mark.parametrize( + "offline_checkpoint,tokenizer_name", + [("bi", "labels_tokenizer"), ("decoder", "decoder_tokenizer")], + indirect=["offline_checkpoint"], +) +def test_missing_auxiliary_tokenizer_fails_without_network(offline_checkpoint, tokenizer_name, tmp_path, no_network): + _, checkpoint = offline_checkpoint + (checkpoint / tokenizer_name / "tokenizer_config.json").unlink() + with pytest.raises(OSError): + GLiNER.from_pretrained(checkpoint, local_files_only=True, cache_dir=tmp_path / "empty-cache") + + +def test_legacy_checkpoint_can_use_local_backbone_config(offline_checkpoint, tmp_path, no_network): + original, checkpoint = offline_checkpoint + backbone_dir = tmp_path / "backbone" + original.config.encoder_config.save_pretrained(backbone_dir) + config_file = checkpoint / "gliner_config.json" + config = json.loads(config_file.read_text()) + config.pop("encoder_config") + config["model_name"] = str(backbone_dir) + config_file.write_text(json.dumps(config)) + loaded = GLiNER.from_pretrained(checkpoint, local_files_only=True, cache_dir=tmp_path / "empty-cache", strict=True) + for name, expected in original.model.state_dict().items(): + torch.testing.assert_close(loaded.model.state_dict()[name], expected) + + +def test_offline_load_preserves_predictions(offline_checkpoint, tmp_path, no_network): + original, checkpoint = offline_checkpoint + loaded = GLiNER.from_pretrained(checkpoint, local_files_only=True, cache_dir=tmp_path / "empty-cache") + expected = original.predict_entities("Alice", ["person"], threshold=0.0) + actual = loaded.predict_entities("Alice", ["person"], threshold=0.0) + assert expected + assert actual == expected diff --git a/tests/test_streaming_label_descriptions.py b/tests/test_streaming_label_descriptions.py new file mode 100644 index 00000000..af992e55 --- /dev/null +++ b/tests/test_streaming_label_descriptions.py @@ -0,0 +1,122 @@ +from unittest.mock import patch + +import pytest +from tokenizers import Tokenizer +from transformers import GPT2Config, PreTrainedTokenizerFast +from tokenizers.models import WordLevel +from tokenizers.pre_tokenizers import Whitespace + +from gliner.model import StreamingSpanGLiNER +from gliner.config import StreamingSpanConfig + + +@pytest.fixture +def model(): + words = ["[UNK]", "[PAD]", "Alice", "Paris", "works", "person", "location", "A", "human", "place"] + backend = Tokenizer(WordLevel({word: index for index, word in enumerate(words)}, unk_token="[UNK]")) + backend.pre_tokenizer = Whitespace() + tokenizer = PreTrainedTokenizerFast(tokenizer_object=backend, unk_token="[UNK]", pad_token="[PAD]") + config = StreamingSpanConfig( + model_name="offline-test/missing-backbone", + decoder_config=GPT2Config( + vocab_size=len(tokenizer), n_embd=16, n_layer=1, n_head=2, bos_token_id=0, eos_token_id=0 + ).to_dict(), + labels_encoder_config={"model_type": "rnn"}, + hidden_size=16, + max_width=2, + num_rnn_layers=0, + dropout=0.0, + ) + model = StreamingSpanGLiNER(config, tokenizer=tokenizer) + model._resize_token_embeddings(model, config, tokenizer) + return model.eval() + + +@pytest.mark.parametrize( + ("labels", "prompts", "names"), + [ + ({"person": "A human", "location": "A place"}, [["A human", "A place"]] * 2, [["person", "location"]] * 2), + ([{"person": "A human"}, {"location": "A place"}], [["A human"], ["A place"]], [["person"], ["location"]]), + (["person", "person", "location"], [["person", "location"]] * 2, [["person", "location"]] * 2), + ([["person"], ["location"]], [["person"], ["location"]], [["person"], ["location"]]), + ("person", [["person"]] * 2, [["person"]] * 2), + ], +) +@pytest.mark.parametrize("batch_size", [1, 2]) +def test_sessions_preserve_prompts_and_public_names(model, labels, prompts, names, batch_size): + for chunks, recompute in [(["Alice", "Paris"], False), ([" works", " works"], False), ([" Alice", " Paris"], True)]: + with patch.object(model.data_processor, "collate_fn", wraps=model.data_processor.collate_fn) as collate: + results = model.inference( + chunks, + labels, + session_id=["a", "b"], + batch_size=batch_size, + threshold=0.0, + return_class_probs=True, + recompute=recompute, + ) + encoded_prompts = [ + list(mapping) for call in collate.call_args_list for mapping in call.args[0]["classes_to_id"] + ] + assert encoded_prompts == prompts + for session_id, entities, row_prompts, row_names in zip(["a", "b"], results, prompts, names, strict=True): + state = model._session_cache.get(session_id) + assert state.labels == tuple(row_prompts) + assert state.label_names == tuple(row_names) + assert entities + assert all(entity["label"] in row_names for entity in entities) + assert all(set(entity["class_probs"]) == set(row_names) for entity in entities) + + +def test_single_session_accepts_per_text_description_mapping(model): + labels = [{"person": "A human"}] + for chunk in ["Alice", " works"]: + entities = model.inference([chunk], labels, session_id=["a"], threshold=0.0)[0] + assert entities + assert all(entity["label"] == "person" for entity in entities) + assert model._session_cache.get("a").text == "Alice works" + + +@pytest.mark.parametrize("updated", [{"person": "A place"}, {"location": "A human"}]) +def test_session_label_changes_require_recompute(model, updated): + model.inference(["Alice"], {"person": "A human"}, session_id=["a"]) + with pytest.raises(ValueError, match="Labels for session 'a' changed"): + model.inference([" works"], updated, session_id=["a"]) + assert model._session_cache.get("a").text == "Alice" + + results = model.inference([" works"], updated, session_id=["a"], recompute=True, threshold=0.0) + assert results[0] + assert all(entity["label"] in updated for entity in results[0]) + state = model._session_cache.get("a") + assert state.labels == tuple(updated.values()) + assert state.label_names == tuple(updated) + assert state.text == "Alice works" + + +def test_per_text_labels_stay_aligned_when_empty_chunks_are_skipped(model): + results = model.inference( + ["", "Paris"], + [{"person": "A human"}, {"location": "A place"}], + session_id=["empty", "city"], + threshold=0.0, + ) + assert results[0] == [] + assert results[1] + assert all(entity["label"] == "location" for entity in results[1]) + assert model._session_cache.get("empty") is None + assert model._session_cache.get("city").labels == ("A place",) + + +@pytest.mark.parametrize( + ("labels", "error", "message"), + [ + ([{"person": "A human"}], ValueError, "Per-text labels must have length 2"), + ({"person": 1}, TypeError, "string keys and values"), + ({"person": "same", "location": "same"}, ValueError, "descriptions must be unique"), + ([{"person": "A human"}, {}], ValueError, "At least one label is required"), + ], +) +def test_session_description_validation(model, labels, error, message): + with pytest.raises(error, match=message): + model.inference(["Alice", "Paris"], labels, session_id=["a", "b"]) + assert model.session_count == 0 diff --git a/tests/test_token_padding.py b/tests/test_token_padding.py new file mode 100644 index 00000000..270cb6cc --- /dev/null +++ b/tests/test_token_padding.py @@ -0,0 +1,74 @@ +from types import SimpleNamespace + +import torch +import pytest + +from gliner.model import UniEncoderTokenGLiNER +from gliner.decoding.decoder import TokenDecoder +from gliner.modeling.outputs import GLiNERBaseOutput +from gliner.data_processing.tokenizer import WordsSplitter + + +@pytest.mark.parametrize("threshold", [0.0, 0.2, 0.5, [0.2, 0.5, 0.2]]) +@pytest.mark.parametrize("flat_ner", [True, False]) +@pytest.mark.parametrize("padding_logit", [-1.0, 5.0]) +def test_token_decoder_ignores_padding(threshold, flat_ner, padding_logit): + decoder = TokenDecoder(SimpleNamespace()) + tokens = [["Alice"], ["works", "in", "Paris"], []] + mappings = [{1: "person"}, {1: "location", 2: "organization"}, {1: "person"}] + logits = torch.full((3, 3, 2, 3), -10.0) + logits[0, 0, 0] = -0.5 # A valid low-confidence entity, retained at threshold 0.2. + logits[0, 1:] = padding_logit + logits[1, 2, 0] = 5.0 + logits[2] = padding_logit + original_logits = logits.clone() + + results = decoder.decode(tokens, mappings, logits, threshold=threshold, flat_ner=flat_ner) + + for i, (row_tokens, spans) in enumerate(zip(tokens, results, strict=True)): + row_threshold = threshold[i] if isinstance(threshold, list) else threshold + unpadded = decoder.decode( + [row_tokens], + mappings[i], + logits[i : i + 1, : len(row_tokens)], + threshold=row_threshold, + flat_ner=flat_ner, + )[0] + assert spans == unpadded + assert all(0 <= span.start <= span.end < len(row_tokens) for span in spans) + assert results[2] == [] + torch.testing.assert_close(logits, original_logits) + + +def test_low_threshold_token_inference_maps_only_real_words(): + model = UniEncoderTokenGLiNER.__new__(UniEncoderTokenGLiNER) + torch.nn.Module.__init__(model) + model.config = SimpleNamespace() + model.data_processor = SimpleNamespace(words_splitter=WordsSplitter("whitespace")) + model._inference_packing_config = None + model.decoder = TokenDecoder(model.config) + + def collator(input_x, entity_types): + return { + "tokens": [item["tokenized_text"] for item in input_x], + "id_to_classes": dict(enumerate(entity_types, start=1)), + } + + def run_batch(batch, **kwargs): + rows = batch["tokens"] + logits = torch.full((len(rows), max(map(len, rows)), 1, 3), -10.0) + for i, row in enumerate(rows): + logits[i, len(row) - 1, 0] = -0.5 + logits[i, len(row) :] = -1.0 # Padding passes 0.2, but not 0.5. + return GLiNERBaseOutput(logits=logits) + + model.create_collator = lambda: collator + model.run_batch = run_batch + texts = ["Alice", "", "works in Paris"] + + assert model.inference(texts, ["entity"], threshold=0.5) == [[], [], []] + results = model.inference(texts, ["entity"], threshold=0.2) + assert results == model.inference(texts, ["entity"], threshold=0.2, batch_size=1) + assert [[entity["text"] for entity in entities] for entities in results] == [["Alice"], [], ["Paris"]] + assert (results[0][0]["start"], results[0][0]["end"]) == (0, 5) + assert (results[2][0]["start"], results[2][0]["end"]) == (9, 14) diff --git a/tests/test_trainer_checkpoints.py b/tests/test_trainer_checkpoints.py new file mode 100644 index 00000000..f9ee306c --- /dev/null +++ b/tests/test_trainer_checkpoints.py @@ -0,0 +1,128 @@ +from types import SimpleNamespace +from unittest.mock import Mock + +import torch +import pytest +import transformers +from torch import nn +from safetensors.torch import save_file + +pytest.importorskip("accelerate") + +from gliner.model import BaseGLiNER +from gliner.config import GLiNERConfig +from gliner.training import Trainer, TrainingArguments + + +class TinyGLiNER(BaseGLiNER): + _create_model = None + _create_data_processor = None + resize_embeddings = None + inference = None + evaluate = None + + def __init__(self): + nn.Module.__init__(self) + self.model = nn.ModuleDict({"encoder": nn.Linear(2, 2), "head": nn.Linear(2, 2)}) + self.model["head"].weight = self.model["encoder"].weight + self.model.register_buffer("count", torch.tensor(7)) + self.config = GLiNERConfig() + self.data_processor = SimpleNamespace(transformer_tokenizer=Mock()) + self._keys_to_ignore_on_save = None + + +def make_trainer(tmp_path, model=None): + args = TrainingArguments(output_dir=str(tmp_path), use_cpu=True, report_to="none") + return Trainer(model=TinyGLiNER() if model is None else model, args=args) + + +@pytest.mark.parametrize("safe_serialization", [False, True]) +@pytest.mark.parametrize("load_method", ["best", "resume", "explicit_model"]) +def test_checkpoint_restores_weights(tmp_path, caplog, safe_serialization, load_method): + trainer = make_trainer(tmp_path) + trainer.args.save_safetensors = safe_serialization + expected = {key: value.clone() for key, value in trainer.model.state_dict().items()} + checkpoint = tmp_path / "checkpoint-1" + trainer.save_model(str(checkpoint)) + assert (checkpoint / "gliner_config.json").is_file() + assert not (checkpoint / "config.json").exists() + + target = TinyGLiNER() if load_method == "explicit_model" else trainer.model + parameters = list(target.parameters()) + with torch.no_grad(): + for tensor in target.state_dict().values(): + tensor.zero_() + + if load_method == "best": + trainer.state.best_model_checkpoint = str(checkpoint) + trainer._load_best_model() + elif load_method == "explicit_model": + trainer._load_from_checkpoint(str(checkpoint), model=target) + else: + trainer._load_from_checkpoint(str(checkpoint)) + + for key, tensor in target.state_dict().items(): + torch.testing.assert_close(tensor, expected[key]) + assert all(before is after for before, after in zip(parameters, target.parameters(), strict=True)) + assert target.model["head"].weight is target.model["encoder"].weight + assert "missing keys" not in caplog.text + assert "unexpected keys" not in caplog.text + + +@pytest.mark.parametrize("load_method", ["best", "resume"]) +@pytest.mark.parametrize("gliner_model", [False, True]) +def test_standard_checkpoint_restores_wrapper_weights(tmp_path, caplog, load_method, gliner_model): + model = TinyGLiNER() if gliner_model else nn.Linear(2, 2) + trainer = make_trainer(tmp_path, model) + expected = {key: tensor.clone() for key, tensor in model.state_dict().items()} + save_file(expected, tmp_path / "model.safetensors") + with torch.no_grad(): + for tensor in model.state_dict().values(): + tensor.zero_() + + if load_method == "best": + trainer.state.best_model_checkpoint = str(tmp_path) + trainer._load_best_model() + else: + trainer._load_from_checkpoint(str(tmp_path)) + + for key, tensor in model.state_dict().items(): + torch.testing.assert_close(tensor, expected[key]) + assert "missing keys" not in caplog.text + assert "unexpected keys" not in caplog.text + + +def test_incomplete_checkpoint_still_reports_incompatible_keys(tmp_path, caplog): + trainer = make_trainer(tmp_path) + trainer.save_model(str(tmp_path)) + state_dict = {key: tensor.clone() for key, tensor in trainer.model.model.state_dict().items()} + del state_dict["encoder.bias"] + state_dict["unexpected"] = torch.zeros(1) + save_file(state_dict, tmp_path / "model.safetensors") + + trainer._load_from_checkpoint(str(tmp_path)) + + assert "missing keys" in caplog.text + assert "encoder.bias" in caplog.text + assert "unexpected keys" in caplog.text + + +@pytest.mark.parametrize("load_method", ["_load_best_model", "_load_from_checkpoint"]) +@pytest.mark.parametrize("backend", ["is_deepspeed_enabled", "is_fsdp_enabled", "sagemaker"]) +def test_distributed_checkpoint_loading_uses_transformers(tmp_path, monkeypatch, load_method, backend): + trainer = make_trainer(tmp_path) + trainer.save_model(str(tmp_path)) + if backend == "sagemaker": + monkeypatch.setattr("transformers.trainer.is_sagemaker_mp_enabled", lambda: True) + else: + setattr(trainer, backend, True) + parent_load = Mock() + monkeypatch.setattr(transformers.Trainer, load_method, parent_load) + + if load_method == "_load_best_model": + trainer.state.best_model_checkpoint = str(tmp_path) + trainer._load_best_model() + parent_load.assert_called_once_with() + else: + trainer._load_from_checkpoint(str(tmp_path), model=trainer.model) + parent_load.assert_called_once_with(str(tmp_path), model=trainer.model) diff --git a/tests/test_training_args.py b/tests/test_training_args.py new file mode 100644 index 00000000..080635fd --- /dev/null +++ b/tests/test_training_args.py @@ -0,0 +1,48 @@ +import pytest +import transformers + +pytest.importorskip("accelerate") + +from gliner.model import BaseGLiNER +from gliner.training import TrainingArguments + + +@pytest.mark.parametrize( + ("kwargs", "expected"), + [ + ({}, 11), + ({"warmup_steps": 7}, 7), + ({"warmup_ratio": 0.25}, 26), + ({"warmup_ratio": 0.25, "warmup_steps": 7}, 7), + ({"warmup_ratio": 0.0}, 0), + ({"warmup_ratio": 1.0}, 101), + ({"max_steps": -1, "num_train_epochs": 2, "warmup_ratio": 0.25}, 26), + ], +) +def test_create_training_args_warmup(tmp_path, kwargs, expected): + args = BaseGLiNER.create_training_args(output_dir=str(tmp_path), use_cpu=True, **kwargs) + + assert args.get_warmup_steps(101) == expected + + +def test_direct_training_arguments_warmup(tmp_path): + args = TrainingArguments(output_dir=str(tmp_path), use_cpu=True, report_to="none", warmup_ratio=0.25) + + assert args.get_warmup_steps(101) == 26 + assert args.to_dict()["warmup_ratio"] == 0.25 + + +@pytest.mark.parametrize("ratio", [-0.1, 1.1, float("nan")]) +def test_invalid_warmup_ratio(tmp_path, ratio): + with pytest.raises(ValueError, match="warmup_ratio"): + BaseGLiNER.create_training_args(output_dir=str(tmp_path), use_cpu=True, warmup_ratio=ratio) + + +@pytest.mark.skipif( + "warmup_ratio" in transformers.TrainingArguments.__dataclass_fields__, + reason="Fractional warmup_steps requires the new Transformers warmup API", +) +def test_fractional_warmup_steps_takes_precedence(tmp_path): + args = BaseGLiNER.create_training_args(output_dir=str(tmp_path), use_cpu=True, warmup_steps=0.2) + + assert args.get_warmup_steps(101) == 21 From 688865c5b98a12f996a94c4cc8d131213d11d7f1 Mon Sep 17 00:00:00 2001 From: Ingvar Date: Tue, 8 Sep 2026 11:53:11 +0300 Subject: [PATCH 2/2] update docs --- docs/architectures.md | 4 +- docs/configs.md | 2 +- docs/index.md | 2 - docs/input_limits.md | 217 --------------- docs/offline.md | 35 --- docs/serving.md | 2 +- docs/streaming.md | 374 ------------------------- docs/usage.md | 624 +++++++++++++++++++++++++++++++++++++++++- 8 files changed, 614 insertions(+), 646 deletions(-) delete mode 100644 docs/input_limits.md delete mode 100644 docs/offline.md delete mode 100644 docs/streaming.md diff --git a/docs/architectures.md b/docs/architectures.md index 7be1d74a..fee9b962 100644 --- a/docs/architectures.md +++ b/docs/architectures.md @@ -32,7 +32,7 @@ All stateless architectures apply the checkpoint's text-only `config.max_len` before encoding labels. Whether labels share a later backbone context with the text depends on the architecture: uni-encoders share it, while bi-encoders keep the two sequences separate. Cached StreamingSpan sessions use their causal -decoder context instead. See [Input limits and truncation](input_limits.md) for +decoder context instead. See [Input limits and truncation](usage.md#input-limits-and-truncation) for the complete behavior and production preflight guidance. ## Vanilla GLiNER (UniEncoderSpan) @@ -660,7 +660,7 @@ complete current entity snapshot with document-relative character offsets. It is not a delta: callers that render live changes should compare consecutive snapshots to identify additions, updates, and removals. -The [StreamingSpan guide](streaming.md) covers the session API, persistent +The [StreamingSpan guide](usage.md#streamingspan-models) covers the session API, persistent batches, asynchronous microbatching, cache cleanup, configuration, and deployment trade-offs. A PII-oriented checkpoint is available at [knowledgator/gliner-stream-pii-v1.0](https://huggingface.co/knowledgator/gliner-stream-pii-v1.0). diff --git a/docs/configs.md b/docs/configs.md index d128624d..c141fe0b 100644 --- a/docs/configs.md +++ b/docs/configs.md @@ -163,7 +163,7 @@ configurations may override the base default. This is not necessarily the backbone's total context length. Uni-encoder prompts and text share the later transformer subword/context limit, while bi-encoders encode labels separately. See [Input limits and -truncation](input_limits.md) for all limits and the StreamingSpan exception. +truncation](usage.md#input-limits-and-truncation) for all limits and the StreamingSpan exception. --- diff --git a/docs/index.md b/docs/index.md index ffe6192c..b4ee1d93 100644 --- a/docs/index.md +++ b/docs/index.md @@ -10,9 +10,7 @@ This documentation includes installation guides, tutorials, advanced topics, and intro instalation quickstart -input_limits usage -streaming configs training architectures diff --git a/docs/input_limits.md b/docs/input_limits.md deleted file mode 100644 index 2ff23fa7..00000000 --- a/docs/input_limits.md +++ /dev/null @@ -1,217 +0,0 @@ -# Input limits and truncation - -GLiNER does not automatically split a long document into windows. For ordinary -stateless inference, each input that exceeds the checkpoint's `config.max_len` -is reduced to a prefix. In GLiNER 0.2.27 this produces a `UserWarning`, but the -prediction call still succeeds and its return value does not say that part of -the input was skipped. - -This behavior applies to span, token, bi-encoder, decoder, and relation -extraction models. Cached StreamingSpan sessions are the exception described -under **Architecture-specific limits** below. - -## What `config.max_len` counts - -`config.max_len` is the maximum number of **text tokens produced by GLiNER's -word splitter**. It is not a character count, a whitespace-word count, a -transformer subword count, or the total prompt-plus-text sequence length. - -For raw-text inference, the processing order is: - -```text -raw text - -> WordsSplitter(config.words_splitter_type) - -> N text tokens - -> if N > config.max_len: warn and keep tokens[:config.max_len] - -> add the entity/relation-type prompt - -> transformer subword tokenization - -> model -``` - -The default `whitespace` splitter also separates punctuation. For example, -`"Acme, Inc."` becomes four splitter tokens (`Acme`, `,`, `Inc`, `.`), even -though a simple `text.split()` returns two strings. Other -`words_splitter_type` values segment text differently. Always inspect the -loaded checkpoint rather than assuming the configuration default: - -```python -print(model.config.max_len) -print(model.config.words_splitter_type) -``` - -The base configuration defaults to `max_len=384`, but checkpoints can save a -different value. For example, -`EmergentMethods/gliner_medium_news-v2.1` saves `max_len=296`. - -## Do labels reduce the text budget? - -At the `config.max_len` stage, **no**. GLiNER splits and truncates the text -first, then prepends or separately encodes the entity-type prompt. Ten labels -and one hundred labels therefore receive the same `config.max_len` allowance -for text. - -There is a second, independent limit to consider. A transformer sees -subtokens, and uni-encoder architectures put the label prompt and retained text -in one transformer sequence. Prompt subtokens therefore consume part of any -finite tokenizer or backbone context capacity. A finite tokenizer limit can -truncate the text tail; a backbone limit can instead reject an oversized -combined sequence. Either way, adding labels reduces the remaining combined -headroom even though it does not change `config.max_len`. - -GLiNER calls the transformer tokenizer with `truncation=True` but without an -explicit tokenizer `max_length`, so the tokenizer's registered limit and -truncation side determine tokenizer-level behavior. Some tokenizers do not -register a finite limit; Transformers may then warn that no maximum was -provided and perform no tokenizer-level truncation. That does not prove that -the backbone supports an unlimited sequence. - -## Behavior when text exceeds `max_len` - -For each input independently, when `num_tokens > model.config.max_len`, GLiNER -0.2.27: - -1. emits a `UserWarning` such as - `Sentence of length 987 has been truncated to 296`; -2. retains only the first `max_len` splitter tokens; -3. runs inference normally on that prefix; and -4. returns ordinary-looking predictions with no truncation metadata. - -Text after the retained prefix is never presented to the model. Entities there -cannot be returned, and entities crossing the boundary are incomplete. The -same prefix truncation occurs during training when a pre-tokenized example is -too long. - -Python warning filters control whether the warning is visible. With the default -filter, repeated warnings from the same location may be shown only once per -process. Warnings may also bypass an application's structured logs. Do not use -the presence or absence of the warning as a per-request completeness signal. - -As of 0.2.27, `predict_entities` and `inference` do not provide a -`return_truncation_info` result and do not have a `truncation="error"` mode. -[Issue #231](https://github.com/urchade/GLiNER/issues/231) tracks the request for -fail-on-truncation behavior. - -## Production preflight without processor internals - -Use the public `prepare_batch` stage to count exactly the splitter tokens that -stateless inference will receive. This avoids depending on -`model.data_processor.words_splitter`: - -```python -def truncation_info(model, text, labels): - prepared = model.prepare_batch(text, labels) - num_tokens = len(prepared["tokens"][0]) if prepared["tokens"] else 0 - max_len = model.config.max_len - return { - "truncated": num_tokens > max_len, - "num_tokens": num_tokens, - "max_len": max_len, - } - - -info = truncation_info(model, text, labels) -if info["truncated"]: - raise ValueError( - f"GLiNER input has {info['num_tokens']} tokens; " - f"the model limit is {info['max_len']}" - ) - -entities = model.predict_entities(text, labels) -``` - -For a batch, `prepared["tokens"]` contains one token list per non-empty input; -`prepared["valid_to_orig_idx"]` maps those lists back to the original batch -indices. Emit the resulting information in the service response or reject the -request before inference, according to the service contract. - -This preflight covers `config.max_len`; it does not measure a combined -uni-encoder subword sequence, an inference-packing limit, or a cached streaming -session's remaining context. - -## Processing long documents - -For full-document coverage, split the text into overlapping windows no longer -than `config.max_len`, predict each window, shift its character offsets back to -document coordinates, and reconcile duplicate predictions from overlaps. The -public preparation result includes the exact token-to-character maps needed to -make splitter-aligned windows: - -```python -def iter_gliner_windows(model, text, labels, overlap): - prepared = model.prepare_batch(text, labels) - if not prepared["tokens"]: - return - - tokens = prepared["tokens"][0] - starts = prepared["start_token_map"][0] - ends = prepared["end_token_map"][0] - window_size = model.config.max_len - - if not 0 <= overlap < window_size: - raise ValueError("overlap must be in [0, model.config.max_len)") - - step = window_size - overlap - for first in range(0, len(tokens), step): - last = min(first + window_size, len(tokens)) - char_start = starts[first] - char_end = ends[last - 1] - yield char_start, text[char_start:char_end] - if last == len(tokens): - break -``` - -For span models, an overlap of at least `max_width - 1` tokens ensures that an -entity no wider than `max_width` is fully contained in some window. A larger -overlap may be useful for contextual accuracy, and token-level models may need -a task-specific overlap because their entity length is not bounded by -`max_width`. Relation extraction also needs application-specific merging; -relations whose endpoints never occur together in a window cannot be inferred. - -GLiNER does not merge window outputs for you. When shifting an entity from a -window beginning at `char_start`, add `char_start` to both `entity["start"]` and -`entity["end"]`. A common overlap policy is to group predictions by -`(start, end, label)` and retain the highest score. - -## Changing `max_len` - -The loader's `max_length` argument overrides the saved `config.max_len`: - -```python -model = GLiNER.from_pretrained( - "urchade/gliner_small-v2.1", - max_length=512, -) -``` - -The name difference is intentional: `max_length=` at load time writes -`model.config.max_len`. Passing `max_length` or `truncation` to -`predict_entities` is not a supported way to change the word-level limit. - -Increasing this value does not resize the backbone context, change the -transformer tokenizer's `model_max_length`, or guarantee quality beyond the -lengths used to train the checkpoint. Subword expansion and the label prompt -can make the combined sequence longer than the word-token count suggests, and -longer inputs require more memory. Prefer windowing unless the checkpoint and -backbone are known to support the larger value. - -## Architecture-specific limits - -| Architecture or path | How the label prompt is encoded | Effective-limit notes | -|---|---|---| -| UniEncoderSpan and UniEncoderToken | Prompt and text share one backbone sequence | `config.max_len` is text-only, but a finite backbone/tokenizer subword limit includes the prompt | -| UniEncoder span/token decoders | Main encoder prompt and text share a sequence; generated labels use an auxiliary decoder | The main encoder has the same two-stage limits as other uni-encoders; the decoder has its own generation limits | -| UniEncoder relation extraction | Entity labels, relation labels, and text share one sequence | Both prompt types can consume the combined backbone context, but neither changes the first-stage text-only `config.max_len` check | -| BiEncoderSpan and BiEncoderToken | Text and labels are encoded separately | Label count does not consume the text encoder sequence; each encoder still has its own tokenizer/backbone limit | -| StreamingSpan without `session_id` | Prompt and text share one causal sequence | Uses the normal stateless `config.max_len` prefix truncation, followed by the causal backbone limit | -| Cached StreamingSpan session | The initial prompt and all appended text share the decoder context | Stateless `config.max_len` truncation is bypassed; exceeding the smaller of `max_cache_length` and the decoder's native limit raises `ValueError` instead of dropping old text | - -Inference packing adds another independent setting: -`InferencePackingConfig.max_length` is measured in already-tokenized backbone -token IDs, not splitter tokens. If one encoded request is longer than that -value, the current packer keeps its first `max_length` token IDs without the -`config.max_len` warning. Set the packing limit for the complete encoded -request (including a uni-encoder prompt), or disable packing for requests that -may exceed it. - -Finally, `max_width` is not an input-length limit. It controls the widest -candidate entity span, in splitter tokens, for span-based architectures. diff --git a/docs/offline.md b/docs/offline.md deleted file mode 100644 index 2fa6be5d..00000000 --- a/docs/offline.md +++ /dev/null @@ -1,35 +0,0 @@ -# Loading models offline - -Prepare a complete model directory on a machine with internet access: - -```python -from gliner import GLiNER - -model = GLiNER.from_pretrained("urchade/gliner_multi-v2.1") -model.save_pretrained("gliner-offline", safe_serialization=True) -``` - -Copy the entire `gliner-offline` directory to the offline machine, then load it: - -```python -from gliner import GLiNER - -model = GLiNER.from_pretrained("gliner-offline", local_files_only=True) -``` - -`save_pretrained` includes model weights, the resolved backbone configuration, -and tokenizers. Models with a separate label encoder or generative decoder also -save their auxiliary tokenizer in `labels_tokenizer/` or `decoder_tokenizer/`. -Keep these subdirectories with the model when copying it. - -Older Hub checkpoints may contain a GLiNER configuration that only names the -backbone, such as `microsoft/mdeberta-v3-base`, without embedding its configuration. -Downloading that checkpoint's files alone may therefore be insufficient. Loading -and saving it with the code above resolves and packages those dependencies. - -`local_files_only=True` restricts loading to local files and the Hugging Face -cache; it does not download missing dependencies. An incomplete legacy checkpoint -still requires its missing backbone configuration or tokenizer to be cached or -provided locally. If the backbone files are in a separate local directory, -`model_name` in `gliner_config.json` can point to that directory. Missing files -raise an error without attempting a network connection. diff --git a/docs/serving.md b/docs/serving.md index 16098705..ecc60a04 100644 --- a/docs/serving.md +++ b/docs/serving.md @@ -94,7 +94,7 @@ overrides the loaded model's text-only `config.max_len`; requests above it can still return HTTP success after GLiNER keeps only the prefix and emits a Python warning inside the replica. If the service contract requires full-input coverage, validate or window requests explicitly. See [Input limits and -truncation](input_limits.md) for a preflight that uses the public model API. +truncation](usage.md#input-limits-and-truncation) for a preflight that uses the public model API. ::: **HTTP request (no client library):** diff --git a/docs/streaming.md b/docs/streaming.md deleted file mode 100644 index 5eb4184d..00000000 --- a/docs/streaming.md +++ /dev/null @@ -1,374 +0,0 @@ -# StreamingSpan - -StreamingSpan is GLiNER's architecture for named entity recognition over text -that arrives incrementally. It uses a causal decoder as the text backbone, -retains reusable state between chunks, and revises recent span predictions when -new right context becomes available. - -Use the -[knowledgator/gliner-stream-pii-v1.0](https://huggingface.co/knowledgator/gliner-stream-pii-v1.0) -checkpoint to get started with streaming PII detection. `GLiNER.from_pretrained` -reads `model_type="gliner_streaming_span"` from the checkpoint and selects -`StreamingSpanGLiNER` automatically. - -```python -from gliner import GLiNER - -model = GLiNER.from_pretrained("knowledgator/gliner-stream-pii-v1.0") -model.eval() -``` - -## Quick start - -Pass one chunk at a time with a stable session ID. The chunks are concatenated -exactly as supplied, so retain spaces and punctuation at chunk boundaries. - -```python -labels = ["person", "email address", "phone number"] -session_id = "support-call-42" -chunks = [ - "Customer Alice Johnson ", - "can be reached at alice@example.com ", - "or +1 202-555-0147.", -] - -try: - for chunk in chunks: - snapshot = model.inference( - [chunk], - labels, - session_id=[session_id], - threshold=0.5, - )[0] - print(snapshot) -finally: - model.clear_session(session_id) -``` - -For `model.inference(..., session_id=[...])`, labels can also be a single -`{"label name": "description"}` dictionary shared by all input chunks, or a list -of dictionaries with one complete label set per input chunk/session. See -[Label descriptions](usage.md#label-descriptions) for the input formats. Keep each -session's label names, descriptions, and order consistent across calls; changing -them requires `recompute=True` or clearing that session. - -Each `snapshot` is the complete set of entities currently active for the -accumulated session text. It is not a list of only the entities detected in the -latest chunk. An entity has the same shape as an ordinary GLiNER prediction: - -```python -{ - "start": 9, # document-relative, inclusive character offset - "end": 22, # document-relative, exclusive character offset - "text": "Alice Johnson", - "label": "person", - "score": 0.93, -} -``` - -Scores and even the active boundaries may change between snapshots as more -context arrives. Consumers that need an event stream should diff consecutive -snapshots by `(start, end, label)`. - -## Architecture - -![GLiNER StreamingSpan architecture](images/gliner-streaming-architecture.svg) - -StreamingSpan uses the following cold and warm paths: - -1. On the first append, labels and text are serialized as - `label<