title englishtospanish translation with kerasnlp abheesht sharmahttps github comabheesht17 date created 20220526 last modified 20221221 description use kerasnlp to train a sequencetosequence transformer model on the machine translation task accelerator gpu introduction kerasnlp provides building blocks for nlp model layers tokenizers metrics etc and makes it convenient to construct nlp pipelines in this example we ll use kerasnlp layers to build an encoderdecoder transformer model and train it on the englishtospanish machine translation task this example is based on the englishtospanish nmt examplehttps keras ioexamplesnlpneuralmachinetranslationwithtransformer by fchollethttps twitter comfchollet the original example is more lowlevel and implements layers from scratch whereas this example uses kerasnlp to show some more advanced approaches such as subword tokenization and using metrics to compute the quality of generated translations you ll learn how to tokenize text using kerasnlp tokenizers wordpiecetokenizer implement a sequencetosequence transformer model using kerasnlp s kerasnlp layers transformerencoder kerasnlp layers transformerdecoder and kerasnlp layers tokenandpositionembedding layers and train it use kerasnlp samplers to generate translations of unseen input sentences using the topp decoding strategy don t worry if you aren t familiar with kerasnlp this tutorial will start with the basics let s dive right in setup before we start implementing the pipeline let s import all the libraries we need shell pip install q rougescore pip install q githttps github comkerasteamkerasnlp git upgrade let s also define our parametershyperparameters downloading the data we ll be working with an englishtospanish translation dataset provided by ankihttps www manythings organki let s download it parsing the data each line contains an english sentence and its corresponding spanish sentence the english sentence is the source sequence and spanish one is the target sequence before adding the text to a list we convert it to lowercase here s what our sentence pairs look like now let s split the sentence pairs into a training set a validation set and a test set tokenizing the data we ll define two tokenizers one for the source language english and the other for the target language spanish we ll be using kerasnlp tokenizers wordpiecetokenizer to tokenize the text kerasnlp tokenizers wordpiecetokenizer takes a wordpiece vocabulary and has functions for tokenizing the text and detokenizing sequences of tokens before we define the two tokenizers we first need to train them on the dataset we have the wordpiece tokenization algorithm is a subword tokenization algorithm training it on a corpus gives us a vocabulary of subwords a subword tokenizer is a compromise between word tokenizers word tokenizers need very large vocabularies for good coverage of input words and character tokenizers characters don t really encode meaning like words do luckily kerasnlp makes it very simple to train wordpiece on a corpus with the kerasnlp tokenizers computewordpiecevocabulary utility every vocabulary has a few special reserved tokens we have four such tokens pad padding token padding tokens are appended to the input sequence length when the input sequence length is shorter than the maximum sequence length unk unknown token start token that marks the start of the input sequence end token that marks the end of the input sequence let s see some tokens now let s define the tokenizers we will configure the tokenizers with the the vocabularies trained above let s try and tokenize a sample from our dataset to verify whether the text has been tokenized correctly we can also detokenize the list of tokens back to the original text format datasets next we ll format our datasets at each training step the model will seek to predict target words n1 and beyond using the source sentence and the target words 0 to n as such the training dataset will yield a tuple inputs targets where inputs is a dictionary with the keys encoderinputs and decoderinputs encoderinputs is the tokenized source sentence and decoderinputs is the target sentence so far that is to say the words 0 to n used to predict word n1 and beyond in the target sentence target is the target sentence offset by one step it provides the next words in the target sentence what the model will try to predict we will add special tokens start and end to the input spanish sentence after tokenizing the text we will also pad the input to a fixed length this can be easily done using kerasnlp layers startendpacker pad eng to maxsequencelength add special tokens start and end to spa and pad it as well let s take a quick look at the sequence shapes we have batches of 64 pairs and all sequences are 40 steps long building the model now let s move on to the exciting part defining our model we first need an embedding layer i e a vector for every token in our input sequence this embedding layer can be initialised randomly we also need a positional embedding layer which encodes the word order in the sequence the convention is to add these two embeddings kerasnlp has a kerasnlp layers tokenandpositionembedding layer which does all of the above steps for us our sequencetosequence transformer consists of a kerasnlp layers transformerencoder layer and a kerasnlp layers transformerdecoder layer chained together the source sequence will be passed to kerasnlp layers transformerencoder which will produce a new representation of it this new representation will then be passed to the kerasnlp layers transformerdecoder together with the target sequence so far target words 0 to n the kerasnlp layers transformerdecoder will then seek to predict the next words in the target sequence n1 and beyond a key detail that makes this possible is causal masking the kerasnlp layers transformerdecoder sees the entire sequence at once and thus we must make sure that it only uses information from target tokens 0 to n when predicting token n1 otherwise it could use information from the future which would result in a model that cannot be used at inference time causal masking is enabled by default in kerasnlp layers transformerdecoder we also need to mask the padding tokens pad for this we can set the maskzero argument of the kerasnlp layers tokenandpositionembedding layer to true this will then be propagated to all subsequent layers encoder decoder training our model we ll use accuracy as a quick way to monitor training progress on the validation data note that machine translation typically uses bleu scores as well as other metrics rather than accuracy however in order to use metrics like rouge bleu etc we will have decode the probabilities and generate the text text generation is computationally expensive and performing this during training is not recommended here we only train for 1 epoch but to get the model to actually converge you should train for at least 10 epochs decoding test sentences qualitative analysis finally let s demonstrate how to translate brand new english sentences we simply feed into the model the tokenized english sentence as well as the target token start the model outputs probabilities of the next token we then we repeatedly generated the next token conditioned on the tokens generated so far until we hit the token end for decoding we will use the kerasnlp samplers module from kerasnlp greedy decoding is a text decoding method which outputs the most likely next token at each time step i e the token with the highest probability tokenize the encoder input define a function that outputs the next token s probability given the input sequence ignore hidden states for now only needed for contrastive search build a prompt of length 40 with a start token and padding tokens evaluating our model quantitative analysis there are many metrics which are used for text generation tasks here to evaluate translations generated by our model let s compute the rouge1 and rouge2 scores essentially rougen is a score based on the number of common ngrams between the reference text and the generated text rouge1 and rouge2 use the number of common unigrams and bigrams respectively we will calculate the score over 30 test samples since decoding is an expensive process after 10 epochs the scores are as follows rouge1 rouge2 precision 0 568 0 374 recall 0 615 0 394 f1 score 0 579 0 381 introduction kerasnlp provides building blocks for nlp model layers tokenizers metrics etc and makes it convenient to construct nlp pipelines in this example we ll use kerasnlp layers to build an encoder decoder transformer model and train it on the english to spanish machine translation task this example is based on the english to spanish nmt example https keras io examples nlp neural_machine_translation_with_transformer by fchollet https twitter com fchollet the original example is more low level and implements layers from scratch whereas this example uses kerasnlp to show some more advanced approaches such as subword tokenization and using metrics to compute the quality of generated translations you ll learn how to tokenize text using keras_nlp tokenizers wordpiecetokenizer implement a sequence to sequence transformer model using kerasnlp s keras_nlp layers transformerencoder keras_nlp layers transformerdecoder and keras_nlp layers tokenandpositionembedding layers and train it use keras_nlp samplers to generate translations of unseen input sentences using the top p decoding strategy don t worry if you aren t familiar with kerasnlp this tutorial will start with the basics let s dive right in setup before we start implementing the pipeline let s import all the libraries we need shell pip install q rouge score pip install q git https github com keras team keras nlp git upgrade let s also define our parameters hyperparameters this should be at least 10 for convergence downloading the data we ll be working with an english to spanish translation dataset provided by anki https www manythings org anki let s download it parsing the data each line contains an english sentence and its corresponding spanish sentence the english sentence is the source sequence and spanish one is the target sequence before adding the text to a list we convert it to lowercase here s what our sentence pairs look like now let s split the sentence pairs into a training set a validation set and a test set tokenizing the data we ll define two tokenizers one for the source language english and the other for the target language spanish we ll be using keras_nlp tokenizers wordpiecetokenizer to tokenize the text keras_nlp tokenizers wordpiecetokenizer takes a wordpiece vocabulary and has functions for tokenizing the text and detokenizing sequences of tokens before we define the two tokenizers we first need to train them on the dataset we have the wordpiece tokenization algorithm is a subword tokenization algorithm training it on a corpus gives us a vocabulary of subwords a subword tokenizer is a compromise between word tokenizers word tokenizers need very large vocabularies for good coverage of input words and character tokenizers characters don t really encode meaning like words do luckily kerasnlp makes it very simple to train wordpiece on a corpus with the keras_nlp tokenizers compute_word_piece_vocabulary utility every vocabulary has a few special reserved tokens we have four such tokens pad padding token padding tokens are appended to the input sequence length when the input sequence length is shorter than the maximum sequence length unk unknown token start token that marks the start of the input sequence end token that marks the end of the input sequence let s see some tokens now let s define the tokenizers we will configure the tokenizers with the the vocabularies trained above let s try and tokenize a sample from our dataset to verify whether the text has been tokenized correctly we can also detokenize the list of tokens back to the original text format datasets next we ll format our datasets at each training step the model will seek to predict target words n 1 and beyond using the source sentence and the target words 0 to n as such the training dataset will yield a tuple inputs targets where inputs is a dictionary with the keys encoder_inputs and decoder_inputs encoder_inputs is the tokenized source sentence and decoder_inputs is the target sentence so far that is to say the words 0 to n used to predict word n 1 and beyond in the target sentence target is the target sentence offset by one step it provides the next words in the target sentence what the model will try to predict we will add special tokens start and end to the input spanish sentence after tokenizing the text we will also pad the input to a fixed length this can be easily done using keras_nlp layers startendpacker pad eng to max_sequence_length add special tokens start and end to spa and pad it as well let s take a quick look at the sequence shapes we have batches of 64 pairs and all sequences are 40 steps long building the model now let s move on to the exciting part defining our model we first need an embedding layer i e a vector for every token in our input sequence this embedding layer can be initialised randomly we also need a positional embedding layer which encodes the word order in the sequence the convention is to add these two embeddings kerasnlp has a keras_nlp layers tokenandpositionembedding layer which does all of the above steps for us our sequence to sequence transformer consists of a keras_nlp layers transformerencoder layer and a keras_nlp layers transformerdecoder layer chained together the source sequence will be passed to keras_nlp layers transformerencoder which will produce a new representation of it this new representation will then be passed to the keras_nlp layers transformerdecoder together with the target sequence so far target words 0 to n the keras_nlp layers transformerdecoder will then seek to predict the next words in the target sequence n 1 and beyond a key detail that makes this possible is causal masking the keras_nlp layers transformerdecoder sees the entire sequence at once and thus we must make sure that it only uses information from target tokens 0 to n when predicting token n 1 otherwise it could use information from the future which would result in a model that cannot be used at inference time causal masking is enabled by default in keras_nlp layers transformerdecoder we also need to mask the padding tokens pad for this we can set the mask_zero argument of the keras_nlp layers tokenandpositionembedding layer to true this will then be propagated to all subsequent layers encoder decoder training our model we ll use accuracy as a quick way to monitor training progress on the validation data note that machine translation typically uses bleu scores as well as other metrics rather than accuracy however in order to use metrics like rouge bleu etc we will have decode the probabilities and generate the text text generation is computationally expensive and performing this during training is not recommended here we only train for 1 epoch but to get the model to actually converge you should train for at least 10 epochs decoding test sentences qualitative analysis finally let s demonstrate how to translate brand new english sentences we simply feed into the model the tokenized english sentence as well as the target token start the model outputs probabilities of the next token we then we repeatedly generated the next token conditioned on the tokens generated so far until we hit the token end for decoding we will use the keras_nlp samplers module from kerasnlp greedy decoding is a text decoding method which outputs the most likely next token at each time step i e the token with the highest probability tokenize the encoder input define a function that outputs the next token s probability given the input sequence ignore hidden states for now only needed for contrastive search build a prompt of length 40 with a start token and padding tokens start sampling after start token evaluating our model quantitative analysis there are many metrics which are used for text generation tasks here to evaluate translations generated by our model let s compute the rouge 1 and rouge 2 scores essentially rouge n is a score based on the number of common n grams between the reference text and the generated text rouge 1 and rouge 2 use the number of common unigrams and bigrams respectively we will calculate the score over 30 test samples since decoding is an expensive process after 10 epochs the scores are as follows rouge 1 rouge 2 precision 0 568 0 374 recall 0 615 0 394 f1 score 0 579 0 381
import keras_nlp
import pathlib
import random

import keras
from keras import ops

import tensorflow.data as tf_data
from tensorflow_text.tools.wordpiece_vocab import (
    bert_vocab_from_dataset as bert_vocab,
)



BATCH_SIZE = 64
EPOCHS = 1  
MAX_SEQUENCE_LENGTH = 40
ENG_VOCAB_SIZE = 15000
SPA_VOCAB_SIZE = 15000

EMBED_DIM = 256
INTERMEDIATE_DIM = 2048
NUM_HEADS = 8



text_file = keras.utils.get_file(
    fname="spa-eng.zip",
    origin="http://storage.googleapis.com/download.tensorflow.org/data/spa-eng.zip",
    extract=True,
)
text_file = pathlib.Path(text_file).parent / "spa-eng" / "spa.txt"



with open(text_file) as f:
    lines = f.read().split("\n")[:-1]
text_pairs = []
for line in lines:
    eng, spa = line.split("\t")
    eng = eng.lower()
    spa = spa.lower()
    text_pairs.append((eng, spa))



for _ in range(5):
    print(random.choice(text_pairs))



random.shuffle(text_pairs)
num_val_samples = int(0.15 * len(text_pairs))
num_train_samples = len(text_pairs) - 2 * num_val_samples
train_pairs = text_pairs[:num_train_samples]
val_pairs = text_pairs[num_train_samples : num_train_samples + num_val_samples]
test_pairs = text_pairs[num_train_samples + num_val_samples :]

print(f"{len(text_pairs)} total pairs")
print(f"{len(train_pairs)} training pairs")
print(f"{len(val_pairs)} validation pairs")
print(f"{len(test_pairs)} test pairs")





def train_word_piece(text_samples, vocab_size, reserved_tokens):
    word_piece_ds = tf_data.Dataset.from_tensor_slices(text_samples)
    vocab = keras_nlp.tokenizers.compute_word_piece_vocabulary(
        word_piece_ds.batch(1000).prefetch(2),
        vocabulary_size=vocab_size,
        reserved_tokens=reserved_tokens,
    )
    return vocab




reserved_tokens = ["[PAD]", "[UNK]", "[START]", "[END]"]

eng_samples = [text_pair[0] for text_pair in train_pairs]
eng_vocab = train_word_piece(eng_samples, ENG_VOCAB_SIZE, reserved_tokens)

spa_samples = [text_pair[1] for text_pair in train_pairs]
spa_vocab = train_word_piece(spa_samples, SPA_VOCAB_SIZE, reserved_tokens)



print("English Tokens: ", eng_vocab[100:110])
print("Spanish Tokens: ", spa_vocab[100:110])



eng_tokenizer = keras_nlp.tokenizers.WordPieceTokenizer(
    vocabulary=eng_vocab, lowercase=False
)
spa_tokenizer = keras_nlp.tokenizers.WordPieceTokenizer(
    vocabulary=spa_vocab, lowercase=False
)



eng_input_ex = text_pairs[0][0]
eng_tokens_ex = eng_tokenizer.tokenize(eng_input_ex)
print("English sentence: ", eng_input_ex)
print("Tokens: ", eng_tokens_ex)
print(
    "Recovered text after detokenizing: ",
    eng_tokenizer.detokenize(eng_tokens_ex),
)

print()

spa_input_ex = text_pairs[0][1]
spa_tokens_ex = spa_tokenizer.tokenize(spa_input_ex)
print("Spanish sentence: ", spa_input_ex)
print("Tokens: ", spa_tokens_ex)
print(
    "Recovered text after detokenizing: ",
    spa_tokenizer.detokenize(spa_tokens_ex),
)




def preprocess_batch(eng, spa):
    batch_size = ops.shape(spa)[0]

    eng = eng_tokenizer(eng)
    spa = spa_tokenizer(spa)

    
    eng_start_end_packer = keras_nlp.layers.StartEndPacker(
        sequence_length=MAX_SEQUENCE_LENGTH,
        pad_value=eng_tokenizer.token_to_id("[PAD]"),
    )
    eng = eng_start_end_packer(eng)

    
    spa_start_end_packer = keras_nlp.layers.StartEndPacker(
        sequence_length=MAX_SEQUENCE_LENGTH + 1,
        start_value=spa_tokenizer.token_to_id("[START]"),
        end_value=spa_tokenizer.token_to_id("[END]"),
        pad_value=spa_tokenizer.token_to_id("[PAD]"),
    )
    spa = spa_start_end_packer(spa)

    return (
        {
            "encoder_inputs": eng,
            "decoder_inputs": spa[:, :-1],
        },
        spa[:, 1:],
    )


def make_dataset(pairs):
    eng_texts, spa_texts = zip(*pairs)
    eng_texts = list(eng_texts)
    spa_texts = list(spa_texts)
    dataset = tf_data.Dataset.from_tensor_slices((eng_texts, spa_texts))
    dataset = dataset.batch(BATCH_SIZE)
    dataset = dataset.map(preprocess_batch, num_parallel_calls=tf_data.AUTOTUNE)
    return dataset.shuffle(2048).prefetch(16).cache()


train_ds = make_dataset(train_pairs)
val_ds = make_dataset(val_pairs)



for inputs, targets in train_ds.take(1):
    print(f'inputs["encoder_inputs"].shape: {inputs["encoder_inputs"].shape}')
    print(f'inputs["decoder_inputs"].shape: {inputs["decoder_inputs"].shape}')
    print(f"targets.shape: {targets.shape}")





encoder_inputs = keras.Input(
    shape=(None,), dtype="int64", name="encoder_inputs"
)

x = keras_nlp.layers.TokenAndPositionEmbedding(
    vocabulary_size=ENG_VOCAB_SIZE,
    sequence_length=MAX_SEQUENCE_LENGTH,
    embedding_dim=EMBED_DIM,
)(encoder_inputs)

encoder_outputs = keras_nlp.layers.TransformerEncoder(
    intermediate_dim=INTERMEDIATE_DIM, num_heads=NUM_HEADS
)(inputs=x)
encoder = keras.Model(encoder_inputs, encoder_outputs)



decoder_inputs = keras.Input(
    shape=(None,), dtype="int64", name="decoder_inputs"
)
encoded_seq_inputs = keras.Input(
    shape=(None, EMBED_DIM), name="decoder_state_inputs"
)

x = keras_nlp.layers.TokenAndPositionEmbedding(
    vocabulary_size=SPA_VOCAB_SIZE,
    sequence_length=MAX_SEQUENCE_LENGTH,
    embedding_dim=EMBED_DIM,
)(decoder_inputs)

x = keras_nlp.layers.TransformerDecoder(
    intermediate_dim=INTERMEDIATE_DIM, num_heads=NUM_HEADS
)(decoder_sequence=x, encoder_sequence=encoded_seq_inputs)
x = keras.layers.Dropout(0.5)(x)
decoder_outputs = keras.layers.Dense(SPA_VOCAB_SIZE, activation="softmax")(x)
decoder = keras.Model(
    [
        decoder_inputs,
        encoded_seq_inputs,
    ],
    decoder_outputs,
)
decoder_outputs = decoder([decoder_inputs, encoder_outputs])

transformer = keras.Model(
    [encoder_inputs, decoder_inputs],
    decoder_outputs,
    name="transformer",
)



transformer.summary()
transformer.compile(
    "rmsprop", loss="sparse_categorical_crossentropy", metrics=["accuracy"]
)
transformer.fit(train_ds, epochs=EPOCHS, validation_data=val_ds)




def decode_sequences(input_sentences):
    batch_size = 1

    
    encoder_input_tokens = ops.convert_to_tensor(eng_tokenizer(input_sentences))
    if len(encoder_input_tokens[0]) < MAX_SEQUENCE_LENGTH:
        pads = ops.full(
            (1, MAX_SEQUENCE_LENGTH - len(encoder_input_tokens[0])), 0
        )
        encoder_input_tokens = ops.concatenate([encoder_input_tokens, pads], 1)

    
    
    def next(prompt, cache, index):
        logits = transformer([encoder_input_tokens, prompt])[:, index - 1, :]
        
        hidden_states = None
        return logits, hidden_states, cache

    
    length = 40
    start = ops.full((batch_size, 1), spa_tokenizer.token_to_id("[START]"))
    pad = ops.full((batch_size, length - 1), spa_tokenizer.token_to_id("[PAD]"))
    prompt = ops.concatenate((start, pad), axis=-1)

    generated_tokens = keras_nlp.samplers.GreedySampler()(
        next,
        prompt,
        end_token_id=spa_tokenizer.token_to_id("[END]"),
        index=1,  
    )
    generated_sentences = spa_tokenizer.detokenize(generated_tokens)
    return generated_sentences


test_eng_texts = [pair[0] for pair in test_pairs]
for i in range(2):
    input_sentence = random.choice(test_eng_texts)
    translated = decode_sequences([input_sentence])
    translated = translated.numpy()[0].decode("utf-8")
    translated = (
        translated.replace("[PAD]", "")
        .replace("[START]", "")
        .replace("[END]", "")
        .strip()
    )
    print(f"** Example {i} **")
    print(input_sentence)
    print(translated)
    print()



rouge_1 = keras_nlp.metrics.RougeN(order=1)
rouge_2 = keras_nlp.metrics.RougeN(order=2)

for test_pair in test_pairs[:30]:
    input_sentence = test_pair[0]
    reference_sentence = test_pair[1]

    translated_sentence = decode_sequences([input_sentence])
    translated_sentence = translated_sentence.numpy()[0].decode("utf-8")
    translated_sentence = (
        translated_sentence.replace("[PAD]", "")
        .replace("[START]", "")
        .replace("[END]", "")
        .strip()
    )

    rouge_1(reference_sentence, translated_sentence)
    rouge_2(reference_sentence, translated_sentence)

print("ROUGE-1 Score: ", rouge_1.result())
print("ROUGE-2 Score: ", rouge_2.result())