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

[jedediahg]

Fixes #383.

inference() accepts labels: List[List[str]] — one label set per text — and the decoder has _get_id_to_class_for_sample for exactly that. Two defects make the path unusable: unequal-length sets raise KeyError, and equal-length sets are silently mislabelled as soon as the batch spans more than one chunk.

Callers passing a single flat label list (the documented path) are unaffected.

1. Padding class slots were never masked

With per-row label sets the class dimension is padded to the batch-wide maximum, so a row with fewer labels still carries scores in slots it never asked for. torch.where returned those class indices and _build_span_tuple raised KeyError, since they have no entry in that row's id_to_class.

Fixed by masking before the candidate search in BaseSpanDecoder._decode_batch and _decode_batch_item, reusing the guard _decode_explicit_spans already applies:

[class_idx + 1 in id_to_class for class_idx in range(num_classes)]

Masking rather than defaulting the lookup is deliberate. A .get(class_idx + 1, ...) at the call site — which is what the neighbouring return_class_probs path does — would stop the exception and emit spans labelled from padding slots: real scores, meaningless labels, silently. That trades a loud failure for a quiet one.

The mask is guarded by if any(len(m) < num_classes ...), so the uniform-label case pays only the length check.

2. id_to_classes was never sliced to the chunk

inference() built the mapping once for the whole batch and closed over it:

def collate_fn(batch):
    return self.collate_batch(batch, prepared["entity_types"], collator)   # whole-batch list

The DataLoader passes collate_fn only the current chunk's rows, but the closure handed the collator the full entity_types every time; collator.py:110 forwarded it unsliced, and the decoder then indexed that global list with a chunk-local index. The first chunk was correct by accident; every later chunk wore another row's labels, with no error.

The loaders now iterate row indices and select each chunk's label sets through _entity_types_for_chunk, mirroring how _process_batches already slices the other per-row input, word_input_spans. A shared flat list is passed through unsliced. Applied at all three loader sites.

Verification

Before:

batch_size=6 [['politician'],['politician'],['politician'],['judge'],['judge'],['judge']]      correct
batch_size=4 [['politician'],['politician'],['politician'],['judge'],['politician'],['politician']]
batch_size=2 [['politician'],['politician'],['politician'],['politician'],['politician'],['politician']]

After: correct at every batch size, and unequal-length sets no longer raise.

• The flat-label path is byte-identical before and after, verified on urchade/gliner_multi-v2.1 at batch sizes 1, 2, 4 and 8.
• Four regression tests added, each of which fails without this change: two in tests/test_decoder.py covering the batch (B > 1) and single-item (B == 1) paths — the existing test_per_sample_id_to_classes only uses equal-length mappings, which is the gap — and two in tests/test_models.py for the chunk helper.
• tests/test_decoder.py and tests/test_models.py pass (58 passed).
• ruff check reports exactly the same set of pre-existing findings before and after. I deliberately did not run ruff format on these files: the repo is not currently format-clean, and reformatting would bury a ~120-line fix in unrelated churn.

Once both are fixed, a fused call carrying a 9-label row set and a 1-label row set returns output equal to the two separate calls — max score delta 9.17e-05 over 3,827 spans on 100 real documents, with no threshold crossings.