gt4sd.frameworks.granular.ml.models.vae_trans.core module

VaeTrans implementation.

Summary

Classes:

VaeTrans

Transformer-based VAE with Gaussian Prior and approx posterior.

Reference

class VaeTrans(name, position, data, vocab_size, tokenizer, hidden_size_enc=256, n_layers_enc=2, hidden_size_dec=256, n_layers_dec=2, kl_coeff=0.1, latent_size=196, feedforward_size=512, heads=4, dropout=0.1, bypass_bottleneck=False, seq_len=127, teacher_forcing=True, loss_function='ce', kl_low=0.0, kl_high=0.1, kl_n_epochs=100, kl_start_epoch=0, inference_check_frequency=50, **kwargs)[source]

Bases: GranularEncoderDecoderModel

Transformer-based VAE with Gaussian Prior and approx posterior.

__init__(name, position, data, vocab_size, tokenizer, hidden_size_enc=256, n_layers_enc=2, hidden_size_dec=256, n_layers_dec=2, kl_coeff=0.1, latent_size=196, feedforward_size=512, heads=4, dropout=0.1, bypass_bottleneck=False, seq_len=127, teacher_forcing=True, loss_function='ce', kl_low=0.0, kl_high=0.1, kl_n_epochs=100, kl_start_epoch=0, inference_check_frequency=50, **kwargs)[source]

Construct VaeRnn.

Parameters

  • name (str) – model name.

  • position (int) – position of the model.

  • data (Dict[str, str]) – data name mappings.

  • vocab_size (int) – size of the vocabulary.

  • tokenizer – tokenizer.

  • hidden_size_enc (int) – encoder hidden size. Defaults to 256.

  • n_layers_enc (int) – number of layers for the encoder. Defaults to 2.

  • hidden_size_dec (int) – decoder hidden size. Defaults to 256.

  • n_layers_dec (int) – number of layers for the decoder. Defaults to 2.

  • kl_coeff (float) – KL coefficient. Defaults to 0.1.

  • latent_size (int) – latent size. Defaults to 196.

  • feedforward_size (int) – size of the feed forward network. Default to 512.

  • heads (int) – number of heads. Defauls to 4.

  • dropout (float) – dropout rate. Defaults to 0.1.

  • bypass_bottleneck (bool) – whether the bottleneck should be by passed. Defaults to False.

  • seq_len (int) – length of the sequence. Defaults to 127.

  • teacher_forcing (bool) – whether to teacher forcing. Defaults to True.

  • loss_function (str) – loss function. Defaults to “ce”.

  • kl_low (float) – low KL weight. Defaults to 0.0.

  • kl_high (float) – high KL weight. Defaults to 0.1.

  • kl_n_epochs (int) – KL number of epochs. Defaults to 100.

  • kl_start_epoch (int) – KL starting epoch. Defaults to 0.

  • inference_check_frequency (int) – frequency for checking inference quality. Defaults to 50.

Raises

ValueError – in case the provided loss function is not supported.

forward(x, tgt, *args, **kwrgs)[source]

Forward pass in the model.

Parameters

  • x (Any) – model input.

  • tgt (Tensor) – target tensor

Return type

Any

Returns

model output.

predict_mask_length(mem)[source]

Predicts mask length from latent memory so mask can be re-created during inference.

Parameters

mem (Tensor) – latent memory.

Return type

Any

Returns

mask length.

_sampling_step(x, *args, **kwargs)[source]

Run a sampling step in the model.

Parameters

x (Any) – model input.

Return type

Any

Returns

model sampling step output.

encode(x, *args, **kwargs)[source]

Encode a sample.

Parameters

x (Any) – input sample.

Return type

Any

Returns

latent encoding.

decode(z, *args, **kwargs)[source]

Decode a latent space point.

Parameters

z (Any) – latent point.

Return type

Any

Returns

tuple with decoded texts and token indices.

encode_decode(x, *args, **kwargs)[source]

Encode and decode a sample.

Parameters

x (Any) – input sample.

Return type

Any

Returns

decoded sample.

inference(x, *args, **kwargs)[source]

Run the model in inference mode.

Parameters

x (Any) – sample.

Return type

Any

Returns

generated output.

_run_step(x, tgt)[source]

Run a step in the model.

Parameters

  • x (Any) – model input.

  • tgt (Tensor) – target tensor

Return type

Any

Returns

model step output.

step(input_data, target_data, device='cpu', current_epoch=0, *args, **kwargs)[source]

Training step for the model.

Parameters

  • input_data (Any) – input for the step.

  • target_data (Any) – target for the step.

  • device (str) – string representing the device to use. Defaults to “cpu”.

  • current_epoch (int) – current epoch. Defaults to 0.

Return type

Tuple[Any, Any, Any]

Returns

a tuple containing the step output, the loss and the logs for the module.

val_step(input_data, target_data, device='cpu', current_epoch=0, *args, **kwargs)[source]

Validation step for the model.

Parameters

  • input_data (Any) – input for the step.

  • target_data (Any) – target for the step.

  • device (str) – string representing the device to use. Defaults to “cpu”.

  • current_epoch (int) – current epoch. Defaults to 0.

Return type

Any

Returns

a tuple containing the step output, the loss and the logs for the module.

static add_model_specific_args(parent_parser, name, *args, **kwargs)[source]

Adding to a parser model specific arguments.

Parameters

  • parent_parser (ArgumentParser) – patent parser.

  • name (str) – model name.

Return type

ArgumentParser

Returns

updated parser.

__annotations__ = {}
__doc__ = 'Transformer-based VAE with Gaussian Prior and approx posterior.'
__module__ = 'gt4sd.frameworks.granular.ml.models.vae_trans.core'