Modules

Attention Bridge

class mammoth.modules.attention_bridge.AttentionBridge(layers)

Bases: Module

N-layered attention-bridge between encoders->decoders

forward(enc_output, mask)

Forward pass for the bridge layers

classmethod from_opts(opts)

Alternate constructor.

class mammoth.modules.attention_bridge.BaseAttentionBridgeLayer(*args: Any, **kwargs: Any)

Bases: Module

Base class for attention bridge layers

property is_fixed_length: bool

whether this layer always produce output of the same size

Encoders

class mammoth.modules.encoder.EncoderBase(*args: Any, **kwargs: Any)

Bases: Module

Base encoder class. Specifies the interface used by different encoder types and required by mammoth.Models.NMTModel.

        graph BT
   A[Input]
   subgraph RNN
     C[Pos 1]
     D[Pos 2]
     E[Pos N]
   end
   F[Memory_Bank]
   G[Final]
   A-->C
   A-->D
   A-->E
   C-->F
   D-->F
   E-->F
   E-->G
    

forward(src, lengths=None)

Parameters:

• src (LongTensor) – padded sequences of sparse indices (src_len, batch, nfeat)

• lengths (LongTensor) – length of each sequence (batch,)

Returns:

• final encoder state, used to initialize decoder

• memory bank for attention, (src_len, batch, hidden)

• lengths

Return type:

(FloatTensor, FloatTensor, FloatTensor)

class mammoth.modules.transformer_encoder.TransformerEncoder(num_layers, d_model, heads, d_ff, dropout, attention_dropout, embeddings, max_relative_positions, pos_ffn_activation_fn='relu', layer_norm_module=None, is_normformer=False)

Bases: EncoderBase

The Transformer encoder from “Attention is All You Need” [VSP+17]

        graph BT
   A[input]
   B[multi-head self-attn]
   C[feed forward]
   O[output]
   A --> B
   B --> C
   C --> O
    

Parameters:

• num_layers (int) – number of encoder layers

• d_model (int) – size of the model

• heads (int) – number of heads

• d_ff (int) – size of the inner FF layer

• dropout (float) – dropout parameters

• embeddings (mammoth.modules.Embeddings) – embeddings to use, should have positional encodings

• pos_ffn_activation_fn (ActivationFunction) – activation function choice for PositionwiseFeedForward layer

• is_normformer (bool) – whether to apply normformer-style normalization

Returns:

• embeddings (src_len, batch_size, model_dim)

• memory_bank (src_len, batch_size, model_dim)

Return type:

(torch.FloatTensor, torch.FloatTensor)

forward(src, lengths=None, skip_embedding=False, mask=None)

See EncoderBase.forward()

classmethod from_opts(opts, embeddings, is_on_top=False)

Alternate constructor.

class mammoth.modules.mean_encoder.MeanEncoder(num_layers, embeddings)

Bases: EncoderBase

A trivial non-recurrent encoder. Simply applies mean pooling.

Parameters:

• num_layers (int) – number of replicated layers

• embeddings (mammoth.modules.Embeddings) – embedding module to use

forward(src, lengths=None)

See EncoderBase.forward()

classmethod from_opts(opts, embeddings)

Alternate constructor.

class mammoth.modules.layer_stack.StackXcoder(task_wrappers: Dict[str, NativeTransformerWrapper], attention_layer_blocks: Dict[int, Dict[str, TransformerStack]], token_embs: Dict[str, Embedding], shared_stacks: ModuleDict, shared_token_embs: ModuleDict, shared_post_emb_norms: ModuleDict, shared_rotary_embs: ModuleDict, shared_to_logits: ModuleDict)

Bases: Module

Switches between different NativeTransformerWrappers depending on the task.

StackXcoder is the single registered owner of all shared parameters. NativeTransformerWrapper instances in task_wrappers hold plain Python references to the shared modules (not nn.Module children), so state_dict() contains exactly one copy of each shared parameter.

Sublayers

class mammoth.modules.average_attn.AverageAttention(model_dim, dropout=0.1, aan_useffn=False, pos_ffn_activation_fn='relu')

Bases: Module

Average Attention module from “Accelerating Neural Transformer via an Average Attention Network” [ZXS18].

Parameters:

• model_dim (int) – the dimension of keys/values/queries, must be divisible by head_count

• dropout (float) – dropout parameter

• pos_ffn_activation_fn (ActivationFunction) – activation function choice for PositionwiseFeedForward layer

cumulative_average(inputs, mask_or_step, layer_cache=None, step=None)

Computes the cumulative average as described in [ZXS18] – Equations (1) (5) (6)

Parameters:

• inputs (FloatTensor) – sequence to average (batch_size, input_len, dimension)

• mask_or_step – if cache is set, this is assumed to be the current step of the dynamic decoding. Otherwise, it is the mask matrix used to compute the cumulative average.

• layer_cache – a dictionary containing the cumulative average of the previous step.

Returns:

a tensor of the same shape and type as inputs.

cumulative_average_mask(batch_size, inputs_len, device)

Builds the mask to compute the cumulative average as described in [ZXS18] – Figure 3

Parameters:

• batch_size (int) – batch size

• inputs_len (int) – length of the inputs

Returns:

• A Tensor of shape (batch_size, input_len, input_len)

Return type:

(FloatTensor)

forward(inputs, mask=None, layer_cache=None, step=None)

Parameters:

inputs (FloatTensor) – (batch_size, input_len, model_dim)

Returns:

• gating_outputs (batch_size, input_len, model_dim)

• average_outputs average attention

  (batch_size, input_len, model_dim)

Return type:

(FloatTensor, FloatTensor)

class mammoth.modules.multi_headed_attn.MultiHeadedAttention(head_count, model_dim, dropout=0.1, max_relative_positions=0)

Bases: Module

Multi-Head Attention module from “Attention is All You Need” [VSP+17].

Similar to standard dot attention but uses multiple attention distributions simulataneously to select relevant items.

        graph BT
   A[key]
   B[value]
   C[query]
   O[output]
   subgraph Attn
     D[Attn 1]
     E[Attn 2]
     F[Attn N]
   end
   A --> D
   C --> D
   A --> E
   C --> E
   A --> F
   C --> F
   D --> O
   E --> O
   F --> O
   B --> O
    

Also includes several additional tricks.

Parameters:

• head_count (int) – number of parallel heads

• model_dim (int) – the dimension of keys/values/queries, must be divisible by head_count

• dropout (float) – dropout parameter

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

Compute the context vector and the attention vectors.

Parameters:

• key (FloatTensor) – set of key_len key vectors (batch, key_len, dim)

• value (FloatTensor) – set of key_len value vectors (batch, key_len, dim)

• query (FloatTensor) – set of query_len query vectors (batch, query_len, dim)

• mask – binary mask 1/0 indicating which keys have zero / non-zero attention (batch, query_len, key_len)

Returns:

• output context vectors (batch, query_len, dim)

• Attention vector in heads (batch, head, query_len, key_len).

Return type:

(FloatTensor, FloatTensor)

class mammoth.modules.position_ffn.PositionwiseFeedForward(d_model, d_ff, dropout=0.1, activation_fn='gelu', is_normformer=False)

Bases: Module

A two-layer Feed-Forward-Network with residual layer norm.

Parameters:

• d_model (int) – the size of input for the first-layer of the FFN.

• d_ff (int) – the hidden layer size of the second-layer of the FNN.

• dropout (float) – dropout probability in \([0, 1)\).

• activation_fn (ActivationFunction) – activation function used.

• is_normformer (bool) – whether to apply normformer-style normalization

forward(x)

Layer definition.

Parameters:

x – (batch_size, input_len, model_dim)

Returns:

Output (batch_size, input_len, model_dim).

Return type:

(FloatTensor)