diff --git a/gliner/data_processing/processor.py b/gliner/data_processing/processor.py index 0e570d9b..08836843 100644 --- a/gliner/data_processing/processor.py +++ b/gliner/data_processing/processor.py @@ -1071,6 +1071,50 @@ def tokenize_inputs(self, texts, entities=None): tokenized_inputs["words_mask"] = torch.tensor(words_masks) return tokenized_inputs + @staticmethod + def _prepare_entity_batch(classes_to_id): + """Prepare one shared label-encoder input without losing per-row order. + + A bi-encoder can encode the distinct labels in a batch once. For + per-example mappings, the returned gather indices restore those shared + embeddings to each example's local (1-indexed) class layout. + """ + if isinstance(classes_to_id, dict): + entities = [label for label, _ in sorted(classes_to_id.items(), key=lambda item: item[1])] + return entities, None, None + + entity_rows = [ + [label for label, _ in sorted(mapping.items(), key=lambda item: item[1])] + for mapping in classes_to_id + ] + if not entity_rows: + return [], None, None + if not any(entity_rows): + # Tokenizers generally reject an empty batch. Encode one internal + # placeholder, then gather it down to a zero-width class layout. + empty_indices = torch.empty(len(entity_rows), 0, dtype=torch.long) + empty_mask = torch.empty(len(entity_rows), 0, dtype=torch.bool) + return [""], empty_indices, empty_mask + if all(row == entity_rows[0] for row in entity_rows[1:]): + return entity_rows[0], None, None + + entities = list(dict.fromkeys(label for row in entity_rows for label in row)) + + max_classes = max((max(mapping.values(), default=0) for mapping in classes_to_id), default=0) + labels_gather_indices = torch.zeros(len(classes_to_id), max_classes, dtype=torch.long) + prompts_embedding_mask = torch.zeros(len(classes_to_id), max_classes, dtype=torch.bool) + + entity_to_index = {label: index for index, label in enumerate(entities)} + for batch_index, mapping in enumerate(classes_to_id): + for label, class_id in mapping.items(): + class_index = class_id - 1 + if class_index < 0: + raise ValueError("Bi-encoder class IDs must be positive") + labels_gather_indices[batch_index, class_index] = entity_to_index[label] + prompts_embedding_mask[batch_index, class_index] = True + + return entities, labels_gather_indices, prompts_embedding_mask + def batch_generate_class_mappings( self, batch_list: List[Dict], *args ) -> Tuple[List[Dict[str, int]], List[Dict[int, str]]]: @@ -1133,14 +1177,18 @@ def tokenize_and_prepare_labels(self, batch, prepare_labels, prepare_entities=Tr Returns: Dictionary containing tokenized inputs, entity encodings, and optionally labels. """ + labels_gather_indices = None + prompts_embedding_mask = None if prepare_entities: - if isinstance(batch["classes_to_id"], dict): - entities = list(batch["classes_to_id"]) - else: - entities = list(batch["classes_to_id"][0]) + entities, labels_gather_indices, prompts_embedding_mask = self._prepare_entity_batch( + batch["classes_to_id"] + ) else: entities = None tokenized_input = self.tokenize_inputs(batch["tokens"], entities) + if labels_gather_indices is not None: + tokenized_input["labels_gather_indices"] = labels_gather_indices + tokenized_input["prompts_embedding_mask"] = prompts_embedding_mask if prepare_labels: labels = self.create_labels(batch) tokenized_input["labels"] = labels @@ -1166,15 +1214,19 @@ def tokenize_and_prepare_labels(self, batch, prepare_labels, prepare_entities=Tr Returns: Dictionary containing tokenized inputs, entity encodings, and optionally labels. """ + labels_gather_indices = None + prompts_embedding_mask = None if prepare_entities: - if isinstance(batch["classes_to_id"], dict): - entities = list(batch["classes_to_id"]) - else: - entities = list(batch["classes_to_id"][0]) + entities, labels_gather_indices, prompts_embedding_mask = self._prepare_entity_batch( + batch["classes_to_id"] + ) else: entities = None tokenized_input = self.tokenize_inputs(batch["tokens"], entities) + if labels_gather_indices is not None: + tokenized_input["labels_gather_indices"] = labels_gather_indices + tokenized_input["prompts_embedding_mask"] = prompts_embedding_mask if prepare_labels: labels = self.create_labels(batch) @@ -1777,24 +1829,19 @@ def collate_raw_batch( else: rel_class_to_ids, rel_id_to_classes = make_mapping(relation_types or []) - if isinstance(class_to_ids, list): - batch = [ - self.preprocess_example( - b["tokenized_text"], - b[key], - class_to_ids[i], - b.get("relations", []), - rel_class_to_ids[i] if isinstance(rel_class_to_ids, list) else rel_class_to_ids, - ) - for i, b in enumerate(batch_list) - ] - else: - batch = [ + batch = [] + for i, example in enumerate(batch_list): + class_to_id_i = class_to_ids[i] if isinstance(class_to_ids, list) else class_to_ids + rel_class_to_id_i = rel_class_to_ids[i] if isinstance(rel_class_to_ids, list) else rel_class_to_ids + batch.append( self.preprocess_example( - b["tokenized_text"], b[key], class_to_ids, b.get("relations", []), rel_class_to_ids + example["tokenized_text"], + example[key], + class_to_id_i, + example.get("relations", []), + rel_class_to_id_i, ) - for b in batch_list - ] + ) return self.create_batch_dict(batch, class_to_ids, id_to_classes, rel_class_to_ids, rel_id_to_classes) diff --git a/gliner/decoding/decoder.py b/gliner/decoding/decoder.py index c48fb0b7..f49336a6 100644 --- a/gliner/decoding/decoder.py +++ b/gliner/decoding/decoder.py @@ -25,6 +25,40 @@ def _threshold_compare_tensor(threshold, batch_size: int, device, dims: int): return threshold +def _get_valid_classes_mask( + num_classes, id_to_classes: Union[Dict[int, str], List[Dict[int, str]]], device +) -> torch.Tensor: + """ + Create a boolean mask indicating valid classes for each batch item. + + Args: + num_classes (int): Total number of classes (C). + id_to_classes (Union[Dict[int, str], List[Dict[int, str]]]): Mapping from class IDs to class names. + device: Device on which to create the tensor. + + Returns: + torch.Tensor: Boolean tensor of shape (C,) for a shared mapping or + (B, C) for per-sample mappings, where True indicates a valid class. + """ + if isinstance(id_to_classes, list): + # For batch-level decoding, we need to create a mask for each batch 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_classes + ], + dtype=torch.bool, + device=device, + ) + else: + valid_classes = torch.tensor( + [class_idx + 1 in id_to_classes for class_idx in range(num_classes)], + dtype=torch.bool, + device=device, + ) + return valid_classes + + @dataclass class Span: """Represents a detected entity span with its properties. @@ -192,21 +226,30 @@ def _get_top_k_class_probs( Dict[str, float]: Dictionary mapping class names to probabilities, sorted by probability in descending order, containing up to k entries. """ - # Get the actual number of classes (might be less than k) num_classes = probs_tensor.shape[0] - k = min(k, num_classes) + valid_indices = [ + class_id - 1 + for class_id in sorted(id_to_class) + if 1 <= class_id <= num_classes + ] + if not valid_indices: + return {} - # Get top-k probabilities and their indices - top_probs, top_indices = torch.topk(probs_tensor, k=k, sorted=True) + valid_indices_tensor = torch.tensor(valid_indices, dtype=torch.long, device=probs_tensor.device) + valid_probs = probs_tensor.index_select(0, valid_indices_tensor) + k = min(k, len(valid_indices)) + + # Rank only real classes from this sample's mapping. The model output may + # include batch-padding slots that must never appear in class_probs. + top_probs, top_positions = torch.topk(valid_probs, k=k, sorted=True) + top_indices = valid_indices_tensor[top_positions] # Convert to dict, mapping class names to probabilities # Note: class indices are 1-indexed (0 is padding), so we add 1 - class_probs = {} - for idx, prob in zip(top_indices.tolist(), top_probs.tolist()): - class_name = id_to_class.get(idx + 1, f"class_{idx}") - class_probs[class_name] = prob - - return class_probs + return { + id_to_class[idx + 1]: prob + for idx, prob in zip(top_indices.tolist(), top_probs.tolist()) + } @abstractmethod def _build_span_tuple( @@ -276,16 +319,24 @@ def _decode_batch_item( Returns: List[tuple]: List of decoded span tuples for this sample. """ + device = probs_i.device # Mask probabilities to only include input spans (for efficiency) if input_spans_i is not None: L, K_dim, _ = probs_i.shape - span_filter = torch.zeros(L, K_dim, dtype=torch.bool, device=probs_i.device) + span_filter = torch.zeros(L, K_dim, dtype=torch.bool, device=device) for word_start, word_end in input_spans_i: width = word_end - word_start if 0 <= width < K_dim and 0 <= word_start < L: 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 any(class_idx + 1 not in id_to_class_i for class_idx in range(num_classes)): + valid_classes = _get_valid_classes_mask(num_classes, id_to_class_i, device) + probs_i = probs_i.masked_fill(~valid_classes, float("-inf")) + span_i = [] # Find all spans above threshold @@ -367,6 +418,8 @@ class IDs to class names. Returns: List[List[Span]]: For each sample in batch, list of Span objects. """ + device = probs.device + B, L, K_dim, C = probs.shape thresholds = _expand_batch_param(threshold, B, "threshold") flat_ner_values = _expand_batch_param(flat_ner, B, "flat_ner") @@ -396,7 +449,7 @@ class IDs to class names. # Apply input_spans mask at batch level (one mask, one multiply) if input_spans is not None: - span_filter = torch.zeros(B, L, K_dim, dtype=torch.bool, device=probs.device) + span_filter = torch.zeros(B, L, K_dim, dtype=torch.bool, device=device) for i, spans_i in enumerate(input_spans): if spans_i is not None: for word_start, word_end in spans_i: @@ -408,6 +461,21 @@ 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( + any(class_idx + 1 not in mapping for class_idx in range(num_classes)) + for mapping in id_to_class_per_item + ): + valid_classes = _get_valid_classes_mask(num_classes, id_to_class_per_item, device) + probs = probs.masked_fill(~valid_classes[:, None, None, :], float("-inf")) + # 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 +516,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)): @@ -458,10 +523,11 @@ class IDs to class names. class_probs = None if return_class_probs: - class_probs = {} - for idx, prob in zip(top_indices_list[j], top_probs_list[j]): - class_name = id_to_class_i.get(idx + 1, f"class_{idx}") - class_probs[class_name] = prob + class_probs = { + id_to_class_i[idx + 1]: prob + for idx, prob in zip(top_indices_list[j], top_probs_list[j]) + if idx + 1 in id_to_class_i + } span = self._build_span_tuple(s, k, c, flat_idx, score, id_to_class_i, span_label_maps[b], class_probs) batch_spans[b].append(span) @@ -641,6 +707,10 @@ def _decode_explicit_spans( top_index_rows = None if return_class_probs: candidate_probabilities = probabilities[batch_indices, span_positions] + candidate_probabilities = candidate_probabilities.masked_fill( + ~valid_classes[batch_indices], + float("-inf"), + ) top_k = min(5, num_classes) top_indices = torch.argsort( candidate_probabilities, @@ -660,11 +730,12 @@ def _decode_explicit_spans( class_probs = None if return_class_probs: class_probs = { - id_to_class.get(index + 1, f"class_{index}"): probability + id_to_class[index + 1]: probability for index, probability in zip( top_index_rows[row_index], top_prob_rows[row_index], ) + if index + 1 in id_to_class } candidates_by_batch[batch_idx].append( Span( @@ -1026,9 +1097,19 @@ def _decode_relations_batch( """ relations: List[List[tuple]] = [[] for _ in range(batch_size)] + rel_id_to_class_per_item = [ + rel_id_to_classes[i] if isinstance(rel_id_to_classes, list) else rel_id_to_classes + for i in range(batch_size) + ] + # 1. Sigmoid — one kernel rel_probs = torch.sigmoid(rel_logits) + num_classes = rel_probs.shape[-1] + if any(len(mapping) < num_classes for mapping in rel_id_to_class_per_item): + valid_classes = _get_valid_classes_mask(num_classes, rel_id_to_class_per_item, rel_probs.device) + rel_probs = rel_probs.masked_fill(~valid_classes[:, None, :], float("-inf")) + # 2. Apply relation mask — zeros out padded relations rel_probs = rel_probs * rel_mask.unsqueeze(-1) @@ -1053,14 +1134,11 @@ def _decode_relations_batch( b_list = b_idx.tolist() c_list = c_idx.tolist() - # 6. Pre-resolve per-sample class mappings - is_list = isinstance(rel_id_to_classes, list) - - # 7. Pure-Python grouping — no more GPU access + # 6. Pure-Python grouping — no more GPU access for k in range(len(b_list)): b = b_list[k] c1 = c_list[k] + 1 # class IDs are 1-indexed - mapping = rel_id_to_classes[b] if is_list else rel_id_to_classes + mapping = rel_id_to_class_per_item[b] if c1 not in mapping: continue relations[b].append((int(head_list[k]), mapping[c1], int(tail_list[k]), scores[k])) @@ -1215,7 +1293,15 @@ def _decode_relations( if rel_mask is None: rel_mask = torch.ones(rel_idx[..., 0].shape, dtype=torch.bool, device=rel_idx.device) + rel_id_to_class_per_item = [ + self._get_id_to_class_for_sample(rel_id_to_classes, i) for i in range(batch_size) + ] + rel_probs = torch.sigmoid(rel_logits) + num_classes = rel_probs.shape[-1] + if any(len(mapping) < num_classes for mapping in rel_id_to_class_per_item): + valid_classes = _get_valid_classes_mask(num_classes, rel_id_to_class_per_item, rel_probs.device) + rel_probs = rel_probs.masked_fill(~valid_classes[:, None, :], float("-inf")) # Batch CPU transfer to avoid per-element .item() sync rel_idx_cpu = rel_idx.tolist() @@ -1230,7 +1316,7 @@ def _decode_relations( # Decode relations for each sample thresholds = _expand_batch_param(threshold, batch_size, "relation_threshold") for i in range(batch_size): - rel_id_to_class_i = rel_id_to_classes[i] if isinstance(rel_id_to_classes, list) else rel_id_to_classes + rel_id_to_class_i = rel_id_to_class_per_item[i] idx_map = idx_mappings[i] num_spans_i = len(spans[i]) threshold_i = thresholds[i] @@ -1583,16 +1669,30 @@ class IDs to class names. # Check if token-level decoding is requested if model_output is not None: + batch_size = len(tokens) + num_classes = model_output.shape[-2] + id_to_class_per_item = [ + self._get_id_to_class_for_sample(id_to_classes, i) for i in range(batch_size) + ] + + # Per-sample label sets share a batch-wide class dimension. Exclude padded + # class slots before candidate search so they cannot produce unmapped spans. + if any(len(mapping) < num_classes for mapping in id_to_class_per_item): + valid_classes = _get_valid_classes_mask(num_classes, id_to_class_per_item, model_output.device) + model_output = model_output.masked_fill( + ~valid_classes[:, None, :, None], + float("-inf"), + ) + model_output = model_output.permute(3, 0, 1, 2) scores_start, scores_end, scores_inside = model_output - batch_size = len(tokens) thresholds = _expand_batch_param(threshold, batch_size, "threshold") flat_ner_values = _expand_batch_param(flat_ner, batch_size, "flat_ner") multi_label_values = _expand_batch_param(multi_label, batch_size, "multi_label") spans = [] for i, _ in enumerate(tokens): - id_to_class_i = self._get_id_to_class_for_sample(id_to_classes, i) + 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] span_scores = self._calculate_span_score( @@ -1716,7 +1816,15 @@ def _decode_relations( if rel_mask is None: rel_mask = torch.ones(rel_idx[..., 0].shape, dtype=torch.bool, device=rel_idx.device) + rel_id_to_class_per_item = [ + self._get_id_to_class_for_sample(rel_id_to_classes, i) for i in range(batch_size) + ] + rel_probs = torch.sigmoid(rel_logits) + num_classes = rel_probs.shape[-1] + if any(len(mapping) < num_classes for mapping in rel_id_to_class_per_item): + valid_classes = _get_valid_classes_mask(num_classes, rel_id_to_class_per_item, rel_probs.device) + rel_probs = rel_probs.masked_fill(~valid_classes[:, None, :], float("-inf")) # Batch CPU transfer to avoid per-element .item() sync rel_idx_cpu = rel_idx.tolist() @@ -1731,7 +1839,7 @@ def _decode_relations( # Decode relations for each sample thresholds = _expand_batch_param(threshold, batch_size, "relation_threshold") for i in range(batch_size): - rel_id_to_class_i = rel_id_to_classes[i] if isinstance(rel_id_to_classes, list) else rel_id_to_classes + rel_id_to_class_i = rel_id_to_class_per_item[i] idx_map = idx_mappings[i] num_spans_i = len(spans[i]) threshold_i = thresholds[i] diff --git a/gliner/evaluation/evaluator.py b/gliner/evaluation/evaluator.py index 399d5e06..8b66ccb1 100644 --- a/gliner/evaluation/evaluator.py +++ b/gliner/evaluation/evaluator.py @@ -265,9 +265,9 @@ def get_predictions(self, ents, rels): else: t_ent_start = t_ent[0] t_ent_end = t_ent[1] - + all_rels.append([lab, (h_ent_start, h_ent_end, t_ent_start, t_ent_end)]) - + return all_rels def transform_data(self): diff --git a/gliner/model.py b/gliner/model.py index 97cabaac..9bc1ebff 100644 --- a/gliner/model.py +++ b/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 @@ -2359,6 +2372,17 @@ def run_batch( Returns: Model output containing logits and span information. """ + labels_gather_indices = batch.get("labels_gather_indices") + if ( + self.onnx_model + and isinstance(labels_gather_indices, torch.Tensor) + and labels_gather_indices.shape[0] > 1 + ): + raise ValueError( + "Batched per-row labels are not supported by existing bi-encoder ONNX graphs; " + "use inference() or collate singleton batches" + ) + if move_to_device and not self.onnx_model: batch = {k: v.to(self.device) if isinstance(v, torch.Tensor) else v for k, v in batch.items()} @@ -2510,12 +2534,31 @@ 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, + ) + + loader_batch_size = batch_size + entity_types = prepared["entity_types"] + has_per_row_entity_types = bool(entity_types and isinstance(entity_types[0], list)) + needs_per_row_label_layout = has_per_row_entity_types and ( + not entity_types[0] or any(row != entity_types[0] for row in entity_types[1:]) + ) + if ( + self.onnx_model + and isinstance(self.data_processor, (BiEncoderSpanProcessor, BiEncoderTokenProcessor)) + and needs_per_row_label_layout + ): + # Existing bi-encoder ONNX graphs expose one shared label matrix and + # cannot consume the per-row gather metadata used by the PyTorch model. + loader_batch_size = 1 data_loader = torch.utils.data.DataLoader( - prepared["input_x"], - batch_size=batch_size, + list(range(len(prepared["input_x"]))), + batch_size=loader_batch_size, shuffle=False, collate_fn=collate_fn, ) @@ -4937,11 +4980,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 +5532,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, + _entity_types_for_chunk(prepared["relation_types"], indices), + ) 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, @@ -5726,7 +5778,8 @@ def evaluate( relation_threshold: The threshold for relation predictions. Defaults to threshold. batch_size: The batch size for evaluation. Defaults to 12. entity_types: Optional list of entity types to evaluate. If None, extracts from test data. Defaults to None. - relation_types: Optional list of relation types to evaluate. If None, extracts from test data. Defaults to None. + relation_types: Optional list of relation types to evaluate. If None, extracts from test + data. Defaults to None. Returns: Tuple of ((ner_output, ner_f1), (rel_output, rel_f1)) containing: diff --git a/gliner/modeling/base.py b/gliner/modeling/base.py index 886cd10f..f0652033 100644 --- a/gliner/modeling/base.py +++ b/gliner/modeling/base.py @@ -1323,6 +1323,8 @@ def get_representations( labels_attention_mask: Optional[torch.LongTensor] = None, text_lengths: Optional[torch.Tensor] = None, words_mask: Optional[torch.LongTensor] = None, + labels_gather_indices: Optional[torch.LongTensor] = None, + prompts_embedding_mask: Optional[torch.LongTensor] = None, **kwargs: Any, ) -> GLiNERRepresentationOutput: """Get entity label and word representations using bi-encoder. @@ -1335,6 +1337,8 @@ def get_representations( labels_attention_mask: Attention mask for labels. text_lengths: Length of each text in batch. words_mask: Word boundary mask. + labels_gather_indices: Per-row indices into the shared label embeddings. + prompts_embedding_mask: Mask for valid labels in each row. **kwargs: Additional arguments for the encoder. Returns: @@ -1365,16 +1369,41 @@ def get_representations( getattr(self.config, "subtoken_pooling", "first"), ) - labels_embeds = labels_embeds.unsqueeze(0) - labels_embeds = labels_embeds.expand(batch_size, -1, -1) - labels_mask = torch.ones(labels_embeds.shape[:-1], dtype=attention_mask.dtype, device=attention_mask.device) + if labels_gather_indices is not None: + if labels_embeds.dim() != 2: + raise ValueError("labels_gather_indices requires 2D shared label embeddings") + if labels_gather_indices.dim() != 2: + raise ValueError("labels_gather_indices must have shape (B, C)") + if labels_gather_indices.shape[0] != batch_size: + raise ValueError("labels_gather_indices must have one row per input text") + + labels_gather_indices = labels_gather_indices.to(device=labels_embeds.device, dtype=torch.long) + labels_embeds = labels_embeds[labels_gather_indices] + elif labels_embeds.dim() == 2: + labels_embeds = labels_embeds.unsqueeze(0).expand(batch_size, -1, -1) + elif labels_embeds.dim() == 3: + if labels_embeds.shape[0] != batch_size: + raise ValueError("Batched label embeddings must have one row per input text") + else: + raise ValueError("labels_embeds must have shape (C, D) or (B, C, D)") + + if prompts_embedding_mask is None: + labels_mask = torch.ones( + labels_embeds.shape[:-1], dtype=attention_mask.dtype, device=labels_embeds.device + ) + else: + if prompts_embedding_mask.shape != labels_embeds.shape[:-1]: + raise ValueError("prompts_embedding_mask must match the batched label layout") + labels_mask = prompts_embedding_mask.to(device=labels_embeds.device, dtype=attention_mask.dtype) labels_embeds = labels_embeds.to(words_embedding.dtype) + labels_embeds = labels_embeds * labels_mask.unsqueeze(-1).to(labels_embeds.dtype) if hasattr(self, "cross_fuser"): words_embedding, labels_embeds = self.features_enhancement( words_embedding, labels_embeds, text_mask=mask, labels_mask=labels_mask ) + labels_embeds = labels_embeds * labels_mask.unsqueeze(-1).to(labels_embeds.dtype) return GLiNERRepresentationOutput( prompts_embedding=labels_embeds, @@ -1429,6 +1458,7 @@ def forward( span_idx: Optional[torch.LongTensor] = None, span_mask: Optional[torch.LongTensor] = None, labels: Optional[torch.FloatTensor] = None, + labels_gather_indices: Optional[torch.LongTensor] = None, **kwargs: Any, ) -> GLiNERBaseOutput: """Forward pass through the bi-encoder span model. @@ -1448,6 +1478,7 @@ def forward( span_idx: Span indices of shape (B, L*K, 2). span_mask: Mask for valid spans of shape (B, L, K). labels: Ground truth labels of shape (B, L, K, C). + labels_gather_indices: Per-row indices into the shared label embeddings. **kwargs: Additional arguments. Returns: @@ -1460,13 +1491,15 @@ def forward( } representations = self.get_representations( - input_ids, - attention_mask, - labels_embeds, - labels_input_ids, - labels_attention_mask, - text_lengths, - words_mask, + input_ids=input_ids, + attention_mask=attention_mask, + labels_embeds=labels_embeds, + labels_input_ids=labels_input_ids, + labels_attention_mask=labels_attention_mask, + text_lengths=text_lengths, + words_mask=words_mask, + labels_gather_indices=labels_gather_indices, + prompts_embedding_mask=prompts_embedding_mask, **encoder_kwargs, ) prompts_embedding = representations.prompts_embedding @@ -1611,6 +1644,7 @@ def forward( text_lengths: Optional[torch.Tensor] = None, labels: Optional[torch.FloatTensor] = None, threshold: Optional[float] = 0.5, + labels_gather_indices: Optional[torch.LongTensor] = None, **kwargs: Any, ) -> GLiNERBaseOutput: """Forward pass through the bi-encoder token model. @@ -1633,6 +1667,7 @@ def forward( text_lengths: Length of each text sequence. labels: Ground truth labels of shape (B, W, C). threshold: float value for filtering spans. + labels_gather_indices: Per-row indices into the shared label embeddings. **kwargs: Additional arguments. Returns: @@ -1645,13 +1680,15 @@ def forward( } representations = self.get_representations( - input_ids, - attention_mask, - labels_embeds, - labels_input_ids, - labels_attention_mask, - text_lengths, - words_mask, + input_ids=input_ids, + attention_mask=attention_mask, + labels_embeds=labels_embeds, + labels_input_ids=labels_input_ids, + labels_attention_mask=labels_attention_mask, + text_lengths=text_lengths, + words_mask=words_mask, + labels_gather_indices=labels_gather_indices, + prompts_embedding_mask=prompts_embedding_mask, **encoder_kwargs, ) prompts_embedding = representations.prompts_embedding diff --git a/gliner/modeling/layers.py b/gliner/modeling/layers.py index 6e44c505..e7bd2079 100644 --- a/gliner/modeling/layers.py +++ b/gliner/modeling/layers.py @@ -179,11 +179,17 @@ def forward( else: value = self.transpose_for_scores(self.value_layer(value)) + attention_mask = attn_mask.to(torch.bool) if attn_mask is not None else head_mask + # CrossFuser builds masks with shape (B, Q, K). SDPA expects them to + # broadcast across the attention-head dimension. + if attention_mask is not None and attention_mask.dim() == 3: + attention_mask = attention_mask.unsqueeze(1) + context_layer = torch.nn.functional.scaled_dot_product_attention( query, key, value, - head_mask, + attention_mask, self.attention_probs_dropout_prob if self.training else 0.0, is_causal=False, scale=None, diff --git a/tests/test_data_processing.py b/tests/test_data_processing.py index ea049b07..6e42ed0b 100644 --- a/tests/test_data_processing.py +++ b/tests/test_data_processing.py @@ -7,6 +7,7 @@ import pytest from transformers import AutoTokenizer +from gliner.data_processing import BiEncoderSpanProcessor, BiEncoderTokenProcessor from gliner.data_processing.utils import make_mapping, get_negatives, pad_2d_tensor, prepare_span_idx, prepare_word_mask @@ -1184,6 +1185,143 @@ def test_batch_generate_class_mappings_with_negatives(self, processor): assert "MISC" in all_types +@pytest.mark.parametrize("processor_class", [BiEncoderSpanProcessor, BiEncoderTokenProcessor]) +@pytest.mark.parametrize( + ("classes_to_id", "expected_entities", "expected_indices", "expected_mask"), + [ + ( + [{"B": 2, "A": 1}, {"D": 2, "C": 1}], + ["A", "B", "C", "D"], + [[0, 1], [2, 3]], + [[True, True], [True, True]], + ), + ( + [{"A": 1, "B": 2}, {"C": 1}], + ["A", "B", "C"], + [[0, 1], [2, 0]], + [[True, True], [True, False]], + ), + ( + [{"A": 1, "B": 2}, {"B": 1, "C": 2}], + ["A", "B", "C"], + [[0, 1], [1, 2]], + [[True, True], [True, True]], + ), + ( + [{}, {"C": 1}], + ["C"], + [[0], [0]], + [[False], [True]], + ), + ( + [{}, {}], + [""], + [[], []], + [[], []], + ), + ], +) +def test_biencoder_processors_preserve_per_row_label_layout( + mock_config, + mock_tokenizer, + mock_words_splitter, + processor_class, + classes_to_id, + expected_entities, + expected_indices, + expected_mask, +): + """Both bi-encoder processors should encode the union and restore each row's order.""" + encoded_values = {"": 0, "A": 10, "B": 20, "C": 30, "D": 40} + encoded_entities = [] + labels_tokenizer = Mock() + labels_tokenizer.unk_token = "[UNK]" + labels_tokenizer.pad_token = "[PAD]" + + def tokenize_labels(entities, **kwargs): + encoded_entities.append(list(entities)) + return { + "input_ids": torch.tensor([[encoded_values[label]] for label in entities]), + "attention_mask": torch.ones(len(entities), 1, dtype=torch.long), + } + + labels_tokenizer.side_effect = tokenize_labels + processor = processor_class(mock_config, mock_tokenizer, mock_words_splitter, labels_tokenizer) + batch = { + "tokens": [["first"], ["second"]], + "classes_to_id": classes_to_id, + } + + result = processor.tokenize_and_prepare_labels(batch, prepare_labels=False) + + assert encoded_entities[-1] == expected_entities + assert result["labels_input_ids"].squeeze(-1).tolist() == [encoded_values[label] for label in expected_entities] + assert result["labels_gather_indices"].tolist() == expected_indices + assert result["prompts_embedding_mask"].tolist() == expected_mask + + +@pytest.mark.parametrize("processor_class", [BiEncoderSpanProcessor, BiEncoderTokenProcessor]) +@pytest.mark.parametrize( + "classes_to_id", + [ + {"A": 1, "B": 2}, + [{"A": 1, "B": 2}, {"A": 1, "B": 2}], + ], +) +def test_biencoder_processors_keep_shared_label_fast_path( + mock_config, mock_tokenizer, mock_words_splitter, processor_class, classes_to_id +): + """A flat shared mapping should retain the existing 2D label-embedding path.""" + labels_tokenizer = Mock() + labels_tokenizer.unk_token = "[UNK]" + labels_tokenizer.pad_token = "[PAD]" + labels_tokenizer.return_value = { + "input_ids": torch.tensor([[10], [20]]), + "attention_mask": torch.ones(2, 1, dtype=torch.long), + } + processor = processor_class(mock_config, mock_tokenizer, mock_words_splitter, labels_tokenizer) + batch = { + "tokens": [["first"], ["second"]], + "classes_to_id": classes_to_id, + } + + result = processor.tokenize_and_prepare_labels(batch, prepare_labels=False) + + assert labels_tokenizer.call_args.args[0] == ["A", "B"] + assert "labels_gather_indices" not in result + assert "prompts_embedding_mask" not in result + + +@pytest.mark.parametrize("processor_class", [BiEncoderSpanProcessor, BiEncoderTokenProcessor]) +def test_biencoder_per_row_training_targets_keep_local_class_ids( + mock_config, mock_tokenizer, mock_words_splitter, processor_class +): + """Per-row embedding gathering must stay aligned with inherited training targets.""" + labels_tokenizer = Mock() + labels_tokenizer.unk_token = "[UNK]" + labels_tokenizer.pad_token = "[PAD]" + labels_tokenizer.return_value = { + "input_ids": torch.tensor([[10], [20], [30]]), + "attention_mask": torch.ones(3, 1, dtype=torch.long), + } + processor = processor_class(mock_config, mock_tokenizer, mock_words_splitter, labels_tokenizer) + batch = { + "tokens": [["first"], ["second"]], + "seq_length": torch.ones(2, 1, dtype=torch.long), + "classes_to_id": [{"A": 1, "B": 2}, {"C": 1}], + "entities": [[(0, 0, "B")], [(0, 0, "C")]], + } + + result = processor.tokenize_and_prepare_labels(batch, prepare_labels=True) + + if processor_class is BiEncoderSpanProcessor: + assert result["labels"][0, 0].tolist() == [0.0, 1.0] + assert result["labels"][1, 0].tolist() == [1.0, 0.0] + else: + assert result["labels"][0, 0, :, 0].tolist() == [0.0, 1.0] + assert result["labels"][1, 0, :, 0].tolist() == [1.0, 0.0] + + class TestUniEncoderSpanDecoderProcessor: """Test suite for UniEncoderSpanDecoderProcessor.""" @@ -1348,6 +1486,31 @@ def test_create_batch_dict_includes_relation_mappings(self, processor): assert result["rel_idx"].shape[0] == 1 # batch size assert result["rel_label"].shape[0] == 1 + def test_collate_raw_batch_selects_relation_mapping_per_sample(self, processor): + """Shared entity labels and per-sample relation labels should remain independently aligned.""" + processor.config.augment_data_prob = 0.0 + batch = [ + { + "tokenized_text": ["Alice", "Acme"], + "ner": [(0, 0, "PER"), (1, 1, "ORG")], + "relations": [(0, 1, "WORKS_FOR")], + }, + { + "tokenized_text": ["Kyiv", "Ukraine"], + "ner": [(0, 0, "PER"), (1, 1, "ORG")], + "relations": [(0, 1, "LOCATED_IN")], + }, + ] + + result = processor.collate_raw_batch( + batch, + entity_types=["PER", "ORG"], + relation_types=[["WORKS_FOR"], ["LOCATED_IN"]], + ) + + assert result["rel_label"].tolist() == [[1], [1]] + assert result["rel_id_to_classes"] == [{1: "WORKS_FOR"}, {1: "LOCATED_IN"}] + def test_prepare_inputs_with_relations(self, processor): """Should add relation tokens to input.""" texts = [["word1", "word2"]] diff --git a/tests/test_decoder.py b/tests/test_decoder.py index eafd1072..73cec7d8 100644 --- a/tests/test_decoder.py +++ b/tests/test_decoder.py @@ -9,6 +9,7 @@ TokenDecoder, BaseSpanDecoder, SpanRelexDecoder, + TokenRelexDecoder, SpanGenerativeDecoder, _decode_relations_batch, ) @@ -207,6 +208,119 @@ 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_return_class_probs_excludes_padded_classes_single_item(self, basic_config): + """The B == 1 path should report probabilities only for mapped classes.""" + decoder = SpanDecoder(basic_config) + logits = torch.tensor([[[[5.0, 10.0, 9.0]]]]) + + result = decoder.decode( + tokens=[["Alice"]], + id_to_classes=[{1: "PERSON"}], + model_output=logits, + threshold=0.5, + return_class_probs=True, + ) + + assert len(result[0]) == 1 + assert result[0][0].class_probs == pytest.approx( + {"PERSON": torch.sigmoid(torch.tensor(5.0)).item()} + ) + + def test_return_class_probs_excludes_padded_classes_batched(self, basic_config): + """The vectorized path should not expose another row's padded class slots.""" + decoder = SpanDecoder(basic_config) + logits = torch.tensor( + [ + [[[5.0, -1.0, 100.0]]], + [[[4.0, 100.0, 100.0]]], + ] + ) + + result = decoder.decode( + tokens=[["Alice"], ["Kyiv"]], + id_to_classes=[ + {1: "PERSON", 2: "ORG"}, + {1: "PLACE"}, + ], + model_output=logits, + threshold=0.5, + return_class_probs=True, + ) + + assert len(result[0]) == 1 + assert len(result[1]) == 1 + assert result[0][0].class_probs == pytest.approx( + { + "PERSON": torch.sigmoid(torch.tensor(5.0)).item(), + "ORG": torch.sigmoid(torch.tensor(-1.0)).item(), + } + ) + assert result[1][0].class_probs == pytest.approx( + {"PLACE": torch.sigmoid(torch.tensor(4.0)).item()} + ) + + def test_return_class_probs_excludes_padded_classes_explicit_spans(self, basic_config): + """Explicit-span decoding should rank only classes present in the row mapping.""" + decoder = SpanDecoder(basic_config) + logits = torch.tensor([[[5.0, 4.0, 100.0, 90.0, 80.0, 70.0, 60.0]]]) + + result = decoder.decode( + tokens=[["Alice"]], + id_to_classes=[{1: "PERSON", 2: "ORG"}], + model_output=logits, + span_idx=torch.tensor([[[0, 0]]]), + span_mask=torch.tensor([[True]]), + threshold=0.5, + return_class_probs=True, + ) + + assert len(result[0]) == 1 + assert result[0][0].class_probs == pytest.approx( + { + "PERSON": torch.sigmoid(torch.tensor(5.0)).item(), + "ORG": torch.sigmoid(torch.tensor(4.0)).item(), + } + ) + def test_empty_predictions(self, basic_config): """Should handle case with no predictions above threshold.""" decoder = SpanDecoder(basic_config) @@ -677,6 +791,37 @@ def test_handles_missing_relation_outputs(self, relex_config, relex_inputs): assert all(len(rels) == 0 for rels in relations) +@pytest.mark.parametrize("decoder_class", [SpanRelexDecoder, TokenRelexDecoder]) +def test_relex_decoders_mask_ragged_relation_classes(decoder_class): + """Padded relation classes must not enter decoding for a row with fewer labels.""" + decoder = decoder_class(Mock()) + spans = [ + [Span(0, 0, "A", 0.9), Span(1, 1, "B", 0.8)], + [Span(0, 0, "X", 0.9), Span(1, 1, "Y", 0.8)], + ] + rel_idx = torch.tensor([[[0, 1]], [[0, 1]]]) + rel_logits = torch.full((2, 1, 2), -10.0) + rel_logits[0, 0, 1] = 5.0 # valid class 2 for sample 0 + rel_logits[1, 0, 0] = 5.0 # valid class 1 for sample 1 + rel_logits[1, 0, 1] = 10.0 # padded class 2 for sample 1 + + decode_kwargs = { + "spans": spans, + "rel_idx": rel_idx, + "rel_logits": rel_logits, + "rel_mask": torch.ones(2, 1, dtype=torch.bool), + "rel_id_to_classes": [{1: "REL_A", 2: "REL_B"}, {1: "REL_C"}], + "threshold": 0.1, + "batch_size": 2, + } + if decoder_class is SpanRelexDecoder: + decode_kwargs["model_output"] = None + + relations = decoder._decode_relations(**decode_kwargs) + + assert [[relation[1] for relation in sample] for sample in relations] == [["REL_B"], ["REL_C"]] + + class TestTokenDecoder: """Test suite for TokenDecoder class.""" @@ -858,6 +1003,44 @@ def test_handles_empty_predictions(self, token_config): assert len(result) == 1 assert len(result[0]) == 0 + def test_ragged_per_sample_id_to_classes(self, token_config): + """Should ignore padded class slots for per-sample mappings of different sizes.""" + decoder = TokenDecoder(token_config) + + model_output = torch.full((2, 2, 2, 3), -10.0) + model_output[0, 0, 0] = 5.0 # valid class 1 for sample 0 + model_output[1, 1, 0] = 5.0 # valid class 1 for sample 1 + model_output[1, 0, 1] = 10.0 # padded class 2 for sample 1 + + result = decoder.decode( + tokens=[["Alice", "x"], ["y", "Kyiv"]], + id_to_classes=[{1: "PERSON", 2: "ORG"}, {1: "LOCATION"}], + model_output=model_output, + threshold=0.1, + ) + + assert [[(span.start, span.end, span.entity_type) for span in spans] for spans in result] == [ + [(0, 0, "PERSON")], + [(1, 1, "LOCATION")], + ] + + def test_ragged_per_sample_id_to_classes_single_item(self, token_config): + """Should also ignore padded class slots for a single decoded item.""" + decoder = TokenDecoder(token_config) + + model_output = torch.full((1, 2, 3, 3), -10.0) + model_output[0, 0, 0] = 5.0 # valid class 1 + model_output[0, 1, 2] = 10.0 # padded class 3 + + result = decoder.decode( + tokens=[["Alice", "x"]], + id_to_classes=[{1: "PERSON"}], + model_output=model_output, + threshold=0.1, + ) + + assert [(span.start, span.end, span.entity_type) for span in result[0]] == [(0, 0, "PERSON")] + def test_per_sample_thresholds(self, token_config, token_inputs): """Should apply token-decoder thresholds independently per batch item.""" decoder = TokenDecoder(token_config) diff --git a/tests/test_modeling.py b/tests/test_modeling.py index 6cc89750..d71033d3 100644 --- a/tests/test_modeling.py +++ b/tests/test_modeling.py @@ -6,6 +6,7 @@ from gliner.config import ( BaseGLiNERConfig, BiEncoderSpanConfig, + BiEncoderTokenConfig, UniEncoderSpanConfig, UniEncoderTokenConfig, UniEncoderSpanRelexConfig, @@ -13,7 +14,9 @@ ) from gliner.modeling.base import ( BaseModel, + BaseBiEncoderModel, BiEncoderSpanModel, + BiEncoderTokenModel, UniEncoderSpanModel, UniEncoderTokenModel, UniEncoderSpanRelexModel, @@ -25,6 +28,7 @@ extract_word_embeddings, extract_prompt_features_and_word_embeddings, ) +from gliner.modeling.layers import MultiheadAttention class TestExtractWordEmbeddings: @@ -739,6 +743,72 @@ def loss(self, x): pass assert losses.shape == (B, N, C) +def test_biencoder_representations_gather_labels_for_each_row(): + """Shared union embeddings should be restored to each row's local class order.""" + + class TokenRepresentationLayer: + def __call__(self, input_ids, attention_mask, labels_input_ids, labels_attention_mask, **kwargs): + token_embeddings = torch.zeros(*input_ids.shape, 1) + return token_embeddings, labels_input_ids.float() + + model = Mock() + model.token_rep_layer = TokenRepresentationLayer() + model.config.subtoken_pooling = "first" + del model.cross_fuser + + result = BaseBiEncoderModel.get_representations( + model, + input_ids=torch.ones(2, 2, dtype=torch.long), + attention_mask=torch.ones(2, 2, dtype=torch.long), + labels_input_ids=torch.tensor([[10], [20], [30]]), + labels_attention_mask=torch.ones(3, 1, dtype=torch.long), + text_lengths=torch.ones(2, 1, dtype=torch.long), + words_mask=torch.tensor([[1, 0], [1, 0]]), + labels_gather_indices=torch.tensor([[0, 1], [2, 0]]), + prompts_embedding_mask=torch.tensor([[True, True], [True, False]]), + ) + + assert result.prompts_embedding.squeeze(-1).tolist() == [[10.0, 20.0], [30.0, 0.0]] + assert result.prompts_embedding_mask.tolist() == [[1, 1], [1, 0]] + + empty_result = BaseBiEncoderModel.get_representations( + model, + input_ids=torch.ones(2, 2, dtype=torch.long), + attention_mask=torch.ones(2, 2, dtype=torch.long), + labels_input_ids=torch.tensor([[0]]), + labels_attention_mask=torch.ones(1, 1, dtype=torch.long), + text_lengths=torch.ones(2, 1, dtype=torch.long), + words_mask=torch.tensor([[1, 0], [1, 0]]), + labels_gather_indices=torch.empty(2, 0, dtype=torch.long), + prompts_embedding_mask=torch.empty(2, 0, dtype=torch.bool), + ) + + assert empty_result.prompts_embedding.shape == (2, 0, 1) + assert empty_result.prompts_embedding_mask.shape == (2, 0) + + +def test_multihead_attention_applies_per_row_attention_mask(): + """Changing masked keys must not change attention output for either row.""" + torch.manual_seed(0) + attention = MultiheadAttention(hidden_size=6, num_heads=3, dropout=0.0).eval() + query = torch.randn(2, 1, 6) + key = torch.randn(2, 2, 6) + value = torch.randn(2, 2, 6) + mask = torch.tensor([[[1, 0]], [[0, 1]]], dtype=torch.long) + + changed_key = key.clone() + changed_value = value.clone() + changed_key[0, 1] += 1000 + changed_value[0, 1] += 1000 + changed_key[1, 0] -= 1000 + changed_value[1, 0] -= 1000 + + expected, _ = attention(query, key, value, attn_mask=mask) + actual, _ = attention(query, changed_key, changed_value, attn_mask=mask) + + assert torch.allclose(actual, expected) + + class TestUniEncoderSpanModel: """Test suite for UniEncoderSpanModel.""" @@ -1139,6 +1209,19 @@ def test_forward_output_shape_without_labels(self, mock_config, model_inputs): assert output.logits.shape[0] == B # Batch dimension assert output.logits.shape[1] == L # Sequence dimension + def test_forward_uses_per_row_label_layout(self, mock_config, model_inputs): + """Span forward should score the gathered local label layout, including padding.""" + inputs = {k: v for k, v in model_inputs.items() if k != "labels"} + inputs["labels_gather_indices"] = torch.tensor([[0, 2], [4, 0]]) + inputs["prompts_embedding_mask"] = torch.tensor([[True, True], [True, False]]) + model = BiEncoderSpanModel(mock_config, from_pretrained=False) + + with torch.no_grad(): + output = model(**inputs) + + assert output.logits.shape[-1] == 2 + assert output.prompts_embedding_mask.tolist() == [[1, 1], [1, 0]] + def test_forward_with_precomputed_labels_embeds(self, mock_config, model_inputs): """Should accept precomputed labels embeddings instead of ids.""" @@ -1248,3 +1331,34 @@ def test_loss_reduction_mean(self, mock_config): assert isinstance(loss, torch.Tensor) assert loss.ndim == 0 assert loss.item() >= 0 + + +def test_biencoder_token_forward_uses_per_row_label_layout(): + """Token forward should propagate the gather indices and per-row label mask.""" + config = BiEncoderTokenConfig( + model_name="bert-base-uncased", + labels_encoder="bert-base-uncased", + hidden_size=64, + dropout=0.1, + max_width=12, + class_token_index=103, + has_rnn=False, + post_fusion_schema="", + embed_ent_token=True, + ) + model = BiEncoderTokenModel(config, from_pretrained=False) + + with torch.no_grad(): + output = model( + input_ids=torch.randint(0, 1000, (2, 4)), + attention_mask=torch.ones(2, 4, dtype=torch.long), + labels_input_ids=torch.randint(0, 1000, (3, 4)), + labels_attention_mask=torch.ones(3, 4, dtype=torch.long), + words_mask=torch.tensor([[0, 1, 2, 0], [0, 1, 2, 0]]), + text_lengths=torch.tensor([[2], [2]]), + labels_gather_indices=torch.tensor([[0, 1], [2, 0]]), + prompts_embedding_mask=torch.tensor([[True, True], [True, False]]), + ) + + assert output.logits.shape == (2, 2, 2, 3) + assert output.prompts_embedding_mask.tolist() == [[1, 1], [1, 0]] diff --git a/tests/test_models.py b/tests/test_models.py index 7de13f9c..32d1c7bb 100644 --- a/tests/test_models.py +++ b/tests/test_models.py @@ -1,7 +1,14 @@ +import torch import pytest from gliner import GLiNER -from gliner.model import BaseEncoderGLiNER, UniEncoderSpanRelexGLiNER +from gliner.model import ( + BaseEncoderGLiNER, + BiEncoderSpanGLiNER, + UniEncoderSpanRelexGLiNER, + _entity_types_for_chunk, +) +from gliner.data_processing import BiEncoderSpanProcessor class _WordsSplitter: @@ -77,3 +84,145 @@ 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]) == [] + + +@pytest.mark.parametrize( + ("onnx_model", "entity_types", "expected_batch_sizes"), + [ + (False, [["a"], ["b"], ["c"]], [2, 1]), + (True, [["a"], ["b"], ["c"]], [1, 1, 1]), + (True, [["a"], ["a"], ["a"]], [2, 1]), + (True, ["a", "b"], [2, 1]), + ], +) +def test_biencoder_onnx_uses_singleton_batches_only_for_per_row_labels( + onnx_model, entity_types, expected_batch_sizes +): + """Legacy ONNX graphs need singleton chunks because they cannot consume gather metadata.""" + model = BiEncoderSpanGLiNER.__new__(BiEncoderSpanGLiNER) + torch.nn.Module.__init__(model) + model.onnx_model = onnx_model + model.data_processor = BiEncoderSpanProcessor.__new__(BiEncoderSpanProcessor) + model._inference_packing_config = None + texts = ["zero", "one", "two"] + prepared = { + "input_x": [{"row": i} for i in range(3)], + "entity_types": entity_types, + "valid_texts": texts, + "valid_to_orig_idx": list(range(3)), + "start_token_map": [[] for _ in texts], + "end_token_map": [[] for _ in texts], + "word_input_spans": None, + "num_original": 3, + } + observed_batch_sizes = [] + + model.prepare_batch = lambda *args, **kwargs: prepared + model.create_collator = object + model.collate_batch = lambda input_x, labels, collator: {"tokens": input_x} + + def process_batches(data_loader, *args, **kwargs): + batches = list(data_loader) + observed_batch_sizes.extend(len(batch["tokens"]) for batch in batches) + return [[] for _ in range(sum(observed_batch_sizes))] + + def map_entities(decoded, *args): + return decoded + + model._process_batches = process_batches + model.map_entities_to_text = map_entities + + model.inference(texts, entity_types, batch_size=2) + + assert observed_batch_sizes == expected_batch_sizes + + +def test_onnx_run_batch_rejects_unaligned_per_row_label_layout(): + """Low-level ONNX batching must fail loudly instead of ignoring gather metadata.""" + model = BiEncoderSpanGLiNER.__new__(BiEncoderSpanGLiNER) + torch.nn.Module.__init__(model) + model.onnx_model = True + batch = { + "input_ids": torch.ones(2, 1, dtype=torch.long), + "labels_gather_indices": torch.tensor([[0], [1]]), + } + + with pytest.raises(ValueError, match="Batched per-row labels are not supported"): + model.run_batch(batch) + + +def test_relex_inference_slices_per_row_relation_types_by_chunk(): + """Every relex DataLoader chunk should receive the relation schemas for its own rows.""" + model = UniEncoderSpanRelexGLiNER.__new__(UniEncoderSpanRelexGLiNER) + torch.nn.Module.__init__(model) + model._inference_packing_config = None + + texts = ["zero", "one", "two", "three"] + prepared = { + "input_x": [{"row": i} for i in range(4)], + "entity_types": [[f"entity_{i}"] for i in range(4)], + "relation_types": [[f"relation_{i}"] for i in range(4)], + "valid_texts": texts, + "valid_to_orig_idx": list(range(4)), + "start_token_map": [[] for _ in texts], + "end_token_map": [[] for _ in texts], + "word_input_spans": None, + "num_original": 4, + } + received_types = [] + + def prepare_batch(*args, **kwargs): + return prepared + + def create_collator(): + return object() + + def collate_batch(input_x, entity_types, collator, relation_types): + received_types.append((entity_types, relation_types)) + return {"tokens": [[str(item["row"])] for item in input_x]} + + def process_batches(data_loader, *args, **kwargs): + batches = list(data_loader) + batch_size = sum(len(batch["tokens"]) for batch in batches) + return ([[] for _ in range(batch_size)], [[] for _ in range(batch_size)]) + + def passthrough(decoded, *args): + return decoded + + model.prepare_batch = prepare_batch + model.create_collator = create_collator + model.collate_batch = collate_batch + model._process_batches = process_batches + model.map_entities_to_text = passthrough + model.map_relations_to_text = passthrough + + model.inference( + texts, + prepared["entity_types"], + relations=prepared["relation_types"], + batch_size=2, + ) + + assert received_types == [ + ([['entity_0'], ['entity_1']], [['relation_0'], ['relation_1']]), + ([['entity_2'], ['entity_3']], [['relation_2'], ['relation_3']]), + ]