lisbet.modeling.heads#

class lisbet.modeling.heads.FrameClassificationHead(output_token_idx, input_dim, num_classes, hidden_dim=None)[source]#

Frame-level classification head.

This head selects a specific token from the sequence (typically the last one) and applies a classification layer to predict frame-level labels.

Parameters:

output_token_idx (int) – Index of the token to use for classification (e.g., -1 for last token).
input_dim (int) – Dimension of the input embeddings (formerly emb_dim).
num_classes (int) – Number of output classes (formerly out_dim).
hidden_dim (int | None) – Dimension of the hidden layer. If None, uses a single linear layer. If provided, uses an MLP with the specified hidden dimension.

output_token_idx#

Index of the token used for classification.

Type:: int

logits#

Classification layer (either Linear or MLP).

Type:: nn.Module

__init__(output_token_idx, input_dim, num_classes, hidden_dim=None)[source]#: Initialize internal Module state, shared by both nn.Module and ScriptModule.

forward(x)[source]#

Forward pass through the frame classification head.

Parameters:: x (Tensor) – Input tensor of shape (batch_size, sequence_length, input_dim).
Returns:: Classification logits of shape (batch_size, num_classes).
Return type:: Tensor

get_config()[source]#

Get the configuration dictionary for this head.

Returns:: Configuration dictionary containing all parameters needed to recreate this head instance.
Return type:: dict[str, Any]

classmethod from_config(config)[source]#

Create a FrameClassificationHead instance from a configuration dictionary.

Parameters:: config (dict[str, Any]) – Configuration dictionary containing all parameters needed to create the head instance.
Returns:: New FrameClassificationHead instance created from the configuration.
Return type:: FrameClassificationHead

class lisbet.modeling.heads.WindowClassificationHead(input_dim, num_classes, hidden_dim=None)[source]#

Window-level classification head.

This head performs global max pooling over the sequence dimension and applies a classification layer to predict window-level labels.

Parameters:

input_dim (int) – Dimension of the input embeddings (formerly emb_dim).
num_classes (int) – Number of output classes (formerly out_dim).
hidden_dim (int | None) – Dimension of the hidden layer. If None, uses a single linear layer. If provided, uses an MLP with the specified hidden dimension.

logits#

Classification layer (either Linear or MLP).

Type:: nn.Module

__init__(input_dim, num_classes, hidden_dim=None)[source]#: Initialize internal Module state, shared by both nn.Module and ScriptModule.

forward(x)[source]#

Forward pass through the window classification head.

Parameters:: x (Tensor) – Input tensor of shape (batch_size, sequence_length, input_dim).
Returns:: Classification logits of shape (batch_size, num_classes).
Return type:: Tensor

get_config()[source]#

Get the configuration dictionary for this head.

Returns:: Configuration dictionary containing all parameters needed to recreate this head instance.
Return type:: dict[str, Any]

classmethod from_config(config)[source]#

Create a WindowClassificationHead instance from a configuration dictionary.

Parameters:: config (dict[str, Any]) – Configuration dictionary containing all parameters needed to create the head instance.
Returns:: New WindowClassificationHead instance created from the configuration.
Return type:: WindowClassificationHead

class lisbet.modeling.heads.EmbeddingHead(output_token_idx)[source]#

Embedding head for extracting behavior embeddings.

This head selects a specific token from the sequence (typically the last one) and returns it as the behavior embedding without any additional transformation.

Parameters:: output_token_idx (int) – Index of the token to use for embedding extraction (e.g., -1 for last token).

output_token_idx#

Index of the token used for embedding extraction.

Type:: int

__init__(output_token_idx)[source]#: Initialize internal Module state, shared by both nn.Module and ScriptModule.

forward(x)[source]#

Forward pass through the embedding head.

Parameters:: x (Tensor) – Input tensor of shape (batch_size, sequence_length, embedding_dim).
Returns:: Embedding tensor of shape (batch_size, embedding_dim).
Return type:: Tensor

get_config()[source]#

Get the configuration dictionary for this head.

Returns:: Configuration dictionary containing all parameters needed to recreate this head instance.
Return type:: dict[str, Any]

classmethod from_config(config)[source]#

Create an EmbeddingHead instance from a configuration dictionary.

Parameters:: config (dict[str, Any]) – Configuration dictionary containing all parameters needed to create the head instance.
Returns:: New EmbeddingHead instance created from the configuration.
Return type:: EmbeddingHead

Modules

`classification`	Classification heads for frame and window classification tasks.
`embedding`	Embedding head for extracting behavior embeddings.