gliner.modeling.utils module

gliner.modeling.utils.extract_first_word_embeddings(token_embeds, words_mask, batch_size, max_text_length, embed_dim)

Select the first marked subtoken representation for each word.

gliner.modeling.utils.extract_last_word_embeddings(token_embeds, words_mask, batch_size, max_text_length, embed_dim)

Select the last marked subtoken representation for each word.

gliner.modeling.utils.extract_mean_word_embeddings(token_embeds, words_mask, attention_mask, batch_size, max_text_length, embed_dim)

Mean-pool all attended subtoken representations for each word.

gliner.modeling.utils.extract_max_word_embeddings(token_embeds, words_mask, attention_mask, batch_size, max_text_length, embed_dim)

Element-wise max-pool all attended subtoken representations for each word.

gliner.modeling.utils.extract_word_embeddings(token_embeds, words_mask, attention_mask, batch_size, max_text_length, embed_dim, text_lengths, subtoken_pooling='first')

Dispatch to the configured subtoken pooling implementation.

first and last expect one marked subtoken per word. mean and max expect every subtoken to carry its 1-based word index.

gliner.modeling.utils.extract_prompt_features(class_token_index, token_embeds, input_ids, attention_mask, batch_size, embed_dim, embed_ent_token=True)

Extract prompt/entity type embeddings from special class tokens.

Extracts embeddings for entity types or other prompt elements that are marked with special class tokens (e.g., [ENT] tokens). These embeddings represent the entity types that the model should extract.

In prompt-based NER, the input is typically:

[ENT] Person [ENT] Organization [SEP] John works at Google

This function extracts the embeddings corresponding to the [ENT] tokens (or the tokens immediately after them if embed_ent_token=False).

Parameters:

• class_token_index (int) – Token ID of the special class token to extract (e.g., token ID for [ENT]).

• token_embeds (Tensor) – Token embeddings from transformer. Shape: (batch_size, seq_len, embed_dim)

• input_ids (Tensor) – Token IDs from tokenizer. Shape: (batch_size, seq_len)

• attention_mask (Tensor) – Standard attention mask from tokenizer. Shape: (batch_size, seq_len)

• batch_size (int) – Size of the batch.

• embed_dim (int) – Embedding dimension size.

• embed_ent_token (bool) – If True, use the [ENT] token embedding itself. If False, use the embedding of the token immediately after [ENT] (i.e., the entity type name token). Default: True.

Returns:

• prompts_embedding: Embeddings for each prompt/entity type. Shape: (batch_size, max_num_types, embed_dim) where max_num_types is the maximum number of entity types across examples in the batch.

• prompts_embedding_mask: Mask indicating valid prompt positions (True) vs padding (False). Shape: (batch_size, max_num_types)

Return type:

Tuple containing

gliner.modeling.utils.extract_prompt_features_and_word_embeddings(class_token_index, token_embeds, input_ids, attention_mask, text_lengths, words_mask, embed_ent_token=True, subtoken_pooling='first', **kwargs)

Extract both prompt embeddings and word embeddings in one call.

Convenience function that combines extract_prompt_features and extract_word_embeddings to get both prompt/entity type embeddings and word-level text embeddings from a single set of token embeddings.

This is the typical use case for prompt-based NER where you need both: 1. Entity type embeddings (from prompt tokens like [ENT]) 2. Word-level text embeddings (from the actual text tokens)

Parameters:

• class_token_index (int) – Token ID of the special class token (e.g., [ENT]).

• token_embeds (Tensor) – Token embeddings from transformer. Shape: (batch_size, seq_len, embed_dim)

• input_ids (Tensor) – Token IDs from tokenizer. Shape: (batch_size, seq_len)

• attention_mask (Tensor) – Standard attention mask from tokenizer. Shape: (batch_size, seq_len)

• text_lengths (Tensor) – Number of words in each example. Shape: (batch_size, 1) or (batch_size,)

• words_mask (Tensor) – Mask mapping subword positions to word indices. Shape: (batch_size, seq_len)

• embed_ent_token (bool) – If True, use [ENT] token embedding. If False, use the token after [ENT] (the entity type name). Default: True.

• subtoken_pooling (str) – Reduction applied to subtokens belonging to the same word. One of first, last, mean, or max.

• **kwargs – Additional keyword arguments passed to extract_prompt_features.

Returns:

• prompts_embedding: Entity type embeddings. Shape: (batch_size, max_num_types, embed_dim)

• prompts_embedding_mask: Mask for valid entity type positions. Shape: (batch_size, max_num_types)

• words_embedding: Word-level text embeddings. Shape: (batch_size, max_text_length, embed_dim)

• mask: Mask for valid word positions. Shape: (batch_size, max_text_length)

Return type:

Tuple containing

gliner.modeling.utils.build_entity_pairs(adj, span_rep, threshold=0.5)

Build entity pairs for relation extraction based on adjacency scores.

Extracts entity pairs (head, tail) where the adjacency score exceeds a threshold, and retrieves their corresponding embeddings. This is used in relation extraction to select which entity pairs should be classified for relation types.

The function considers ALL directed pairs (i,j) where i≠j, not just the upper triangle, since relation direction matters (e.g., “founded” vs “founded_by” have opposite directions).

Parameters:

• adj (Tensor) – Adjacency matrix with scores or probabilities for entity pairs. Shape: (batch_size, num_entities, num_entities) The diagonal (self-pairs) is ignored. Values > threshold indicate potential relations.

• span_rep (Tensor) – Entity/span embeddings for each entity in the batch. Shape: (batch_size, num_entities, embed_dim)

• threshold (float) – Minimum adjacency score to consider a pair as a potential relation. Pairs with adj[i,j] > threshold are kept. Default: 0.5.

Returns:

• pair_idx: Indices of (head, tail) entity pairs. Shape: (batch_size, max_pairs, 2) Values are entity indices, or -1 for padding positions.

• pair_mask: Boolean mask indicating valid pairs (True) vs padding (False). Shape: (batch_size, max_pairs)

• head_rep: Embeddings of head entities for each pair. Shape: (batch_size, max_pairs, embed_dim)

• tail_rep: Embeddings of tail entities for each pair. Shape: (batch_size, max_pairs, embed_dim)

Return type:

Tuple containing

gliner.modeling.utils.build_all_entity_pairs(span_rep, span_mask)

Build all possible entity pairs for single-step relation extraction.

Generates all directed pairs (i, j) where i != j for valid entities (those with span_mask == 1), without any adjacency filtering.

Parameters:

• span_rep (Tensor) – Entity/span embeddings. Shape: (batch_size, num_entities, embed_dim)

• span_mask (Tensor) – Mask for valid entities. Shape: (batch_size, num_entities)

Returns:

• pair_idx: Indices of (head, tail) entity pairs. Shape: (B, max_pairs, 2)

• pair_mask: Boolean mask for valid pairs. Shape: (B, max_pairs)

• head_rep: Head entity embeddings. Shape: (B, max_pairs, embed_dim)

• tail_rep: Tail entity embeddings. Shape: (B, max_pairs, embed_dim)

Return type:

Tuple containing

gliner.modeling.utils.extract_spans_from_tokens(scores, labels=None, threshold=0.5)

Extract entity spans from BIO-style token predictions.

Parameters:

• scores (Tensor) – (B, W, C, 3) - logits for [start, end, inside]

• labels (Tensor | None) – Optional (B, W, C, 3) - ground truth labels

• threshold (float) – Confidence threshold (used when labels is None)

Returns:

(B, N, 2) - [start, end] indices, padded span_mask: (B, N) - validity mask

Return type:

span_idx