lisbet.modeling.heads.embedding

Embedding head for extracting behavior embeddings.

Classes

EmbeddingHead(output_token_idx)

Embedding head for extracting behavior embeddings.

class lisbet.modeling.heads.embedding.EmbeddingHead(output_token_idx)

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)

Initialize internal Module state, shared by both nn.Module and ScriptModule.

forward(x)

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()

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)

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