espnet2.tok.token_embedding.TokenEmbedding

source

class espnet2.tok.token_embedding.TokenEmbedding(input_size: int, embed_dim: int, use_positional_encoding: bool = True, pos_enc_class: ~typing.Type[~torch.nn.modules.module.Module] = <class 'espnet2.legacy.nets.pytorch_backend.transformer.embedding.PositionalEncoding'>, positional_dropout_rate: float = 0.1)

Bases: AbsFrontend

Embed token IDs or differentiable assignments with shared weights.

Integer inputs of shape (B, T) use ordinary embedding lookup. A floating-point input of shape (B, T, K) is interpreted as a one-hot or soft distribution over the K tokens and multiplied by the embedding matrix. The latter path preserves gradients to a differentiable quantizer, including hard straight-through assignments.

Unlike espnet2.mt.frontend.embedding.Embedding, which is intended for discrete text token IDs, this frontend is used when tokenizer losses must backpropagate through soft distributions or hard one-hot straight-through assignments. It still accepts integer token IDs for inference and other discrete inputs.

• Parameters:
  • input_size – Number of input tokens K.
  • embed_dim – Embedding dimension.
  • use_positional_encoding – Apply positional encoding after embedding.
  • pos_enc_class – Positional-encoding class. It must accept the embedding dimension and dropout rate as its first two arguments.
  • positional_dropout_rate – Dropout rate used by positional encoding.

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

forward(input: Tensor, input_lengths: Tensor) → Tuple[Tensor, Tensor]

Embed token IDs or distributions without changing their lengths.

• Parameters:
  • input – Integer token IDs of shape (B, T) or floating-point token distributions of shape (B, T, K). Padded distributions must be all-zero vectors.
  • input_lengths – Valid lengths of shape (B,).
• Returns: Embedded features of shape (B, T, embed_dim) and the unchanged input lengths.

output_size() → int

Return the embedding dimension.