Wiener Transformer

This is the implementation of the Transformer model used with Wiener Loss.

class wiener_transformer.transformer.Decoder(layer, N)

Bases: Module

Generic N layer decoder with masking to prevent attending to future positions.

forward(x, memory, src_mask, tgt_mask)

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

class wiener_transformer.transformer.DecoderLayer(size, self_attn, src_attn, feed_forward, dropout)

Bases: Module

Decoder layer is made up of self-attention, source-attention, and feed-forward networks.

forward(x, memory, src_mask, tgt_mask)

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

class wiener_transformer.transformer.Embeddings(d_model, vocab_size, embedding_type='learned', word2vec_model=None, fasttext_model=None, glove_weights=None)

Bases: Module

Embedding layer that can be initialized with pre-trained embeddings (Word2Vec, FastText, GloVe) or learned embeddings.

forward(x)

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

init_fasttext_weights(fasttext_model)

init_glove_weights(glove_weights)

init_learned_weights()

init_word2vec_weights(word2vec_model)

class wiener_transformer.transformer.Encoder(layer, N)

Bases: Module

Core Encoder is a stack of N identical layers.

forward(x, mask)

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

class wiener_transformer.transformer.EncoderDecoder(encoder, decoder, src_embed, tgt_embed, generator)

Bases: Module

A standard Encoder-Decoder architecture. The Encoder processes the input sequence, and the Decoder generates the output sequence.

decode(memory, src_mask, tgt, tgt_mask)

encode(src, src_mask)

forward(src, tgt, src_mask, tgt_mask)

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

class wiener_transformer.transformer.EncoderLayer(size, self_attn, feed_forward, dropout)

Bases: Module

Encoder layer is made up of self-attention and feed-forward networks.

forward(x, mask)

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

class wiener_transformer.transformer.Generator(d_model, vocab)

Bases: Module

Linear layer followed by softmax to generate output probabilities over the target vocabulary.

forward(x)

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

class wiener_transformer.transformer.LayerNorm(features, eps=1e-06)

Bases: Module

Layer Normalization as introduced by Ba et al.

forward(x)

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

class wiener_transformer.transformer.MultiHeadedAttention(h, d_model, dropout=0.1)

Bases: Module

Multi-Head Attention mechanism, which allows the model to jointly attend to information from different representation subspaces.

forward(query, key, value, mask=None)

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

class wiener_transformer.transformer.PositionalEncoding(d_model, dropout, max_len=5000)

Bases: Module

Add positional encoding to the input embeddings to provide information about the position of the tokens in the sequence.

forward(x)

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

class wiener_transformer.transformer.PositionwiseFeedForward(d_model, d_ff, dropout=0.1)

Bases: Module

Implements the position-wise feed-forward network.

forward(x)

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

class wiener_transformer.transformer.SublayerConnection(size, dropout)

Bases: Module

A residual connection followed by a layer normalization.

forward(x, sublayer)

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

wiener_transformer.transformer.clones(module, N)

wiener_transformer.transformer.make_model(src_vocab, tgt_vocab, N=6, d_model=512, d_ff=2048, h=8, dropout=0.1, embedding_type='learned')

Helper function to construct a model from hyperparameters.

Parameters:

• src_vocab – Size of source vocabulary.

• tgt_vocab – Size of target vocabulary.

• N – Number of layers in the encoder and decoder.

• d_model – Dimensionality of the embeddings.

• d_ff – Dimensionality of the feed-forward network.

• h – Number of attention heads.

• dropout – Dropout rate.

• embedding_type – Type of embeddings (‘learned’, ‘word2vec’, ‘fasttext’, ‘glove’).

Returns:

A constructed EncoderDecoder model.

wiener_transformer.transformer.subsequent_mask(size)

Utilities

These are the utilities associated with the Wiener Loss Transformer.

Train

wiener_transformer.utils.train.load_trained_model(config)

Load a trained model from disk. If no trained model exists, train it first.

Parameters:

config – Configuration dictionary.

Returns:

Trained model.

wiener_transformer.utils.train.train_model(train, valid, test, src_tokenizer, tgt_tokenizer, config, accelerator=None, model=None, optimizer=None, lr_scheduler=None)

Train the model using the specified configuration and dataset.

Parameters:

• train – Training dataset.

• valid – Validation dataset.

• test – Test dataset.

• src_tokenizer – Source tokenizer.

• tgt_tokenizer – Target tokenizer.

• config – Configuration dictionary.

• accelerator – Accelerator object for distributed training (optional).

• model – Preloaded model (optional).

• optimizer – Preloaded optimizer (optional).

• lr_scheduler – Preloaded learning rate scheduler (optional).

Embeddings

wiener_transformer.utils.embeddings.create_embedding_weights(model_type='word2vec', vector_size=512, batch_size=1)

wiener_transformer.utils.embeddings.extract_sentences(dataloader)

wiener_transformer.utils.embeddings.load_glove_embeddings(vector_size=512)

Wiener Loss

class wiener_transformer.utils.wienerloss.WienerLoss(method='fft', filter_dim=2, filter_scale=2, reduction='mean', mode='reverse', penalty_function=None, store_filters=False, epsilon=0.0001, std=0.0001, clamp_min=None, corr_norm=None, rel_epsilon=False)

Bases: Module

This loss was implemented by Cruz et al in the paper Convolve and Conquer: Data Comparison with Wiener Filters. I have modified the implementation:

1. Ensured that the identity function can work with any number of dimensions

Original implementation: https://github.com/dpelacani/WienerLoss

The AWLoss class implements the adaptive Wiener criterion, which aims to compare two data samples through a convolutional filter. The methodology is inspired by the paper `Adaptive Waveform Inversion: Theory`_ (Warner and Guasch, 2014).

A matching filter w can be computed such that it transforms a targetsignal p into the data d under an L2 norm principle:

g = || p*w - d ||^2

Let ‘Z’ be the Toeplitz matrix formulation of p such that the equation above is equivalent to: g = || Zw - d ||^2

Minimizing this functional: dgdw = Z^T (Zw - d) dgdw –> 0 : w = (Z^T @ Z)^(-1) @ Z^T @ d

To stabilize the matrix inversion, an amount is added to the diagonal of (Z^T @ Z) based on a value epsilon such that the inverted matrix is (Z^T @ Z) + max(diagonal(Z^T @ Z)) * epsilon

In 2D, convolving p with w (or w with p) is equivalent to the matrix vector multiplication Zd where Z is the doubly block Toeplitz of the reconstructed image P and w is the flattened array of the 2D kernel W.

Therefore, the system is equivalent to solving || Zw - d ||^2, and the solution to w is given by w = (Z^T @ Z + max(diagonal(Z^T @ Z)) * epsilon)^(-1) @ Z^T @ d

This composes the direct method.

Alternatively, convolution can be performed in the frequency domain with multiplication and division operations. This tends to be much more computationally efficient.

The criterion is evaluated through a symmetrical monotonically decreasing function T and a dirac delta function that rewards when the filter kernel w is close to the identity kernel, and penalizes otherwise.

f = 1/2 ||T * (v - delta)||^2

Parameters:

• method – “fft” for Fast Fourier Transform or “direct” for the Levinson-Durbin recurssion algorithm. Defaults to “fft”

• optional – “fft” for Fast Fourier Transform or “direct” for the Levinson-Durbin recurssion algorithm. Defaults to “fft”

• filter_dim – the dimensionality of the filter. This parameter should be upper-bounded by the dimensionality of the data. If data is 3-dimensional and filter_dim is set to 2, one filter is computed per channel dimension assuming format [B, NC, H , W]. Current implementation only supports filter dimensions for 1D, 2D and 3D. Defaults to 2

• optional – the dimensionality of the filter. This parameter should be upper-bounded by the dimensionality of the data. If data is 3-dimensional and filter_dim is set to 2, one filter is computed per channel dimension assuming format [B, NC, H , W]. Current implementation only supports filter dimensions for 1D, 2D and 3D. Defaults to 2

• filter_scale – the scale of the filters compared to the size of the data. Defaults to 2

• optional – the scale of the filters compared to the size of the data. Defaults to 2

• reduction – specifies the reduction to apply to the output, “mean” or “sum”. Defaults to mean

• optional – specifies the reduction to apply to the output, “mean” or “sum”. Defaults to mean

• mode – “forward” or “reverse” computation of the filter. For details of the difference, refer to the original paper. Default “reverse”

• optional – “forward” or “reverse” computation of the filter. For details of the difference, refer to the original paper. Default “reverse”

• penalty_function – the penalty function to apply to the filter. If None, the penalty function is the identity. Takes “identity”, “gaussian” or custom penalty function. Default None

• optional – the penalty function to apply to the filter. If None, the penalty function is the identity. Takes “identity”, “gaussian” or custom penalty function. Default None

• std – the standard deviation of the gaussian when penalty_function=”gaussian”. Mean is always zero. Default None

• optional – the standard deviation of the gaussian when penalty_function=”gaussian”. Mean is always zero. Default None

• store_filters – whether to store the filters in memory, useful for debugging. Option to store the filers before or after normalisation with “norm” and “unorm”. Default False.

• optional – whether to store the filters in memory, useful for debugging. Option to store the filers before or after normalisation with “norm” and “unorm”. Default False.

• epsilon – the stabilization value to compute the filter. Default 1e-4.

• optional – the stabilization value to compute the filter. Default 1e-4.

• clamp_min – filters are clipped to this minimum value after computation. If None, operation is disabled. Default none

• optional – filters are clipped to this minimum value after computation. If None, operation is disabled. Default none

forward(recon, target, epsilon=None, gamma=0.0, eta=0.0)

> The function takes in a reconstructed signal, a target signal, and a few other parameters, and returns the loss

Args

recon

the reconstructed signal

target

the target signal

epsilon, optional

the stabilization value to compute the filter. If passed, overwrites the class attribute of same name. Default None.

gamma, optional

noise to add to both target and reconstructed signals for training stabilization. Default 0.

eta, optional

noise to add to penalty function. Default 0.

get_filter_shape(input_shape)

identity(mesh, val=1, **kwargs)

make_delta(shape)

make_doubly_block(X)

Makes Doubly Blocked Toeplitz of a matrix X [r, c]

make_penalty(shape, eta=0.0, penalty_function=None, std=None, device='cpu')

make_toeplitz(a)

Makes toeplitz matrix of a vector A

multigauss(mesh, mean, covmatrix)

Multivariate gaussian of N dimensions on evenly spaced hypercubed grid. Mesh should be stacked along the last axis E.g. for a 3D gaussian of 20 grid points in each axis mesh should be of shape (20, 20, 20, 3)

pad_signal(x, shape, val=0)

x must be a multichannel signal of shape [batch_size, nchannels, width, height]

rms(x)

wiener(x, y, fs, epsilon=1e-09)

calculates the optimal least squares convolutional Wiener filter that transforms signal x into signal y using the direct Toeplitz matrix implementation

wienerfft(x, y, fs, prwh=1e-09)

George Strong (geowstrong@gmail.com) calculates the optimal least squares convolutional Wiener filter that transforms signal x into signal y using FFT

Vocabulary

wiener_transformer.utils.vocab.batch_iterator(dataset, lang='en', batch_size=1000)

Generator function that yields batches of text data from the dataset.

Parameters:

• dataset – The dataset to iterate over.

• lang – The language to extract text from (‘en’ or ‘de’).

• batch_size – The size of each batch to yield.

Yields:

Batches of text data.

wiener_transformer.utils.vocab.load_tokenizers(dataset)

Loads or trains tokenizers for the source and target languages. If pre-trained tokenizers are available on disk, they are loaded; otherwise, new tokenizers are trained on the dataset.

Parameters:

dataset – The dataset to use for training the tokenizers if needed.

Returns:

The source and target tokenizers.

wiener_transformer.utils.vocab.load_wmt(language_pair=('en', 'de'))

Loads the WMT dataset for the specified language pair using the HuggingFace datasets library.

Parameters:

language_pair – A tuple specifying the source and target languages (default is (‘en’, ‘de’)).

Returns:

The training, validation, and test splits of the WMT dataset.

wiener_transformer.utils.vocab.transform(examples)

Transforms the dataset examples into a format suitable for tokenization.

Parameters:

examples – A dictionary containing ‘translation’ key with text data.

Returns:

A dictionary with separate keys for source (‘en’) and target (‘de’) languages.

Helpers

class wiener_transformer.utils.helpers.Batch(src, tgt=None, pad=2)

Bases: object

Object for holding a batch of data with masks during training.

Parameters:

• src – Source sequences for the batch.

• tgt – Target sequences for the batch (default: None).

• pad – Index used for padding (default: 2).

src

Source sequences for the batch.

src_mask

Mask indicating non-padding elements in the source sequences.

tgt

Target sequences without the last token (teacher forcing).

tgt_y

Target sequences without the first token (teacher forcing).

tgt_mask

Mask indicating non-padding elements and future positions in the target sequences.

ntokens

Total number of non-padding tokens in the target sequences.

static make_std_mask(tgt, pad)

Create a mask to hide padding and future words in the target sequences.

Parameters:

• tgt – Target sequences.

• pad – Index used for padding.

Returns:

Tensor representing the mask for the target sequences.

class wiener_transformer.utils.helpers.DummyOptimizer

Bases: Optimizer

A dummy optimizer that does nothing. Used for testing or evaluation without updates.

step()

Placeholder method to simulate an optimizer step.

zero_grad(set_to_none=False)

Placeholder method to simulate zeroing gradients.

step()

Performs a single optimization step (parameter update).

Parameters:

closure (Callable) – A closure that reevaluates the model and returns the loss. Optional for most optimizers.

Note

Unless otherwise specified, this function should not modify the .grad field of the parameters.

zero_grad(set_to_none=False)

Resets the gradients of all optimized torch.Tensor s.

Parameters:

set_to_none (bool) – instead of setting to zero, set the grads to None. This will in general have lower memory footprint, and can modestly improve performance. However, it changes certain behaviors. For example: 1. When the user tries to access a gradient and perform manual ops on it, a None attribute or a Tensor full of 0s will behave differently. 2. If the user requests zero_grad(set_to_none=True) followed by a backward pass, .grads are guaranteed to be None for params that did not receive a gradient. 3. torch.optim optimizers have a different behavior if the gradient is 0 or None (in one case it does the step with a gradient of 0 and in the other it skips the step altogether).

class wiener_transformer.utils.helpers.DummyScheduler

Bases: object

A dummy learning rate scheduler that does nothing. Used for testing or evaluation without updates.

step()

Placeholder method to simulate a scheduler step.

step()

class wiener_transformer.utils.helpers.TrainState

Bases: object

Track the number of steps, examples, and tokens processed during training.

step

Number of steps in the current epoch.

Type:

int

accum_step

Number of gradient accumulation steps.

Type:

int

samples

Total number of examples used.

Type:

int

tokens

Total number of tokens processed.

Type:

int

accum_step: int = 0

samples: int = 0

step: int = 0

tokens: int = 0

wiener_transformer.utils.helpers.calculate_bleu(model, test_dataloader, pad_idx, tgt_tokenizer, max_len=50)

Calculate BLEU scores and BERTScore for model predictions on the test dataset.

Parameters:

• model – The trained model to be evaluated.

• test_dataloader – DataLoader providing batches of test data.

• pad_idx – The index used for padding in the sequences.

• tgt_tokenizer – Tokenizer for the target language.

• max_len – Maximum length for generated sequences (default: 50).

Returns:

• corpus_bleu_score: BLEU score for the entire corpus.

• average_sentence_bleu_score: Average BLEU score for individual sentences.

• sacrebleu_score: SacreBLEU score for the corpus.

• recall: Average BERTScore recall.

• precision: Average BERTScore precision.

• f1: Average BERTScore F1 score.

Return type:

Tuple containing

wiener_transformer.utils.helpers.pad(tensor, pad, value)

Pad a tensor with a specified value.

Parameters:

• tensor – The tensor to be padded.

• pad – A tuple specifying the padding to be applied to each dimension.

• value – The value to use for padding.

Returns:

Padded tensor.

wiener_transformer.utils.helpers.rate(step, model_size, factor, warmup)

Compute the learning rate according to the warmup strategy.

Parameters:

• step – Current step number.

• model_size – Dimensionality of the model.

• factor – Scaling factor for the learning rate.

• warmup – Number of warmup steps.

Returns:

Learning rate value for the current step.

wiener_transformer.utils.helpers.run_epoch(data_iter, model, wiener_fn, classic_loss, optimizer, scheduler, accelerator, mode='train', loss_weight=1, accum_iter=1, train_state=<wiener_transformer.utils.helpers.TrainState object>, max_batches=inf, gamma=0.1, awl=None)

Train or evaluate the model for a single epoch.

Parameters:

• data_iter – DataLoader providing batches of data.

• model – The model being trained or evaluated.

• wiener_fn – Function to compute the Wiener loss (if applicable).

• classic_loss – Function to compute the standard loss (e.g., cross-entropy).

• optimizer – Optimizer used for training.

• scheduler – Learning rate scheduler.

• accelerator – Distributed training utility (e.g., from Hugging Face’s Accelerate).

• mode – Mode of operation, either “train” or “eval” (default: “train”).

• loss_weight – Weighting factor for the loss function (default: 1).

• accum_iter – Number of steps for gradient accumulation (default: 1).

• train_state – Object tracking training state (default: TrainState()).

• max_batches – Maximum number of batches to process (default: inf).

• gamma – Regularization parameter for the Wiener loss (default: 0.1).

• awl – Automatic weighting of losses (if applicable).

Returns:

• avg_loss: Average combined loss over the epoch.

• avg_wiener: Average Wiener loss over the epoch.

• avg_kldiv: Average Kullback-Leibler divergence loss over the epoch.

• train_state: Updated training state after the epoch.

Return type:

Tuple containing

wiener_transformer.utils.helpers.subsequent_mask(size)

Create a mask to hide subsequent positions in a sequence.

Parameters:

size – The length of the sequence.

Returns:

Tensor of shape (1, size, size) with True values in the upper triangular part.

wiener_transformer.utils.helpers.tokenize(text, tokenizer)

Tokenize a given text using a specified tokenizer.

Parameters:

• text – The text to be tokenized.

• tokenizer – The tokenizer to be used for tokenization.

Returns:

List of tokens.

Data Loader

class wiener_transformer.utils.data_loader.WMT14Dataset(data_iter)

Bases: Dataset

A custom Dataset class for the WMT14 dataset.

Parameters:

data_iter – An iterable containing the data samples.

__len__()

Returns the length of the dataset.

__getitem__(idx)

Returns the data sample at the given index.

wiener_transformer.utils.data_loader.collate_batch(batch, src_pipeline, tgt_pipeline, device, max_padding=128, pad_id=2)

Collate a batch of data for the DataLoader, applying tokenization, padding, and conversion to tensors.

Parameters:

• batch – A list of data examples where each example is a dictionary containing ‘translation’ with ‘en’ and ‘de’ keys.

• src_pipeline – A tokenization function for the source text.

• tgt_pipeline – A tokenization function for the target text.

• device – The device to place the tensors on (e.g., ‘cpu’ or ‘cuda’).

• max_padding – The maximum length to pad the sequences to (default: 128).

• pad_id – The token ID used for padding (default: 2).

Returns:

• src: Padded and tokenized source sequences.

• tgt: Padded and tokenized target sequences.

Return type:

A tuple of tensors (src, tgt) where

wiener_transformer.utils.data_loader.create_dataloaders(train_data, valid_data, test_data, src_tokenizer, tgt_tokenizer, device, batch_size=32, max_padding=128)

Create DataLoader objects for the training, validation, and test datasets.

Parameters:

• train_data – Iterable containing the training data.

• valid_data – Iterable containing the validation data.

• test_data – Iterable containing the test data.

• src_tokenizer – Tokenizer for the source language.

• tgt_tokenizer – Tokenizer for the target language.

• device – The device to place the tensors on (e.g., ‘cpu’ or ‘cuda’).

• batch_size – Number of samples per batch (default: 32).

• max_padding – The maximum length to pad the sequences to (default: 128).

Returns:

Tuple containing DataLoader objects for the training, validation, and test datasets.