gliner.modeling.base module

Base model classes for GLiNER neural network architectures.

This module provides abstract base classes and concrete implementations for various encoder-decoder architectures used in named entity recognition (NER) and relation extraction tasks.

Classes:

BaseModel: Abstract base class for all models. BaseUniEncoderModel: Base class for uni-encoder architectures. UniEncoderSpanModel: Span-based NER model with uni-encoder. UniEncoderTokenModel: Token-based NER model with uni-encoder. BaseBiEncoderModel: Base class for bi-encoder architectures. BiEncoderSpanModel: Span-based NER model with bi-encoder. BiEncoderTokenModel: Token-based NER model with bi-encoder. UniEncoderSpanDecoderModel: Span model with decoder for label generation. UniEncoderSpanRelexModel: Span model with relation extraction capabilities.

class gliner.modeling.base.BaseModel(config, from_pretrained=False, cache_dir=None)

Bases: ABC, Module

Abstract base class for all GLiNER models.

This class defines the common interface and shared functionality for all model architectures. It includes methods for padding/truncating sequences and computing losses with various masking strategies.

data_processor

Data processor for handling input preprocessing.

config

Model configuration object.

from_pretrained

Whether model was loaded from pretrained weights.

Type:

bool

cache_dir

Directory for caching pretrained models.

Type:

Optional[Path]

Initialize the base model.

Parameters:

• config (Any) – Configuration object containing model hyperparameters.

• from_pretrained (bool) – Whether to load from pretrained weights.

• cache_dir (str | Path | None) – Directory path for caching pretrained models.

data_processor = None

__init__(config, from_pretrained=False, cache_dir=None)

Initialize the base model.

Parameters:

• config (Any) – Configuration object containing model hyperparameters.

• from_pretrained (bool) – Whether to load from pretrained weights.

• cache_dir (str | Path | None) – Directory path for caching pretrained models.

abstract get_representations()

Get intermediate representations from the model.

Returns:

Tuple of tensors representing intermediate model outputs.

Return type:

Tuple[Tensor, …]

abstract forward(x)

Forward pass through the model.

Parameters:

x (Any) – Input data.

Returns:

Model outputs.

Return type:

Any

abstract loss(x)

Compute the total loss for the model.

Parameters:

x (Any) – Input data and labels.

Returns:

Scalar loss tensor.

Return type:

Tensor

class gliner.modeling.base.BaseUniEncoderModel(config, from_pretrained=False, cache_dir=None)

Bases: BaseModel

Base class for uni-encoder model architectures.

Uni-encoder models use a single encoder for both text and entity labels, embedding them in the same semantic space.

token_rep_layer

Token-level representation encoder.

Type:

Encoder

rnn

Optional LSTM layer for sequence modeling.

Type:

Optional[LstmSeq2SeqEncoder]

cross_fuser

Optional cross-attention fusion layer.

Type:

Optional[CrossFuser]

Initialize the uni-encoder model.

Parameters:

• config (Any) – Model configuration object.

• from_pretrained (bool) – Whether to load from pretrained weights.

• cache_dir (str | Path | None) – Directory for caching pretrained models.

__init__(config, from_pretrained=False, cache_dir=None)

Initialize the uni-encoder model.

Parameters:

• config (Any) – Model configuration object.

• from_pretrained (bool) – Whether to load from pretrained weights.

• cache_dir (str | Path | None) – Directory for caching pretrained models.

get_representations(input_ids=None, attention_mask=None, text_lengths=None, words_mask=None, **kwargs)

Get entity label and word representations from input.

Parameters:

• input_ids (FloatTensor | None) – Input token IDs of shape (B, L).

• attention_mask (LongTensor | None) – Attention mask of shape (B, L).

• text_lengths (Tensor | None) – Length of each text in batch.

• words_mask (LongTensor | None) – Word boundary mask.

• **kwargs (Any) – Additional arguments for the encoder.

Returns:

• prompts_embedding: Entity label embeddings of shape (B, C, D).

• prompts_embedding_mask: Mask for prompts of shape (B, C).

• words_embedding: Word embeddings of shape (B, W, D).

• mask: Mask for words of shape (B, W).

Return type:

Tuple containing

class gliner.modeling.base.UniEncoderSpanModel(config, from_pretrained=False, cache_dir=None)

Bases: BaseUniEncoderModel

Span-based NER model using uni-encoder architecture.

This model identifies entity spans by scoring all possible spans against entity type embeddings.

span_rep_layer

Layer for computing span representations.

Type:

SpanRepLayer

prompt_rep_layer

Projection layer for entity label embeddings.

Type:

nn.Module

Initialize the span-based uni-encoder model.

Parameters:

• config (Any) – Model configuration object.

• from_pretrained (bool) – Whether to load from pretrained weights.

• cache_dir (str | Path | None) – Directory for caching pretrained models.

__init__(config, from_pretrained=False, cache_dir=None)

Initialize the span-based uni-encoder model.

Parameters:

• config (Any) – Model configuration object.

• from_pretrained (bool) – Whether to load from pretrained weights.

• cache_dir (str | Path | None) – Directory for caching pretrained models.

forward(input_ids=None, attention_mask=None, words_embedding=None, mask=None, prompts_embedding=None, prompts_embedding_mask=None, words_mask=None, text_lengths=None, span_idx=None, span_mask=None, labels=None, **kwargs)

Forward pass through the span-based model.

Parameters:

• input_ids (FloatTensor | None) – Input token IDs of shape (B, L).

• attention_mask (LongTensor | None) – Attention mask of shape (B, L).

• words_embedding (FloatTensor | None) – Pre-computed word embeddings of shape (B, W, D).

• mask (LongTensor | None) – Mask for words of shape (B, W).

• prompts_embedding (FloatTensor | None) – Pre-computed entity label embeddings of shape (B, C, D).

• prompts_embedding_mask (LongTensor | None) – Mask for prompts of shape (B, C).

• words_mask (LongTensor | None) – Word boundary mask.

• text_lengths (Tensor | None) – Length of each text sequence.

• span_idx (LongTensor | None) – Span indices of shape (B, L*K, 2).

• span_mask (LongTensor | None) – Mask for valid spans of shape (B, L, K).

• labels (FloatTensor | None) – Ground truth labels of shape (B, L, K, C).

• **kwargs (Any) – Additional arguments.

Returns:

GLiNERBaseOutput containing logits, loss, and intermediate representations.

Return type:

GLiNERBaseOutput

loss(scores, labels, prompts_embedding_mask, mask_label, alpha=-1.0, gamma=0.0, prob_margin=0.0, label_smoothing=0.0, reduction='sum', negatives=1.0, masking='none', **kwargs)

Compute span classification loss.

Parameters:

• scores (Tensor) – Predicted scores of shape (B, L, K, C).

• labels (Tensor) – Ground truth labels of shape (B, L, K, C).

• prompts_embedding_mask (Tensor) – Mask for valid entity types of shape (B, C).

• mask_label (Tensor) – Mask for valid spans of shape (B, L, K).

• alpha (float) – Focal loss alpha parameter.

• gamma (float) – Focal loss gamma parameter.

• prob_margin (float) – Margin for probability adjustment.

• label_smoothing (float) – Label smoothing factor.

• reduction (str) – Loss reduction method (‘sum’ or ‘mean’).

• negatives (float) – Negative sampling probability.

• masking (str) – Masking strategy for negative sampling.

• **kwargs (Any) – Additional arguments.

Returns:

Scalar loss tensor.

Return type:

Tensor

class gliner.modeling.base.UniEncoderTokenModel(config, from_pretrained=False, cache_dir=None)

Bases: BaseUniEncoderModel

Token-based NER model using uni-encoder architecture.

This model classifies each word independently as entity type or non-entity.

scorer

Scoring layer for computing token-label compatibility.

Type:

Scorer

Initialize the token-based uni-encoder model.

Parameters:

• config (Any) – Model configuration object.

• from_pretrained (bool) – Whether to load from pretrained weights.

• cache_dir (str | Path | None) – Directory for caching pretrained models.

__init__(config, from_pretrained=False, cache_dir=None)

Initialize the token-based uni-encoder model.

Parameters:

• config (Any) – Model configuration object.

• from_pretrained (bool) – Whether to load from pretrained weights.

• cache_dir (str | Path | None) – Directory for caching pretrained models.

forward(input_ids=None, attention_mask=None, words_embedding=None, mask=None, prompts_embedding=None, prompts_embedding_mask=None, words_mask=None, text_lengths=None, labels=None, **kwargs)

Forward pass through the token-based model.

Parameters:

• input_ids (FloatTensor | None) – Input token IDs of shape (B, L).

• attention_mask (LongTensor | None) – Attention mask of shape (B, L).

• words_embedding (FloatTensor | None) – Pre-computed word embeddings of shape (B, W, D).

• mask (LongTensor | None) – Mask for words of shape (B, W).

• prompts_embedding (FloatTensor | None) – Pre-computed entity label embeddings of shape (B, C, D).

• prompts_embedding_mask (LongTensor | None) – Mask for prompts of shape (B, C).

• words_mask (LongTensor | None) – Word boundary mask.

• text_lengths (Tensor | None) – Length of each text sequence.

• labels (FloatTensor | None) – Ground truth labels of shape (B, W, C).

• **kwargs (Any) – Additional arguments.

Returns:

GLiNERBaseOutput containing logits, loss, and intermediate representations.

Return type:

GLiNERBaseOutput

loss(scores, labels, prompts_embedding_mask, mask, alpha=-1.0, gamma=0.0, prob_margin=0.0, label_smoothing=0.0, reduction='sum', negatives=1.0, **kwargs)

Compute token classification loss.

Parameters:

• scores (Tensor) – Predicted scores of shape (B, W, C).

• labels (Tensor) – Ground truth labels of shape (B, W, C).

• prompts_embedding_mask (Tensor) – Mask for valid entity types of shape (B, C).

• mask (Tensor) – Mask for valid tokens of shape (B, W).

• alpha (float) – Focal loss alpha parameter.

• gamma (float) – Focal loss gamma parameter.

• prob_margin (float) – Margin for probability adjustment.

• label_smoothing (float) – Label smoothing factor.

• reduction (str) – Loss reduction method (‘sum’ or ‘mean’).

• negatives (float) – Negative sampling probability.

• **kwargs (Any) – Additional arguments.

Returns:

Scalar loss tensor.

Return type:

Tensor

class gliner.modeling.base.BaseBiEncoderModel(config, from_pretrained=False, cache_dir=None)

Bases: BaseModel

Base class for bi-encoder model architectures.

Bi-encoder models use separate encoders for text and entity labels, allowing independent encoding before fusion.

token_rep_layer

Bi-encoder for text and labels.

Type:

BiEncoder

rnn

Optional LSTM layer.

Type:

Optional[LstmSeq2SeqEncoder]

cross_fuser

Optional cross-attention fusion layer.

Type:

Optional[CrossFuser]

Initialize the bi-encoder model.

Parameters:

• config (Any) – Model configuration object.

• from_pretrained (bool) – Whether to load from pretrained weights.

• cache_dir (str | Path | None) – Directory for caching pretrained models.

__init__(config, from_pretrained=False, cache_dir=None)

Initialize the bi-encoder model.

Parameters:

• config (Any) – Model configuration object.

• from_pretrained (bool) – Whether to load from pretrained weights.

• cache_dir (str | Path | None) – Directory for caching pretrained models.

features_enhancement(text_embeds, labels_embeds, text_mask=None, labels_mask=None)

Enhance features using cross-attention fusion.

Parameters:

• text_embeds (Tensor) – Text embeddings of shape (B, W, D).

• labels_embeds (Tensor) – Label embeddings of shape (B, C, D).

• text_mask (Tensor | None) – Mask for text of shape (B, W).

• labels_mask (Tensor | None) – Mask for labels of shape (B, C).

Returns:

• Enhanced text embeddings of shape (B, W, D).

• Enhanced label embeddings of shape (B, C, D).

Return type:

Tuple containing

get_representations(input_ids=None, attention_mask=None, labels_embeds=None, labels_input_ids=None, labels_attention_mask=None, text_lengths=None, words_mask=None, **kwargs)

Get entity label and word representations using bi-encoder.

Parameters:

• input_ids (FloatTensor | None) – Input token IDs of shape (B, L).

• attention_mask (LongTensor | None) – Attention mask of shape (B, L).

• labels_embeds (FloatTensor | None) – Pre-computed label embeddings of shape (C, D).

• labels_input_ids (FloatTensor | None) – Label token IDs for encoding.

• labels_attention_mask (LongTensor | None) – Attention mask for labels.

• text_lengths (Tensor | None) – Length of each text in batch.

• words_mask (LongTensor | None) – Word boundary mask.

• **kwargs (Any) – Additional arguments for the encoder.

Returns:

• labels_embeds: Label embeddings of shape (B, C, D).

• labels_mask: Mask for labels of shape (B, C).

• words_embedding: Word embeddings of shape (B, W, D).

• mask: Mask for words of shape (B, W).

Return type:

Tuple containing

class gliner.modeling.base.BiEncoderSpanModel(config, from_pretrained=False, cache_dir=None)

Bases: BaseBiEncoderModel

Span-based NER model using bi-encoder architecture.

span_rep_layer

Layer for computing span representations.

Type:

SpanRepLayer

prompt_rep_layer

Projection layer for entity label embeddings.

Type:

nn.Module

Initialize the span-based bi-encoder model.

Parameters:

• config (Any) – Model configuration object.

• from_pretrained (bool) – Whether to load from pretrained weights.

• cache_dir (str | Path | None) – Directory for caching pretrained models.

__init__(config, from_pretrained=False, cache_dir=None)

Initialize the span-based bi-encoder model.

Parameters:

• config (Any) – Model configuration object.

• from_pretrained (bool) – Whether to load from pretrained weights.

• cache_dir (str | Path | None) – Directory for caching pretrained models.

forward(input_ids=None, attention_mask=None, labels_embeds=None, labels_input_ids=None, labels_attention_mask=None, words_embedding=None, mask=None, prompts_embedding=None, prompts_embedding_mask=None, words_mask=None, text_lengths=None, span_idx=None, span_mask=None, labels=None, **kwargs)

Forward pass through the bi-encoder span model.

Parameters:

• input_ids (FloatTensor | None) – Input token IDs of shape (B, L).

• attention_mask (LongTensor | None) – Attention mask of shape (B, L).

• labels_embeds (FloatTensor | None) – Pre-computed label embeddings of shape (C, D).

• labels_input_ids (FloatTensor | None) – Label token IDs for encoding.

• labels_attention_mask (LongTensor | None) – Attention mask for labels.

• words_embedding (FloatTensor | None) – Pre-computed word embeddings.

• mask (LongTensor | None) – Mask for words.

• prompts_embedding (FloatTensor | None) – Pre-computed entity label embeddings.

• prompts_embedding_mask (LongTensor | None) – Mask for prompts.

• words_mask (LongTensor | None) – Word boundary mask.

• text_lengths (Tensor | None) – Length of each text sequence.

• span_idx (LongTensor | None) – Span indices of shape (B, L*K, 2).

• span_mask (LongTensor | None) – Mask for valid spans of shape (B, L, K).

• labels (FloatTensor | None) – Ground truth labels of shape (B, L, K, C).

• **kwargs (Any) – Additional arguments.

Returns:

GLiNERBaseOutput containing logits, loss, and intermediate representations.

Return type:

GLiNERBaseOutput

loss(scores, labels, prompts_embedding_mask, mask_label, alpha=-1.0, gamma=0.0, label_smoothing=0.0, reduction='sum', negatives=1.0, masking='none', **kwargs)

Compute span classification loss for bi-encoder.

Parameters:

• scores (Tensor) – Predicted scores of shape (B, L, K, C).

• labels (Tensor) – Ground truth labels of shape (B, L, K, C).

• prompts_embedding_mask (Tensor) – Mask for valid entity types of shape (B, C).

• mask_label (Tensor) – Mask for valid spans of shape (B, L, K).

• alpha (float) – Focal loss alpha parameter.

• gamma (float) – Focal loss gamma parameter.

• label_smoothing (float) – Label smoothing factor.

• reduction (str) – Loss reduction method (‘sum’ or ‘mean’).

• negatives (float) – Negative sampling probability.

• masking (str) – Masking strategy for negative sampling.

• **kwargs (Any) – Additional arguments.

Returns:

Scalar loss tensor.

Return type:

Tensor

class gliner.modeling.base.BiEncoderTokenModel(config, from_pretrained=False, cache_dir=None)

Bases: BaseBiEncoderModel

Token-based NER model using bi-encoder architecture.

scorer

Scoring layer for computing token-label compatibility.

Type:

Scorer

Initialize the token-based bi-encoder model.

Parameters:

• config (Any) – Model configuration object.

• from_pretrained (bool) – Whether to load from pretrained weights.

• cache_dir (str | Path | None) – Directory for caching pretrained models.

__init__(config, from_pretrained=False, cache_dir=None)

Initialize the token-based bi-encoder model.

Parameters:

• config (Any) – Model configuration object.

• from_pretrained (bool) – Whether to load from pretrained weights.

• cache_dir (str | Path | None) – Directory for caching pretrained models.

forward(input_ids=None, attention_mask=None, labels_embeds=None, labels_input_ids=None, labels_attention_mask=None, words_embedding=None, mask=None, prompts_embedding=None, prompts_embedding_mask=None, words_mask=None, text_lengths=None, labels=None, **kwargs)

Forward pass through the bi-encoder token model.

Parameters:

• input_ids (FloatTensor | None) – Input token IDs of shape (B, L).

• attention_mask (LongTensor | None) – Attention mask of shape (B, L).

• labels_embeds (FloatTensor | None) – Pre-computed label embeddings of shape (C, D).

• labels_input_ids (FloatTensor | None) – Label token IDs for encoding.

• labels_attention_mask (LongTensor | None) – Attention mask for labels.

• words_embedding (FloatTensor | None) – Pre-computed word embeddings.

• mask (LongTensor | None) – Mask for words.

• prompts_embedding (FloatTensor | None) – Pre-computed entity label embeddings.

• prompts_embedding_mask (LongTensor | None) – Mask for prompts.

• words_mask (LongTensor | None) – Word boundary mask.

• text_lengths (Tensor | None) – Length of each text sequence.

• labels (FloatTensor | None) – Ground truth labels of shape (B, W, C).

• **kwargs (Any) – Additional arguments.

Returns:

GLiNERBaseOutput containing logits, loss, and intermediate representations.

Return type:

GLiNERBaseOutput

loss(scores, labels, prompts_embedding_mask, mask, alpha=-1.0, gamma=0.0, prob_margin=0.0, label_smoothing=0.0, reduction='sum', negatives=1.0, **kwargs)

Compute token classification loss for bi-encoder.

Parameters:

• scores (Tensor) – Predicted scores of shape (B, W, C).

• labels (Tensor) – Ground truth labels of shape (B, W, C).

• prompts_embedding_mask (Tensor) – Mask for valid entity types of shape (B, C).

• mask (Tensor) – Mask for valid tokens of shape (B, W).

• alpha (float) – Focal loss alpha parameter.

• gamma (float) – Focal loss gamma parameter.

• prob_margin (float) – Margin for probability adjustment.

• label_smoothing (float) – Label smoothing factor.

• reduction (str) – Loss reduction method (‘sum’ or ‘mean’).

• negatives (float) – Negative sampling probability.

• **kwargs (Any) – Additional arguments.

Returns:

Scalar loss tensor.

Return type:

Tensor

class gliner.modeling.base.UniEncoderSpanDecoderModel(config, from_pretrained=False, cache_dir=None)

Bases: UniEncoderSpanModel

Span-based model with decoder for generating entity type labels.

This model extends the span-based approach by adding a decoder that can generate entity type labels as sequences, enabling more flexible entity typing.

decoder

Decoder module for label generation.

Type:

Decoder

_enc2dec_proj

Projection layer if encoder and decoder dimensions differ.

Type:

Optional[nn.Module]

Initialize the span-decoder model.

Parameters:

• config (Any) – Model configuration object.

• from_pretrained (bool) – Whether to load from pretrained weights.

• cache_dir (str | Path | None) – Directory for caching pretrained models.

__init__(config, from_pretrained=False, cache_dir=None)

Initialize the span-decoder model.

Parameters:

• config (Any) – Model configuration object.

• from_pretrained (bool) – Whether to load from pretrained weights.

• cache_dir (str | Path | None) – Directory for caching pretrained models.

select_decoder_embedding(representations, rep_mask)

Select and pack valid representations for decoder input.

Keeps only representations where mask == 1 and optionally projects them to the decoder hidden size.

Parameters:

• representations (FloatTensor) – Tensor of shape (B, N, D) containing embeddings.

• rep_mask (LongTensor) – Tensor of shape (B, N) where 1 indicates valid position.

Returns:

• target_rep: FloatTensor of shape (B, M, D) with kept representations.

• target_mask: LongTensor of shape (B, M) with 1 for valid, 0 for pad.

• sel_idx: LongTensor of shape (B, M) with original column indices (-1 for padding positions).

Return type:

Tuple containing

get_raw_decoder_inputs(representations, rep_mask)

Extract valid span tokens for decoder input.

Parameters:

• representations (Tensor) – Span representations of shape (B, S, T, D).

• rep_mask (Tensor) – Mask of shape (B, S, T).

Returns:

• span_tokens: Valid span tokens of shape (M, T, D).

• span_tokens_mask: Mask for span tokens of shape (M, T).

Return type:

Tuple containing

decode_labels(decoder_embedding=None, decoder_embedding_mask=None, decoder_labels_ids=None, decoder_labels_mask=None, decoder_labels=None, **kwargs)

Decode labels using the decoder in teacher forcing mode.

Parameters:

• decoder_embedding (FloatTensor | None) – Span token embeddings of shape (B, N, T, D).

• decoder_embedding_mask (LongTensor | None) – Mask for span tokens of shape (B, N, T).

• decoder_labels_ids (FloatTensor | None) – Label token IDs of shape (M, L).

• decoder_labels_mask (LongTensor | None) – Mask for labels of shape (M, L).

• decoder_labels (FloatTensor | None) – Ground truth labels for loss computation of shape (M, L).

• **kwargs (Any) – Additional arguments.

Returns:

• loss: Cross-entropy loss scalar.

• decoder_outputs: Decoder logits of shape (M, S+L-1, V).

Return type:

Tuple containing

generate_labels(decoder_embedding=None, decoder_embedding_mask=None, max_new_tokens=32, eos_token_id=None, pad_token_id=None, temperature=1.0, do_sample=False, num_return_sequences=1, labels_trie=None, **kwargs)

Generate entity type labels from decoder embeddings.

Parameters:

• decoder_embedding (FloatTensor | None) – Span token embeddings of shape (B, N, D).

• decoder_embedding_mask (LongTensor | None) – Mask for span tokens of shape (B, N).

• max_new_tokens (int) – Maximum number of tokens to generate.

• eos_token_id (int | None) – End-of-sequence token ID.

• pad_token_id (int | None) – Padding token ID.

• temperature (float) – Sampling temperature.

• do_sample (bool) – Whether to use sampling instead of greedy decoding.

• num_return_sequences (int) – Number of sequences to generate per input.

• labels_trie (Any | None) – Optional trie for constrained decoding.

• **kwargs (Any) – Additional generation arguments.

Returns:

Generated token IDs of shape (M, L).

Return type:

Tensor

select_span_decoder_embedding(prompts_embedding, prompts_embedding_mask, span_rep, span_scores, span_mask, decoder_text_embeds=None, decoder_words_mask=None, span_labels=None, threshold=0.5, top_k=None, decoder_input_ids=None, decoder_labels_ids=None)

Select span embeddings for decoder input based on predictions or labels.

This method selects which spans to provide to the decoder, either based on ground truth labels (during training) or model predictions (during inference). It can operate in two modes: 1. “prompt” mode: Uses entity type embeddings as decoder input. 2. “span” mode: Uses contextualized tokens within each span as decoder input.

Parameters:

• prompts_embedding (Tensor) – Entity type embeddings of shape (B, C, D).

• prompts_embedding_mask (Tensor) – Mask for prompts of shape (B, C).

• span_rep (Tensor) – Span representations of shape (B, L, K, D).

• span_scores (Tensor) – Span classification scores of shape (B, L, K, C).

• span_mask (Tensor) – Mask for valid spans of shape (B, L, K).

• decoder_text_embeds (Tensor | None) – Text embeddings for span mode of shape (B, T, D).

• decoder_words_mask (Tensor | None) – Word position mask of shape (B, T).

• span_labels (Tensor | None) – Ground truth labels of shape (B, L, K, C).

• threshold (float) – Confidence threshold for selecting spans.

• top_k (int | None) – Optional limit on number of spans to select.

• decoder_input_ids (Tensor | None) – Debugging parameter for input IDs.

• decoder_labels_ids (Tensor | None) – Debugging parameter for label IDs.

Returns:

• span_rep_kept: Selected span embeddings for decoder of shape (B, S, T, D) or None if no valid spans.

• span_msk: Mask for selected spans of shape (B, S, T) or None.

• span_sel_idx: Original indices of selected spans of shape (B, S) or None.

Return type:

Tuple containing

forward(input_ids=None, attention_mask=None, decoder_input_ids=None, decoder_attention_mask=None, decoder_labels_ids=None, decoder_labels_mask=None, decoder_words_mask=None, words_embedding=None, mask=None, prompts_embedding=None, prompts_embedding_mask=None, words_mask=None, text_lengths=None, span_idx=None, span_mask=None, labels=None, decoder_labels=None, threshold=0.5, **kwargs)

Forward pass through the span-decoder model.

Parameters:

• input_ids (FloatTensor | None) – Input token IDs of shape (B, L).

• attention_mask (LongTensor | None) – Attention mask of shape (B, L).

• decoder_input_ids (FloatTensor | None) – Decoder input IDs for span mode.

• decoder_attention_mask (LongTensor | None) – Decoder attention mask.

• decoder_labels_ids (FloatTensor | None) – Label token IDs for decoding of shape (M, L).

• decoder_labels_mask (LongTensor | None) – Mask for decoder labels of shape (M, L).

• decoder_words_mask (LongTensor | None) – Word position mask for span mode.

• words_embedding (FloatTensor | None) – Pre-computed word embeddings.

• mask (LongTensor | None) – Mask for words.

• prompts_embedding (FloatTensor | None) – Pre-computed entity type embeddings.

• prompts_embedding_mask (LongTensor | None) – Mask for prompts.

• words_mask (LongTensor | None) – Word boundary mask.

• text_lengths (Tensor | None) – Length of each text sequence.

• span_idx (LongTensor | None) – Span indices of shape (B, L*K, 2).

• span_mask (LongTensor | None) – Mask for valid spans of shape (B, L, K).

• labels (FloatTensor | None) – Ground truth span labels of shape (B, L, K, C).

• decoder_labels (FloatTensor | None) – Ground truth decoder labels of shape (M, L).

• threshold (float | None) – Confidence threshold for span selection.

• **kwargs (Any) – Additional arguments.

Returns:

GLiNERDecoderOutput containing logits, losses, and decoder information.

Return type:

GLiNERDecoderOutput

loss(scores, labels, prompts_embedding_mask, mask_label, alpha=-1.0, gamma=0.0, label_smoothing=0.0, reduction='sum', negatives=1.0, masking='none', decoder_loss=None, **kwargs)

Compute combined loss for span classification and decoder.

Parameters:

• scores (Tensor) – Predicted span scores of shape (B, L, K, C).

• labels (Tensor) – Ground truth span labels of shape (B, L, K, C).

• prompts_embedding_mask (Tensor) – Mask for valid entity types of shape (B, C).

• mask_label (Tensor) – Mask for valid spans of shape (B, L, K).

• alpha (float) – Focal loss alpha parameter.

• gamma (float) – Focal loss gamma parameter.

• label_smoothing (float) – Label smoothing factor.

• reduction (str) – Loss reduction method (‘sum’ or ‘mean’).

• negatives (float) – Negative sampling probability.

• masking (str) – Masking strategy for negative sampling.

• decoder_loss (Tensor | None) – Optional decoder loss to combine with span loss.

• **kwargs (Any) – Additional arguments.

Returns:

Scalar combined loss tensor.

Return type:

Tensor

class gliner.modeling.base.UniEncoderSpanRelexModel(config, from_pretrained=False, cache_dir=None)

Bases: UniEncoderSpanModel

Span-based NER model with relation extraction capabilities.

This model extends span-based NER to also extract relations between identified entities, predicting both entity types and relation types in a joint model.

relations_rep_layer

Layer for computing pairwise entity relations (adjacency matrix).

Type:

Optional[RelationsRepLayer]

triples_score_layer

Layer for scoring (head, relation, tail) triples.

Type:

Optional[TriplesScoreLayer]

pair_rep_layer

Alternative layer for relation scoring via concatenation.

Type:

Optional[nn.Module]

Initialize the span-based relation extraction model.

Parameters:

• config (Any) – Model configuration object.

• from_pretrained (bool) – Whether to load from pretrained weights.

• cache_dir (str | Path | None) – Directory for caching pretrained models.

__init__(config, from_pretrained=False, cache_dir=None)

Initialize the span-based relation extraction model.

Parameters:

• config (Any) – Model configuration object.

• from_pretrained (bool) – Whether to load from pretrained weights.

• cache_dir (str | Path | None) – Directory for caching pretrained models.

select_span_target_embedding(span_rep, span_scores, span_mask, span_labels=None, threshold=0.5, top_k=None)

Select entity spans for relation extraction.

Filters spans based on entity classification scores or ground truth labels, keeping only high-confidence or positive entity spans for relation modeling.

Parameters:

• span_rep (FloatTensor) – Span representations of shape (B, L, K, D).

• span_scores (FloatTensor) – Span classification scores of shape (B, L, K, C).

• span_mask (LongTensor) – Mask for valid spans of shape (B, L, K).

• span_labels (FloatTensor | None) – Optional ground truth labels of shape (B, L, K, C).

• threshold (float) – Confidence threshold for selecting spans.

• top_k (int | None) – Optional limit on number of spans to select.

Returns:

• target_rep: Selected span representations of shape (B, E, D).

• target_mask: Mask for selected spans of shape (B, E).

Return type:

Tuple containing

select_target_embedding(representations=None, rep_mask=None)

Pack valid representations by removing masked positions.

Parameters:

• representations (FloatTensor | None) – Tensor of shape (B, N, D).

• rep_mask (LongTensor | None) – Mask of shape (B, N) where 1 indicates valid position.

Returns:

• target_rep: Packed representations of shape (B, M, D).

• target_mask: Packed mask of shape (B, M).

Return type:

Tuple containing

forward(input_ids=None, attention_mask=None, words_embedding=None, mask=None, prompts_embedding=None, prompts_embedding_mask=None, words_mask=None, text_lengths=None, span_idx=None, span_mask=None, labels=None, adj_matrix=None, rel_matrix=None, threshold=0.5, adjacency_threshold=0.5, **kwargs)

Forward pass through the relation extraction model.

Parameters:

• input_ids (FloatTensor | None) – Input token IDs of shape (B, L).

• attention_mask (LongTensor | None) – Attention mask of shape (B, L).

• words_embedding (FloatTensor | None) – Pre-computed word embeddings.

• mask (LongTensor | None) – Mask for words.

• prompts_embedding (FloatTensor | None) – Pre-computed entity type embeddings.

• prompts_embedding_mask (LongTensor | None) – Mask for entity types.

• words_mask (LongTensor | None) – Word boundary mask.

• text_lengths (Tensor | None) – Length of each text sequence.

• span_idx (LongTensor | None) – Span indices of shape (B, L*K, 2).

• span_mask (LongTensor | None) – Mask for valid spans of shape (B, L, K).

• labels (FloatTensor | None) – Ground truth entity labels of shape (B, L, K, C).

• adj_matrix (FloatTensor | None) – Ground truth adjacency matrix of shape (B, E, E).

• rel_matrix (FloatTensor | None) – Ground truth relation labels of shape (B, N, C_rel).

• threshold (float | None) – Confidence threshold for entity selection.

• adjacency_threshold (float | None) – Threshold for relation adjacency.

• **kwargs (Any) – Additional arguments.

Returns:

GLiNERRelexOutput containing entity and relation predictions.

Return type:

GLiNERRelexOutput

adj_loss(logits, labels, adj_mask, alpha=-1.0, gamma=0.0, prob_margin=0.0, label_smoothing=0.0, reduction='sum', masking='span', negatives=1.0, **kwargs)

Compute adjacency matrix loss for entity pair relations.

Parameters:

• logits (Tensor) – Predicted adjacency scores of shape (B, E, E).

• labels (Tensor) – Ground truth adjacency of shape (B, E, E).

• adj_mask (Tensor) – Mask for valid entity pairs of shape (B, E, E).

• alpha (float) – Focal loss alpha parameter.

• gamma (float) – Focal loss gamma parameter.

• prob_margin (float) – Margin for probability adjustment.

• label_smoothing (float) – Label smoothing factor.

• reduction (str) – Loss reduction method (‘sum’ or ‘mean’).

• masking (str) – Masking strategy for negative sampling.

• negatives (float) – Negative sampling probability.

• **kwargs (Any) – Additional arguments.

Returns:

Scalar adjacency loss tensor.

Return type:

Tensor

rel_loss(logits, labels, pair_mask, class_mask, alpha=-1.0, gamma=0.0, prob_margin=0.0, label_smoothing=0.0, reduction='sum', negatives=1.0, masking='span', **kwargs)

Compute relation classification loss for selected entity pairs.

Parameters:

• logits (Tensor) – Predicted relation scores of shape (B, N, C).

• labels (Tensor) – Ground truth relation labels of shape (B, N, C).

• pair_mask (Tensor) – Mask for valid pairs of shape (B, N, C).

• class_mask (Tensor) – Mask for valid relation classes of shape (B, N, C).

• alpha (float) – Focal loss alpha parameter.

• gamma (float) – Focal loss gamma parameter.

• prob_margin (float) – Margin for probability adjustment.

• label_smoothing (float) – Label smoothing factor.

• reduction (str) – Loss reduction method (‘sum’ or ‘mean’).

• negatives (float) – Negative sampling probability.

• masking (str) – Masking strategy for negative sampling.

• **kwargs (Any) – Additional arguments.

Returns:

Scalar relation classification loss tensor.

Return type:

Tensor