alibi_detect.models.pytorch.embedding module

class alibi_detect.models.pytorch.embedding.TransformerEmbedding(model_name_or_path, embedding_type, layers=None)[source]

Bases: Module

forward(tokens)[source]
Return type:

Tensor

alibi_detect.models.pytorch.embedding.hidden_state_embedding(hidden_states, layers, use_cls, reduce_mean=True)[source]

Extract embeddings from hidden attention state layers.

Parameters:

  • hidden_states (Tensor) – Attention hidden states in the transformer model.

  • layers (List[int]) – List of layers to use for the embedding.

  • use_cls (bool) – Whether to use the next sentence token (CLS) to extract the embeddings.

  • reduce_mean (bool) – Whether to take the mean of the output tensor.

Return type:

Tensor

Returns:

Tensor with embeddings.