Fix per-row label sets in inference(): KeyError on unequal lengths, silent mislabelling across chunks

gliner/decoding/decoder.py

@@ -286,6 +286,17 @@ def _decode_batch_item(
                     span_filter[word_start, width] = True
             probs_i = probs_i * span_filter.unsqueeze(-1)
 
+        # Same padding guard as _decode_batch: with per-row label sets this row's scores may span
+        # more class slots than its own label set has.
+        num_classes = probs_i.shape[-1]
+        if len(id_to_class_i) < num_classes:
+            valid_classes = torch.tensor(
+                [class_idx + 1 in id_to_class_i for class_idx in range(num_classes)],
+                dtype=torch.bool,
+                device=probs_i.device,
+            )
+            probs_i = probs_i * valid_classes
+
         span_i = []
 
         # Find all spans above threshold
@@ -408,6 +419,25 @@ class IDs to class names.
                     span_filter[i] = True
             probs = probs * span_filter.unsqueeze(-1)
 
+        # Pre-resolve id_to_class mappings per batch item
+        id_to_class_per_item = [self._get_id_to_class_for_sample(id_to_classes, i) for i in range(B)]
+
+        # With per-row label sets the class dimension is padded to the batch-wide maximum, so a row
+        # with fewer labels carries scores in class slots it never asked for. Mask them out — the
+        # same guard _decode_explicit_spans already applies — or torch.where returns class indices
+        # that are absent from that row's id_to_class and _build_span_tuple raises KeyError.
+        num_classes = probs.shape[-1]
+        if any(len(m) < num_classes for m in id_to_class_per_item):
+            valid_classes = torch.tensor(
+                [
+                    [class_idx + 1 in id_to_class for class_idx in range(num_classes)]
+                    for id_to_class in id_to_class_per_item
+                ],
+                dtype=torch.bool,
+                device=probs.device,
+            )
+            probs = probs * valid_classes[:, None, None, :]
+
         # ONE torch.where on the full (B, L, K, C) tensor
         threshold_tensor = _threshold_compare_tensor(threshold, B, probs.device, probs.dim())
         b_idx, s_idx, k_idx, c_idx = torch.where(probs > threshold_tensor)
@@ -448,9 +478,6 @@ class IDs to class names.
             top_probs_list = all_top_probs.tolist()
             top_indices_list = all_top_indices.tolist()
 
-        # Pre-resolve id_to_class mappings per batch item
-        id_to_class_per_item = [self._get_id_to_class_for_sample(id_to_classes, i) for i in range(B)]
-
         # Group by batch item and build Span objects (pure Python)
         batch_spans: List[List[Span]] = [[] for _ in range(B)]
         for j, (b, s, k, c, flat_idx, score) in enumerate(zip(b_list, s_list, k_list, c_list, flat_idxs, scores)):

gliner/model.py

@@ -115,6 +115,19 @@
 )
 from .data_processing.tokenizer import WordsSplitter
 
+
+def _entity_types_for_chunk(entity_types, indices):
+    """Select the label sets belonging to the rows in one DataLoader chunk.
+
+    ``entity_types`` is either a single flat list shared by every row, or -- when ``inference`` is
+    called with ``labels=List[List[str]]`` -- one list per row. Only the per-row form may be
+    sliced; slicing the shared form would silently drop labels.
+    """
+    if entity_types and isinstance(entity_types[0], list):
+        return [entity_types[i] for i in indices]
+    return entity_types
+
+
 if is_module_available("onnxruntime"):
     import onnxruntime as ort
 
@@ -2510,11 +2523,15 @@ def inference(
 
         collator = self.create_collator()
 
-        def collate_fn(batch):
-            return self.collate_batch(batch, prepared["entity_types"], collator)
+        def collate_fn(indices):
+            return self.collate_batch(
+                [prepared["input_x"][i] for i in indices],
+                _entity_types_for_chunk(prepared["entity_types"], indices),
+                collator,
+            )
 
         data_loader = torch.utils.data.DataLoader(
-            prepared["input_x"],
+            list(range(len(prepared["input_x"]))),
             batch_size=batch_size,
             shuffle=False,
             collate_fn=collate_fn,
@@ -4937,11 +4954,15 @@ def inference(
 
         collator = self.create_collator()
 
-        def collate_fn(batch):
-            return self.collate_batch(batch, prepared["entity_types"], collator)
+        def collate_fn(indices):
+            return self.collate_batch(
+                [prepared["input_x"][i] for i in indices],
+                _entity_types_for_chunk(prepared["entity_types"], indices),
+                collator,
+            )
 
         data_loader = torch.utils.data.DataLoader(
-            prepared["input_x"],
+            list(range(len(prepared["input_x"]))),
             batch_size=batch_size,
             shuffle=False,
             collate_fn=collate_fn,
@@ -5485,11 +5506,16 @@ def inference(
 
         collator = self.create_collator()
 
-        def collate_fn(batch):
-            return self.collate_batch(batch, prepared["entity_types"], collator, prepared["relation_types"])
+        def collate_fn(indices):
+            return self.collate_batch(
+                [prepared["input_x"][i] for i in indices],
+                _entity_types_for_chunk(prepared["entity_types"], indices),
+                collator,
+                prepared["relation_types"],
+            )
 
         data_loader = torch.utils.data.DataLoader(
-            prepared["input_x"],
+            list(range(len(prepared["input_x"]))),
             batch_size=batch_size,
             shuffle=False,
             collate_fn=collate_fn,

tests/test_decoder.py

@@ -207,6 +207,45 @@ def test_per_sample_id_to_classes(self, basic_config, basic_inputs):
             assert batch_0_types.issubset({'PERSON', 'ORG'})
             assert batch_1_types.issubset({'LOCATION', 'DATE'})
 
+    def test_ragged_per_sample_id_to_classes(self, basic_config, basic_inputs):
+        """Per-sample mappings of DIFFERENT sizes must decode without raising.
+
+        The class dimension is padded to the batch-wide maximum, so a sample with fewer labels
+        still carries scores in class slots it never asked for. Those slots have no entry in that
+        sample's id_to_class and previously raised KeyError.
+        """
+        decoder = SpanDecoder(basic_config)
+
+        # sample 0 has two classes, sample 1 only one -- logits still have C=2
+        id_to_classes_list = [{1: 'PERSON', 2: 'ORG'}, {1: 'LOCATION'}]
+
+        result = decoder.decode(
+            tokens=basic_inputs['tokens'],
+            id_to_classes=id_to_classes_list,
+            model_output=basic_inputs['logits'],
+            threshold=0.5,
+        )
+
+        assert {span.entity_type for span in result[0]}.issubset({'PERSON', 'ORG'})
+        assert {span.entity_type for span in result[1]}.issubset({'LOCATION'})
+
+    def test_ragged_per_sample_id_to_classes_single_item(self, basic_config):
+        """Same guard on the B == 1 path, which decodes per item rather than per batch."""
+        decoder = SpanDecoder(basic_config)
+
+        logits = torch.full((1, 3, 2, 3), -10.0)
+        logits[0, 0, 0, 0] = 5.0   # class 1 -- present in the mapping
+        logits[0, 1, 0, 2] = 5.0   # class 3 -- a padding slot, must be ignored
+
+        result = decoder.decode(
+            tokens=[['Alice', 'met', 'Bob']],
+            id_to_classes=[{1: 'PERSON'}],
+            model_output=logits,
+            threshold=0.5,
+        )
+
+        assert {span.entity_type for span in result[0]} == {'PERSON'}
+
     def test_empty_predictions(self, basic_config):
         """Should handle case with no predictions above threshold."""
         decoder = SpanDecoder(basic_config)

tests/test_models.py

@@ -1,7 +1,7 @@
 import pytest
 
 from gliner import GLiNER
-from gliner.model import BaseEncoderGLiNER, UniEncoderSpanRelexGLiNER
+from gliner.model import BaseEncoderGLiNER, UniEncoderSpanRelexGLiNER, _entity_types_for_chunk
 
 
 class _WordsSplitter:
@@ -77,3 +77,22 @@ def test_relex_prepare_batch_validates_per_text_relation_count():
             [["person"], ["organization"]],
             relations=[["works_at"]],
         )
+
+
+def test_entity_types_for_chunk_slices_per_row_label_sets():
+    """Per-row label sets must follow their rows into each DataLoader chunk.
+
+    The loader hands collate_fn one chunk at a time; without this the decoder indexes the
+    whole-batch list with a chunk-local index and rows past the first chunk are decoded against
+    another row's labels.
+    """
+    per_row = [["a"], ["b"], ["c"], ["d"]]
+    assert _entity_types_for_chunk(per_row, [2, 3]) == [["c"], ["d"]]
+    assert _entity_types_for_chunk(per_row, [0, 1]) == [["a"], ["b"]]
+
+
+def test_entity_types_for_chunk_passes_shared_label_list_through():
+    """A single flat list is shared by every row and must NOT be sliced."""
+    shared = ["person", "organization", "location"]
+    assert _entity_types_for_chunk(shared, [1, 2]) == shared
+    assert _entity_types_for_chunk([], [0]) == []