From 3dd3567f6830599859613754fb8b92cd47041e1a Mon Sep 17 00:00:00 2001 From: Yurii Havrylko <10372700+yuriihavrylko@users.noreply.github.com> Date: Sat, 4 Jul 2026 23:48:46 +0200 Subject: [PATCH] Fix token_lengths kwarg leaking into the labels encoder (#370) --- gliner/modeling/encoder.py | 9 +++++++-- 1 file changed, 7 insertions(+), 2 deletions(-) diff --git a/gliner/modeling/encoder.py b/gliner/modeling/encoder.py index 2852e3c..1717053 100644 --- a/gliner/modeling/encoder.py +++ b/gliner/modeling/encoder.py @@ -943,8 +943,8 @@ def encode_labels( input_ids: Label token IDs of shape (batch_size, seq_len). attention_mask: Attention mask of shape (batch_size, seq_len). *args: Additional positional arguments. - **kwargs: Additional keyword arguments (packing_config and pair_attention_mask - are removed as they're not supported for labels). + **kwargs: Additional keyword arguments (packing_config, pair_attention_mask + and token_lengths are removed as they're not supported for labels). Returns: Pooled label embeddings of shape (batch_size, hidden_size). @@ -952,6 +952,11 @@ def encode_labels( label_kwargs = dict(kwargs) label_kwargs.pop("packing_config", None) label_kwargs.pop("pair_attention_mask", None) + # token_lengths carries precomputed lengths for the *text* inputs (see + # collator); forwarding it to the labels encoder crashes plain HF + # backbones ("BertModel.forward() got an unexpected keyword argument + # 'token_lengths'", #370). + label_kwargs.pop("token_lengths", None) label_kwargs["attention_mask"] = attention_mask labels_embeddings = self.labels_encoder(input_ids, *args, **label_kwargs) if hasattr(self, "labels_projection"):