langvae.encoders package

Submodules

langvae.encoders.automodel_presets module

class langvae.encoders.automodel_presets.AutoModelPreset(cls: str = 'AutoModelForTextEncoding', pooling_method: PoolingMethod = PoolingMethod.MEAN, normalize: bool = False)

Bases: object

Predefined settings class for encoder models loaded with AutoModel.

cls

Name of the class used for loading thge model [AutoModel | AutoModelForTextEncoding].

Type:

str

pooling_method

Method used for pooling the token embeddings [MEAN | LAST | CLS].

Type:

PoolingMethod

normalize

Whether the embeddings are to be normalized or not.

Type:

bool

cls: str = 'AutoModelForTextEncoding'

property cls_type: Type

normalize: bool = False

pooling_method: PoolingMethod = 'mean'

class langvae.encoders.automodel_presets.PoolingMethod(value, names=None, *, module=None, qualname=None, type=None, start=1, boundary=None)

Bases: StrEnum

CLS = 'cls'

LAST = 'last'

MEAN = 'mean'

langvae.encoders.sentence module

class langvae.encoders.sentence.SentenceEncoder(model_path: str, latent_size: int, decoder_tokenizer: PreTrainedTokenizer, caching: bool = False, automodel_preset: dict = None, device: str = 'cpu', args=None)

Bases: BaseEncoder

Encoder class for producing sentence embeddings.

This encoder uses a pre-trained transformer model to encode input sentences and then projects the encoded representations into a latent space defined by the specified latent size. The encoder is designed to work with one-hot encoded token tensors and utilizes a linear transformation to produce mean and log-variance for variational inference. The input tokenization is given by the tokenizer of the decoder model selected for the VAE model where this encoder will be connected to. So, for example in a T5-Llama setup the inputs are tokenized using the Llama tokenizer. SentenceEncoder takes care of converting the token representations. Inputs are given as tensors (B x S x V), where \(B\) is the batch size, \(S\) is the maximum sentence length and \(V\) is the decoder vocabulary size.

encoder

A HuggingFace pre-trained transformer model.

Type:

AutoModel

linear

Linear layer to project encoder outputs to latent space.

Type:

nn.Linear

tokenizer

Tokenizer corresponding to the encoder model.

Type:

AutoTokenizer

decoder_tokenizer

Tokenizer used for decoding input tokens.

Type:

PreTrainedTokenizer

caching

Whether the pre-trained encoder embeddings will be cached (in memory) or not.

Type:

bool

automodel_preset

Predefined settings for encoder models: Automodel loader class, pooling method, normalization.

Type:

dict

device

Device on which the model and data are allocated (e.g., ‘cpu’, ‘cuda’).

Type:

str

debug

Flag to enable/disable debugging output.

Type:

bool

property annotated: bool

property decoder_tokenizer: PreTrainedTokenizer

property encoder: Module

forward(x: Tensor, c: Dict[str, Tensor] = None) → ModelOutput

Processes the input tensor through the encoder and linear transformation to produce latent variables.

Parameters:

• x (Tensor) – Input tensor containing token IDs. Sparse tensors will be interpreted as one-hot.

• c (Tensor) – Conditional variable.

Returns:

Object containing the latent embedding and log covariance.

Return type:

ModelOutput

init_pretrained_model()

recode(tok_ids: list[torch.Tensor] | Tensor) → Tensor

to(device, include_pretrained: bool = True)

Move and/or cast the parameters and buffers.

This can be called as

to(device=None, dtype=None, non_blocking=False)

to(dtype, non_blocking=False)

to(tensor, non_blocking=False)

to(memory_format=torch.channels_last)

Its signature is similar to torch.Tensor.to(), but only accepts floating point or complex dtypes. In addition, this method will only cast the floating point or complex parameters and buffers to dtype (if given). The integral parameters and buffers will be moved device, if that is given, but with dtypes unchanged. When non_blocking is set, it tries to convert/move asynchronously with respect to the host if possible, e.g., moving CPU Tensors with pinned memory to CUDA devices.

See below for examples.

Note

This method modifies the module in-place.

Parameters:

• device (torch.device) – the desired device of the parameters and buffers in this module

• dtype (torch.dtype) – the desired floating point or complex dtype of the parameters and buffers in this module

• tensor (torch.Tensor) – Tensor whose dtype and device are the desired dtype and device for all parameters and buffers in this module

• memory_format (torch.memory_format) – the desired memory format for 4D parameters and buffers in this module (keyword only argument)

Returns:

self

Return type:

Module

Examples:

>>> # xdoctest: +IGNORE_WANT("non-deterministic")
>>> linear = nn.Linear(2, 2)
>>> linear.weight
Parameter containing:
tensor([[ 0.1913, -0.3420],
        [-0.5113, -0.2325]])
>>> linear.to(torch.double)
Linear(in_features=2, out_features=2, bias=True)
>>> linear.weight
Parameter containing:
tensor([[ 0.1913, -0.3420],
        [-0.5113, -0.2325]], dtype=torch.float64)
>>> # xdoctest: +REQUIRES(env:TORCH_DOCTEST_CUDA1)
>>> gpu1 = torch.device("cuda:1")
>>> linear.to(gpu1, dtype=torch.half, non_blocking=True)
Linear(in_features=2, out_features=2, bias=True)
>>> linear.weight
Parameter containing:
tensor([[ 0.1914, -0.3420],
        [-0.5112, -0.2324]], dtype=torch.float16, device='cuda:1')
>>> cpu = torch.device("cpu")
>>> linear.to(cpu)
Linear(in_features=2, out_features=2, bias=True)
>>> linear.weight
Parameter containing:
tensor([[ 0.1914, -0.3420],
        [-0.5112, -0.2324]], dtype=torch.float16)

>>> linear = nn.Linear(2, 2, bias=None).to(torch.cdouble)
>>> linear.weight
Parameter containing:
tensor([[ 0.3741+0.j,  0.2382+0.j],
        [ 0.5593+0.j, -0.4443+0.j]], dtype=torch.complex128)
>>> linear(torch.ones(3, 2, dtype=torch.cdouble))
tensor([[0.6122+0.j, 0.1150+0.j],
        [0.6122+0.j, 0.1150+0.j],
        [0.6122+0.j, 0.1150+0.j]], dtype=torch.complex128)

property tokenizer: PreTrainedTokenizer

langvae.encoders.sentence.cls_token_pooling(model_output: Tensor, attention_mask: Tensor) → Tensor

langvae.encoders.sentence.last_token_pooling(model_output: Tensor, attention_mask: Tensor) → Tensor

langvae.encoders.sentence.mean_pooling(model_output: Tensor, attention_mask: Tensor)

Module contents