gliner.modeling.utils module

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

Extract word-level embeddings from subword token embeddings.

Maps subword token embeddings back to word-level embeddings using a word mask that indicates which subword token corresponds to which word. Only the first subword token of each word is typically used for the word representation.

This is essential for span-based NER where predictions are made at the word level but the transformer operates on subword tokens.

Parameters:

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

• words_mask (Tensor) – Mask mapping subword positions to word indices. Non-zero values indicate the word index (1-indexed). Zero values are special tokens or continuation subwords to ignore. 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.

• max_text_length (int) – Maximum number of words across all examples in batch.

• embed_dim (int) – Embedding dimension size.

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

Returns:

• words_embedding: Word-level embeddings extracted from token embeddings. Shape: (batch_size, max_text_length, embed_dim)

• mask: Boolean mask indicating valid word positions (True) vs padding (False). Shape: (batch_size, max_text_length)

Return type:

Tuple containing

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, **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.

• **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