lrnnx.architectures.embedding module

Embedding modules for sequence models.

class PositionEmbedding[source]

Bases: Module

Learned positional embeddings (position indices -> vectors).

Parameters:

  • max_position_embeddings (int) – Maximum sequence length supported.

  • embedding_dim (int) – Dimension of the embedding vectors.

forward(positions: torch.Tensor) → torch.Tensor[source]

Forward pass for positional embeddings.

Parameters:

positions (torch.Tensor) – Tensor of position indices.

Returns:

Positional embeddings.

Return type:

torch.Tensor

class TokenEmbedding[source]

Bases: Module

Token embedding module. Positional embeddings are optional and explicit.

By default this returns token lookups only. Enable learned positional embeddings with use_position=True and providing max_position_embeddings.

Parameters:

  • vocab_size (int) – Size of the vocabulary.

  • embedding_dim (int) – Dimension of the embedding vectors.

  • padding_idx (int, optional) – Index for padding tokens. Defaults to None.

  • max_position_embeddings (int, optional) – Max sequence length for positional embeddings. Required if use_position=True. Defaults to None.

  • use_position (bool, optional) – Whether to include learned positional embeddings. Defaults to False.

  • dropout (float, optional) – Dropout probability. Defaults to 0.1.

forward(token_ids: torch.Tensor) → torch.Tensor[source]

Convert token IDs to embeddings.

Parameters:

token_ids (torch.Tensor) – Tensor of token IDs of shape (batch_size, seq_len).

Returns:

Embedded tokens of shape (batch_size, seq_len, embedding_dim).

Return type:

torch.Tensor