Skip to main content

Quickly initialize bespoke Transformer models

Project description


The goal of the xfmers library is to provide a simple API to quickly initialise Transformers with specified hyperparameters and features. The library generates standard TF 2.0 Keras models that can be used with, Automatic Mixed Precision (AMP), XLA, Horovod and tf.distribute APIs etc.

Included layers/features:

  • Multi-head attention
    • Encoder or Decoder (causal) mode
  • Transformer layers
    • Spatial convolution
    • Reversible layers
  • Transformer Stack (Encoder/Decoder)
    • Weight sharing (ALBERT-like)
  • Embedding Layer
    • Learnable positional embeddings
    • Factorized embedding parameterization (ALBERT-like)
  • Misc
    • Training schedules
    • Activation functions


Creating an ALBERT-like Transformer

Models can be created using Keras layers and trained using or Gradient Tape.

inputs = tf.keras.Input(shape=(None, ), name="inputs")
padding_mask = layers.PaddingMaskGenerator()(inputs)
embeddings = layers.TokenPosEmbedding(d_vocab=vocab_size, d_model=128, pos_length=512,
                                      # project embedding from 128 -> 512
encoder_block = layers.TransformerStack(layers=3,
                                        causal=False,        # attend pair-wise between all positons
                                        weight_sharing=True, # share weights between all encoder layers
enc_outputs = encoder_block({"token_inputs": embeddings,
                             "mask_inputs": padding_mask})
preds = layers.LMHead(vocab_size=vocab_size, name="outputs")(enc_outputs)

model = tf.keras.Model(inputs=inputs, outputs=preds, name=name)

Installing Xfmers

Install from Pip

pip install xfmers


  • Core Maintainer: Timothy Liu (tlkh)
  • Please open an issue if you encounter problems or have a feature request

Project details

Download files

Download the file for your platform. If you're not sure which to choose, learn more about installing packages.

Files for xfmers, version 0.0.3
Filename, size File type Python version Upload date Hashes
Filename, size xfmers-0.0.3-py3-none-any.whl (9.0 kB) File type Wheel Python version py3 Upload date Hashes View
Filename, size xfmers-0.0.3.tar.gz (7.8 kB) File type Source Python version None Upload date Hashes View

Supported by

Pingdom Pingdom Monitoring Google Google Object Storage and Download Analytics Sentry Sentry Error logging AWS AWS Cloud computing DataDog DataDog Monitoring Fastly Fastly CDN DigiCert DigiCert EV certificate StatusPage StatusPage Status page