<jupyter_start><jupyter_text>Concept - TransformerLM Quick Start and Guide > In this notebook, we will Use a pre-trained TransformerLM, Train a TransformerLM model, and Looking inside the Trax TransformerLM- toc: true- badges: true- comments: true- categories: [Concept, NLP, Transformer]- author: "<NAME>"- image: Language models are machine learning models that power some of the most impressive applications involving text and language (e.g. machine translation, sentiment analysis, chatbots, summarization). At the time of this writing, some of the largest ML models in existence are language models. They are also based on the [transformer](https://arxiv.org/abs/1706.03762) architecture. The transformer language model (TransformerLM) is a simpler [variation](https://arxiv.org/pdf/1801.10198.pdf) of the original transformer architecture and is useful for plenty of tasks.![](https://storage.googleapis.com/ml-intro/t/transformerLM-1.png)The [Trax](https://trax-ml.readthedocs.io/en/latest/) implementation of TransformerLM focuses on clear code and speed.  It runs without any changes on CPUs, GPUs and TPUs.In this notebook, we will:1. Use a pre-trained TransformerLM2. Train a TransformerLM model3. Looking inside the Trax TransformerLM<jupyter_code>import os
import numpy as np
! pip install -q -U trax
import trax<jupyter_output>[K     |████████████████████████████████| 419kB 2.8MB/s 
[K     |████████████████████████████████| 1.5MB 8.4MB/s 
[K     |████████████████████████████████| 163kB 21.2MB/s 
[K     |████████████████████████████████| 2.6MB 18.7MB/s 
[K     |████████████████████████████████| 194kB 35.5MB/s 
[K     |████████████████████████████████| 368kB 37.9MB/s 
[K     |████████████████████████████████| 307kB 49.1MB/s 
[K     |████████████████████████████████| 983kB 47.3MB/s 
[K     |████████████████████████████████| 358kB 49.9MB/s 
[K     |████████████████████████████████| 81kB 9.3MB/s 
[K     |████████████████████████████████| 5.3MB 49.0MB/s 
[K     |████████████████████████████████| 655kB 50.9MB/s 
[K     |████████████████████████████████| 71kB 8.3MB/s 
[K     |████████████████████████████████| 1.1MB 49.3MB/s 
[K     |████████████████████████████████| 3.5MB 49.2MB/s 
[K     |████████████████████████████████| 1.1MB 34.8MB/s 
[K     |████████████████████████████████| 245kB 51.3MB/s 
[K[...]<jupyter_text>Using a pre-trained TransformerLMThe following cell loads a pre-trained TransformerLM that sorts a list of four integers.<jupyter_code># Create a Transformer model.
# Have to use the same configuration of the pre-trained model we'll load next
model = trax.models.TransformerLM(  
    d_model=32, d_ff=128, n_layers=2, 
    vocab_size=32, mode='predict')

# Initialize using pre-trained weights.
model.init_from_file('gs://ml-intro/models/sort-transformer.pkl.gz',
                     weights_only=True, 
                     input_signature=trax.shapes.ShapeDtype((1,1), dtype=np.int32))

# Input sequence
# The 0s indicate the beginning and end of the input sequence
input = [0, 3, 15, 14, 9, 0]


# Run the model
output = trax.supervised.decoding.autoregressive_sample(
    model, np.array([input]), temperature=0.0, max_length=4)

# Show us the output
output<jupyter_output><empty_output><jupyter_text>This is a trivial example to get you started and put a toy transformer into your hands. Language models get their name from their ability to assign probabilities to sequences of words. This property makes them useful for generating text (and other types of sequences) by probabilistically choosing the next item in the sequence (often the highest probability one)  -- exactly like the next-word suggestion feature of your smartphone keyboard.In Trax, TransformerLM is a series of [Layers]() combined using the [Serial]() combinator. A high level view of the TransformerLM we've declared above can look like this:![](https://storage.googleapis.com/ml-intro/t/transformerLM-layers-1.png)The model has two decoder layers because we set `n_layers` to 2. TransformerLM makes predictions by being fed one token at a time, with output tokens typically fed back as inputs (that's the `autoregressive` part of the `autoregressive_sample` method we used to generate the output from the model). If we're to think of a simple model trained to generate the fibonacci sequence, we can give it a number in the sequence and it would continue to generate the next items in the sequence:![](https://storage.googleapis.com/ml-intro/t/transformerLM-input-output-fib.gif) Train a TransformerLM ModelLet's train a TransformerLM model. We'll train this one to reverse a list of integers. This is another toy task that we can train a small transformer to do. But using the concepts we'll go over, you'll be able to train proper language models on larger dataset.**Example**: This model is to take a sequence like `[1, 2, 3, 4]` and return `[4, 3, 2, 1]`.1. Create the Model1. Prepare the Dataset1. Train the model using `Trainer`  Create the Model<jupyter_code># Create a Transformer model.
def tiny_transformer_lm(mode='train'):
  return trax.models.TransformerLM(  
          d_model=32, d_ff=128, n_layers=2, 
          vocab_size=32, mode=mode)<jupyter_output><empty_output><jupyter_text>Refer to [TransferLM in the API reference](https://trax-ml.readthedocs.io/en/latest/trax.models.htmltrax.models.transformer.TransformerLM) to understand each of its parameters and their default values. We have chosen to create a small model using these values for `d_model`, `d_ff`, and `n_layers` to be able to train the model more quickly on this simple task.![](https://storage.googleapis.com/ml-intro/t/untrained-transformer.png) Prepare the DatasetTrax models are trained on streams of data represented as python iterators. [`trax.data`](https://trax-ml.readthedocs.io/en/latest/trax.data.html) gives you the tools to construct your datapipeline. Trax also gives you readily available access to [TensorFlow Datasets](https://www.tensorflow.org/datasets).For this simple task, we will create a python generator. Every time we invoke it, it returns a batch of training examples.<jupyter_code>def reverse_ints_task(batch_size, length=4):
  while True:
    random_ints = m = np.random.randint(1, 31, (batch_size,length))
    source = random_ints

    target = np.flip(source, 1)

    zero = np.zeros([batch_size, 1], np.int32)
    x = np.concatenate([zero, source, zero, target], axis=1)

    loss_weights = np.concatenate([np.zeros((batch_size, length+2)),
                                    np.ones((batch_size, length))], axis=1)
    yield (x, x, loss_weights)  # Here inputs and targets are the same.

reverse_ints_inputs =  reverse_ints_task(16)<jupyter_output><empty_output><jupyter_text>This function prepares a dataset and returns one batch at a time. If we ask for a batch size of 8, for example, it returns the following:<jupyter_code>a = reverse_ints_task(8)
sequence_batch, _ , masks = next(a)
sequence_batch<jupyter_output><empty_output><jupyter_text>You can see that each example starts with 0, then a list of integers, then another 0, then the reverse of the list of integers. The function will give us as many examples and batches as we request.In addition to the example, the generator returns a mask vector. During the training process, the model is challenged to predict the tokens hidden by the mask (which have a value of 1 associated with that position. So for example, if the first element in the batch is the following vector:0567808765 And the associated mask vector for this example is:0000001111 Then the model will only be presented with the following prefix items, and it has to predict the rest:056780___ _ It's important here to note that while `5, 6, 7, 8` constitute the input sequence, the **zeros** serve a different purpose. We are using them as special tokens to delimit where the source sequence begins and ends. With this, we now have a method that streams the dataset in addition to the method that creates the model.![](https://storage.googleapis.com/ml-intro/t/untrained-transformer-and-dataset.png) Train the modelTrax's [training](https://trax-ml.readthedocs.io/en/latest/notebooks/trax_intro.htmlSupervised-training) takes care of the training process. We hand it the model, define training and eval tasks, and create the training loop. We then start the training loop.<jupyter_code>from trax.supervised import training
from trax import layers as tl

# Training task.
train_task = training.TrainTask(
    labeled_data=reverse_ints_inputs,
    loss_layer=tl.CrossEntropyLoss(),
    optimizer=trax.optimizers.Adam(0.01),
    n_steps_per_checkpoint=500,
)


# Evaluaton task.
eval_task = training.EvalTask(
    labeled_data=reverse_ints_inputs,
    metrics=[tl.CrossEntropyLoss(), tl.Accuracy()],
    n_eval_batches=20  # For less variance in eval numbers.
)

output_dir = os.path.expanduser('~/train_dir/')
!rm -f ~/train_dir/model.pkl.gz  # Remove old model.

# Train tiny model with Loop.
training_loop = training.Loop(
    tiny_transformer_lm(),
    train_task,
    eval_tasks=[eval_task],
    output_dir=output_dir)

# run 1000 steps (batches)
training_loop.run(1000)<jupyter_output>Step      1: Ran 1 train steps in 17.93 secs
Step      1: train CrossEntropyLoss |  4.14618683
Step      1: eval  CrossEntropyLoss |  3.74931383
Step      1: eval          Accuracy |  0.03359375

Step    500: Ran 499 train steps in 23.67 secs
Step    500: train CrossEntropyLoss |  0.62780923
Step    500: eval  CrossEntropyLoss |  0.01693780
Step    500: eval          Accuracy |  0.99609375

Step   1000: Ran 500 train steps in 5.34 secs
Step   1000: train CrossEntropyLoss |  0.00926041
Step   1000: eval  CrossEntropyLoss |  0.00390428
Step   1000: eval          Accuracy |  0.99921875<jupyter_text>The Trainer is the third key component in this process that helps us arrive at the trained model.![](https://storage.googleapis.com/ml-intro/t/transformerLM-training.png) Make predictionsLet's take our newly minted model for a ride. To do that, we load it up, and use the handy `autoregressive_sample` method to feed it our input sequence and return the output sequence. These components now look like this:![](https://storage.googleapis.com/ml-intro/t/transformerLM-sampling-prediction.png)And this is the code to do just that:<jupyter_code>input = np.array([[0, 4, 6, 8, 10, 0]])

# Initialize model for inference.
predict_model = tiny_transformer_lm(mode='predict')
predict_signature = trax.shapes.ShapeDtype((1,1), dtype=np.int32)
predict_model.init_from_file(os.path.join(output_dir, "model.pkl.gz"),
                             weights_only=True, input_signature=predict_signature)

# Run the model
output = trax.supervised.decoding.autoregressive_sample(
    predict_model, input, temperature=0.0, max_length=4)

# Print the contents of output
print(output)<jupyter_output>[[10  8  6  4]]
{'fraction_non_alphanumeric': 0.13713759779107226, 'fraction_numerical': 0.03414634146341464, 'mean_word_length': 4.46579476861167, 'pattern_counts': {'":': 0, '<': 27, '<?xml version=': 0, '>': 28, 'https://': 14, 'lorem ipsum': 0, 'www.': 1, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 1, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '7', 'dir': 'jupyter-structured-clean-dedup', 'id': '134037', 'license': 'Apache-2.0', 'n_tokens_mistral': 4029, 'n_tokens_neox': 3387, 'n_words': 1167, 'path': '_notebooks/2021-07-09-concept-transformer-lm.ipynb', 'repo_name': 'recohut/notebook', 'size': '10870'}