CINXE.COM

<!DOCTYPE html> <html lang="en"> <head> <meta charset="utf-8"> <meta name="viewport" content="width=device-width, initial-scale=1"> <meta name="description" content="Keras documentation"> <meta name="author" content="Keras Team"> <link rel="shortcut icon" href="https://keras.io/img/favicon.ico"> <link rel="canonical" href="https://keras.io/examples/nlp/neural_machine_translation_with_transformer/" />  <meta property="og:title" content="Keras documentation: English-to-Spanish translation with a sequence-to-sequence Transformer"> <meta property="og:image" content="https://keras.io/img/logo-k-keras-wb.png"> <meta name="twitter:title" content="Keras documentation: English-to-Spanish translation with a sequence-to-sequence Transformer"> <meta name="twitter:image" content="https://keras.io/img/k-keras-social.png"> <meta name="twitter:card" content="summary"> <title>English-to-Spanish translation with a sequence-to-sequence Transformer</title>  <link href="/css/bootstrap.min.css" rel="stylesheet">  <link href="https://fonts.googleapis.com/css2?family=Open+Sans:wght@400;600;700;800&display=swap" rel="stylesheet">  <link href="/css/docs.css" rel="stylesheet"> <link href="/css/monokai.css" rel="stylesheet">  <script>(function(w,d,s,l,i){w[l]=w[l]||[];w[l].push({'gtm.start': new Date().getTime(),event:'gtm.js'});var f=d.getElementsByTagName(s)[0], j=d.createElement(s),dl=l!='dataLayer'?'&l='+l:'';j.async=true;j.src= 'https://www.googletagmanager.com/gtm.js?id='+i+dl;f.parentNode.insertBefore(j,f); })(window,document,'script','dataLayer','GTM-5DNGF4N'); </script> <script> (function(i,s,o,g,r,a,m){i['GoogleAnalyticsObject']=r;i[r]=i[r]||function(){ (i[r].q=i[r].q||[]).push(arguments)},i[r].l=1*new Date();a=s.createElement(o), m=s.getElementsByTagName(o)[0];a.async=1;a.src=g;m.parentNode.insertBefore(a,m) })(window,document,'script','https://www.google-analytics.com/analytics.js','ga'); ga('create', 'UA-175165319-128', 'auto'); ga('send', 'pageview'); </script>  <script async defer src="https://buttons.github.io/buttons.js"></script> </head> <body>  <noscript><iframe src="https://www.googletagmanager.com/ns.html?id=GTM-5DNGF4N" height="0" width="0" style="display:none;visibility:hidden"></iframe></noscript>  <div class='k-page'> <div class="k-nav" id="nav-menu"> <a href='/'><img src='/img/logo-small.png' class='logo-small' /></a> <div class="nav flex-column nav-pills" role="tablist" aria-orientation="vertical"> <a class="nav-link" href="/about/" role="tab" aria-selected="">About Keras</a> <a class="nav-link" href="/getting_started/" role="tab" aria-selected="">Getting started</a> <a class="nav-link" href="/guides/" role="tab" aria-selected="">Developer guides</a> <a class="nav-link active" href="/examples/" role="tab" aria-selected="">Code examples</a> <a class="nav-sublink" href="/examples/vision/">Computer Vision</a> <a class="nav-sublink active" href="/examples/nlp/">Natural Language Processing</a> <a class="nav-sublink2" href="/examples/nlp/text_classification_from_scratch/">Text classification from scratch</a> <a class="nav-sublink2" href="/examples/nlp/active_learning_review_classification/">Review Classification using Active Learning</a> <a class="nav-sublink2" href="/examples/nlp/fnet_classification_with_keras_hub/">Text Classification using FNet</a> <a class="nav-sublink2" href="/examples/nlp/multi_label_classification/">Large-scale multi-label text classification</a> <a class="nav-sublink2" href="/examples/nlp/text_classification_with_transformer/">Text classification with Transformer</a> <a class="nav-sublink2" href="/examples/nlp/text_classification_with_switch_transformer/">Text classification with Switch Transformer</a> <a class="nav-sublink2" href="/examples/nlp/tweet-classification-using-tfdf/">Text classification using Decision Forests and pretrained embeddings</a> <a class="nav-sublink2" href="/examples/nlp/pretrained_word_embeddings/">Using pre-trained word embeddings</a> <a class="nav-sublink2" href="/examples/nlp/bidirectional_lstm_imdb/">Bidirectional LSTM on IMDB</a> <a class="nav-sublink2" href="/examples/nlp/data_parallel_training_with_keras_hub/">Data Parallel Training with KerasHub and tf.distribute</a> <a class="nav-sublink2" href="/examples/nlp/neural_machine_translation_with_keras_hub/">English-to-Spanish translation with KerasHub</a> <a class="nav-sublink2 active" href="/examples/nlp/neural_machine_translation_with_transformer/">English-to-Spanish translation with a sequence-to-sequence Transformer</a> <a class="nav-sublink2" href="/examples/nlp/lstm_seq2seq/">Character-level recurrent sequence-to-sequence model</a> <a class="nav-sublink2" href="/examples/nlp/multimodal_entailment/">Multimodal entailment</a> <a class="nav-sublink2" href="/examples/nlp/ner_transformers/">Named Entity Recognition using Transformers</a> <a class="nav-sublink2" href="/examples/nlp/text_extraction_with_bert/">Text Extraction with BERT</a> <a class="nav-sublink2" href="/examples/nlp/addition_rnn/">Sequence to sequence learning for performing number addition</a> <a class="nav-sublink2" href="/examples/nlp/semantic_similarity_with_keras_hub/">Semantic Similarity with KerasHub</a> <a class="nav-sublink2" href="/examples/nlp/semantic_similarity_with_bert/">Semantic Similarity with BERT</a> <a class="nav-sublink2" href="/examples/nlp/sentence_embeddings_with_sbert/">Sentence embeddings using Siamese RoBERTa-networks</a> <a class="nav-sublink2" href="/examples/nlp/masked_language_modeling/">End-to-end Masked Language Modeling with BERT</a> <a class="nav-sublink2" href="/examples/nlp/abstractive_summarization_with_bart/">Abstractive Text Summarization with BART</a> <a class="nav-sublink2" href="/examples/nlp/pretraining_BERT/">Pretraining BERT with Hugging Face Transformers</a> <a class="nav-sublink2" href="/examples/nlp/parameter_efficient_finetuning_of_gpt2_with_lora/">Parameter-efficient fine-tuning of GPT-2 with LoRA</a> <a class="nav-sublink2" href="/examples/nlp/mlm_training_tpus/">Training a language model from scratch with 🤗 Transformers and TPUs</a> <a class="nav-sublink2" href="/examples/nlp/multiple_choice_task_with_transfer_learning/">MultipleChoice Task with Transfer Learning</a> <a class="nav-sublink2" href="/examples/nlp/question_answering/">Question Answering with Hugging Face Transformers</a> <a class="nav-sublink2" href="/examples/nlp/t5_hf_summarization/">Abstractive Summarization with Hugging Face Transformers</a> <a class="nav-sublink" href="/examples/structured_data/">Structured Data</a> <a class="nav-sublink" href="/examples/timeseries/">Timeseries</a> <a class="nav-sublink" href="/examples/generative/">Generative Deep Learning</a> <a class="nav-sublink" href="/examples/audio/">Audio Data</a> <a class="nav-sublink" href="/examples/rl/">Reinforcement Learning</a> <a class="nav-sublink" href="/examples/graph/">Graph Data</a> <a class="nav-sublink" href="/examples/keras_recipes/">Quick Keras Recipes</a> <a class="nav-link" href="/api/" role="tab" aria-selected="">Keras 3 API documentation</a> <a class="nav-link" href="/2.18/api/" role="tab" aria-selected="">Keras 2 API documentation</a> <a class="nav-link" href="/keras_tuner/" role="tab" aria-selected="">KerasTuner: Hyperparam Tuning</a> <a class="nav-link" href="/keras_hub/" role="tab" aria-selected="">KerasHub: Pretrained Models</a> </div> </div> <div class='k-main'> <div class='k-main-top'> <script> function displayDropdownMenu() { e = document.getElementById("nav-menu"); if (e.style.display == "block") { e.style.display = "none"; } else { e.style.display = "block"; document.getElementById("dropdown-nav").style.display = "block"; } } function resetMobileUI() { if (window.innerWidth <= 840) { document.getElementById("nav-menu").style.display = "none"; document.getElementById("dropdown-nav").style.display = "block"; } else { document.getElementById("nav-menu").style.display = "block"; document.getElementById("dropdown-nav").style.display = "none"; } var navmenu = document.getElementById("nav-menu"); var menuheight = navmenu.clientHeight; var kmain = document.getElementById("k-main-id"); kmain.style.minHeight = (menuheight + 100) + 'px'; } window.onresize = resetMobileUI; window.addEventListener("load", (event) => { resetMobileUI() }); </script> <div id='dropdown-nav' onclick="displayDropdownMenu();"> <svg viewBox="-20 -20 120 120" width="60" height="60"> <rect width="100" height="20"></rect> <rect y="30" width="100" height="20"></rect> <rect y="60" width="100" height="20"></rect> </svg> </div> <form class="bd-search d-flex align-items-center k-search-form" id="search-form"> <input type="search" class="k-search-input" id="search-input" placeholder="Search Keras documentation..." aria-label="Search Keras documentation..." autocomplete="off"> <button class="k-search-btn"> <svg width="13" height="13" viewBox="0 0 13 13"><title>search</title><path d="m4.8495 7.8226c0.82666 0 1.5262-0.29146 2.0985-0.87438 0.57232-0.58292 0.86378-1.2877 0.87438-2.1144 0.010599-0.82666-0.28086-1.5262-0.87438-2.0985-0.59352-0.57232-1.293-0.86378-2.0985-0.87438-0.8055-0.010599-1.5103 0.28086-2.1144 0.87438-0.60414 0.59352-0.8956 1.293-0.87438 2.0985 0.021197 0.8055 0.31266 1.5103 0.87438 2.1144 0.56172 0.60414 1.2665 0.8956 2.1144 0.87438zm4.4695 0.2115 3.681 3.6819-1.259 1.284-3.6817-3.7 0.0019784-0.69479-0.090043-0.098846c-0.87973 0.76087-1.92 1.1413-3.1207 1.1413-1.3553 0-2.5025-0.46363-3.4417-1.3909s-1.4088-2.0686-1.4088-3.4239c0-1.3553 0.4696-2.4966 1.4088-3.4239 0.9392-0.92727 2.0864-1.3969 3.4417-1.4088 1.3553-0.011889 2.4906 0.45771 3.406 1.4088 0.9154 0.95107 1.379 2.0924 1.3909 3.4239 0 1.2126-0.38043 2.2588-1.1413 3.1385l0.098834 0.090049z"></path></svg> </button> </form> <script> var form = document.getElementById('search-form'); form.onsubmit = function(e) { e.preventDefault(); var query = document.getElementById('search-input').value; window.location.href = '/search.html?query=' + query; return False } </script> </div> <div class='k-main-inner' id='k-main-id'> <div class='k-location-slug'> ► <a href='/examples/'>Code examples</a> / <a href='/examples/nlp/'>Natural Language Processing</a> / English-to-Spanish translation with a sequence-to-sequence Transformer </div> <div class='k-content'> <h1 id="englishtospanish-translation-with-a-sequencetosequence-transformer">English-to-Spanish translation with a sequence-to-sequence Transformer</h1> Author: <a href="https://twitter.com/fchollet">fchollet</a> Date created: 2021/05/26 Last modified: 2023/02/25 Description: Implementing a sequence-to-sequence Transformer and training it on a machine translation task. <div class='example_version_banner keras_3'>ⓘ This example uses Keras 3</div> <img class="k-inline-icon" src="https://colab.research.google.com/img/colab_favicon.ico"/> <a href="https://colab.research.google.com/github/keras-team/keras-io/blob/master/examples/nlp/ipynb/neural_machine_translation_with_transformer.ipynb">View in Colab</a> •<img class="k-inline-icon" src="https://github.com/favicon.ico"/> <a href="https://github.com/keras-team/keras-io/blob/master/examples/nlp/neural_machine_translation_with_transformer.py">GitHub source</a> <hr /> <h2 id="introduction">Introduction</h2> In this example, we'll build a sequence-to-sequence Transformer model, which we'll train on an English-to-Spanish machine translation task. You'll learn how to: <ul> <li>Vectorize text using the Keras <code>TextVectorization</code> layer.</li> <li>Implement a <code>TransformerEncoder</code> layer, a <code>TransformerDecoder</code> layer, and a <code>PositionalEmbedding</code> layer.</li> <li>Prepare data for training a sequence-to-sequence model.</li> <li>Use the trained model to generate translations of never-seen-before input sentences (sequence-to-sequence inference).</li> </ul> The code featured here is adapted from the book <a href="https://www.manning.com/books/deep-learning-with-python-second-edition">Deep Learning with Python, Second Edition</a> (chapter 11: Deep learning for text). The present example is fairly barebones, so for detailed explanations of how each building block works, as well as the theory behind Transformers, I recommend reading the book. <hr /> <h2 id="setup">Setup</h2> <div class="codehilite"><pre><code># We set the backend to TensorFlow. The code works with # both `tensorflow` and `torch`. It does not work with JAX # due to the behavior of `jax.numpy.tile` in a jit scope # (used in `TransformerDecoder.get_causal_attention_mask()`: # `tile` in JAX does not support a dynamic `reps` argument. # You can make the code work in JAX by wrapping the # inside of the `get_causal_attention_mask` method in # a decorator to prevent jit compilation: # `with jax.ensure_compile_time_eval():`. import os os.environ["KERAS_BACKEND"] = "tensorflow" import pathlib import random import string import re import numpy as np import tensorflow.data as tf_data import tensorflow.strings as tf_strings import keras from keras import layers from keras import ops from keras.layers import TextVectorization </code></pre></div> <hr /> <h2 id="downloading-the-data">Downloading the data</h2> We'll be working with an English-to-Spanish translation dataset provided by <a href="https://www.manythings.org/anki/">Anki</a>. Let's download it: <div class="codehilite"><pre><code>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" </code></pre></div> <hr /> <h2 id="parsing-the-data">Parsing the data</h2> 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. We prepend the token <code>"[start]"</code> and we append the token <code>"[end]"</code> to the Spanish sentence. <div class="codehilite"><pre><code>with open(text_file) as f: lines = f.read().split("\n")[:-1] text_pairs = [] for line in lines: eng, spa = line.split("\t") spa = "[start] " + spa + " [end]" text_pairs.append((eng, spa)) </code></pre></div> Here's what our sentence pairs look like: <div class="codehilite"><pre><code>for _ in range(5): print(random.choice(text_pairs)) </code></pre></div> <div class="k-default-codeblock"> <div class="codehilite"><pre><code>("On Saturday nights, it's difficult to find parking around here.", '[start] Los sábados por la noche es difícil encontrar aparcamiento por aquí. [end]') ('I was the worst student in the class.', '[start] Fui el peor estudiante en la clase. [end]') ('There is nothing to do today.', '[start] No hay nada que hacer hoy. [end]') ('The twins do resemble each other.', '[start] Los gemelos se parecen mutuamente. [end]') ('They found Tom in the crowd.', '[start] Encontraron a Tom entre la multitud. [end]') </code></pre></div> </div> Now, let's split the sentence pairs into a training set, a validation set, and a test set. <div class="codehilite"><pre><code>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") </code></pre></div> <div class="k-default-codeblock"> <div class="codehilite"><pre><code>118964 total pairs 83276 training pairs 17844 validation pairs 17844 test pairs </code></pre></div> </div> <hr /> <h2 id="vectorizing-the-text-data">Vectorizing the text data</h2> We'll use two instances of the <code>TextVectorization</code> layer to vectorize the text data (one for English and one for Spanish), that is to say, to turn the original strings into integer sequences where each integer represents the index of a word in a vocabulary. The English layer will use the default string standardization (strip punctuation characters) and splitting scheme (split on whitespace), while the Spanish layer will use a custom standardization, where we add the character <code>"¿"</code> to the set of punctuation characters to be stripped. Note: in a production-grade machine translation model, I would not recommend stripping the punctuation characters in either language. Instead, I would recommend turning each punctuation character into its own token, which you could achieve by providing a custom <code>split</code> function to the <code>TextVectorization</code> layer. <div class="codehilite"><pre><code>strip_chars = string.punctuation + "¿" strip_chars = strip_chars.replace("[", "") strip_chars = strip_chars.replace("]", "") vocab_size = 15000 sequence_length = 20 batch_size = 64 def custom_standardization(input_string): lowercase = tf_strings.lower(input_string) return tf_strings.regex_replace(lowercase, "[%s]" % re.escape(strip_chars), "") eng_vectorization = TextVectorization( max_tokens=vocab_size, output_mode="int", output_sequence_length=sequence_length, ) spa_vectorization = TextVectorization( max_tokens=vocab_size, output_mode="int", output_sequence_length=sequence_length + 1, standardize=custom_standardization, ) train_eng_texts = [pair[0] for pair in train_pairs] train_spa_texts = [pair[1] for pair in train_pairs] eng_vectorization.adapt(train_eng_texts) spa_vectorization.adapt(train_spa_texts) </code></pre></div> 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 <code>(inputs, targets)</code>, where: <ul> <li><code>inputs</code> is a dictionary with the keys <code>encoder_inputs</code> and <code>decoder_inputs</code>. <code>encoder_inputs</code> is the vectorized source sentence and <code>encoder_inputs</code> 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.</li> <li><code>target</code> is the target sentence offset by one step: it provides the next words in the target sentence – what the model will try to predict.</li> </ul> <div class="codehilite"><pre><code>def format_dataset(eng, spa): eng = eng_vectorization(eng) spa = spa_vectorization(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(format_dataset) return dataset.cache().shuffle(2048).prefetch(16) train_ds = make_dataset(train_pairs) val_ds = make_dataset(val_pairs) </code></pre></div> Let's take a quick look at the sequence shapes (we have batches of 64 pairs, and all sequences are 20 steps long): <div class="codehilite"><pre><code>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}") </code></pre></div> <div class="k-default-codeblock"> <div class="codehilite"><pre><code>inputs["encoder_inputs"].shape: (64, 20) inputs["decoder_inputs"].shape: (64, 20) targets.shape: (64, 20) </code></pre></div> </div> <hr /> <h2 id="building-the-model">Building the model</h2> Our sequence-to-sequence Transformer consists of a <code>TransformerEncoder</code> and a <code>TransformerDecoder</code> chained together. To make the model aware of word order, we also use a <code>PositionalEmbedding</code> layer. The source sequence will be pass to the <code>TransformerEncoder</code>, which will produce a new representation of it. This new representation will then be passed to the <code>TransformerDecoder</code>, together with the target sequence so far (target words 0 to N). The <code>TransformerDecoder</code> 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 (see method <code>get_causal_attention_mask()</code> on the <code>TransformerDecoder</code>). The <code>TransformerDecoder</code> sees the entire sequences 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). <div class="codehilite"><pre><code>import keras.ops as ops class TransformerEncoder(layers.Layer): def __init__(self, embed_dim, dense_dim, num_heads, **kwargs): super().__init__(**kwargs) self.embed_dim = embed_dim self.dense_dim = dense_dim self.num_heads = num_heads self.attention = layers.MultiHeadAttention( num_heads=num_heads, key_dim=embed_dim ) self.dense_proj = keras.Sequential( [ layers.Dense(dense_dim, activation="relu"), layers.Dense(embed_dim), ] ) self.layernorm_1 = layers.LayerNormalization() self.layernorm_2 = layers.LayerNormalization() self.supports_masking = True def call(self, inputs, mask=None): if mask is not None: padding_mask = ops.cast(mask[:, None, :], dtype="int32") else: padding_mask = None attention_output = self.attention( query=inputs, value=inputs, key=inputs, attention_mask=padding_mask ) proj_input = self.layernorm_1(inputs + attention_output) proj_output = self.dense_proj(proj_input) return self.layernorm_2(proj_input + proj_output) def get_config(self): config = super().get_config() config.update( { "embed_dim": self.embed_dim, "dense_dim": self.dense_dim, "num_heads": self.num_heads, } ) return config class PositionalEmbedding(layers.Layer): def __init__(self, sequence_length, vocab_size, embed_dim, **kwargs): super().__init__(**kwargs) self.token_embeddings = layers.Embedding( input_dim=vocab_size, output_dim=embed_dim ) self.position_embeddings = layers.Embedding( input_dim=sequence_length, output_dim=embed_dim ) self.sequence_length = sequence_length self.vocab_size = vocab_size self.embed_dim = embed_dim def call(self, inputs): length = ops.shape(inputs)[-1] positions = ops.arange(0, length, 1) embedded_tokens = self.token_embeddings(inputs) embedded_positions = self.position_embeddings(positions) return embedded_tokens + embedded_positions def compute_mask(self, inputs, mask=None): if mask is None: return None else: return ops.not_equal(inputs, 0) def get_config(self): config = super().get_config() config.update( { "sequence_length": self.sequence_length, "vocab_size": self.vocab_size, "embed_dim": self.embed_dim, } ) return config class TransformerDecoder(layers.Layer): def __init__(self, embed_dim, latent_dim, num_heads, **kwargs): super().__init__(**kwargs) self.embed_dim = embed_dim self.latent_dim = latent_dim self.num_heads = num_heads self.attention_1 = layers.MultiHeadAttention( num_heads=num_heads, key_dim=embed_dim ) self.attention_2 = layers.MultiHeadAttention( num_heads=num_heads, key_dim=embed_dim ) self.dense_proj = keras.Sequential( [ layers.Dense(latent_dim, activation="relu"), layers.Dense(embed_dim), ] ) self.layernorm_1 = layers.LayerNormalization() self.layernorm_2 = layers.LayerNormalization() self.layernorm_3 = layers.LayerNormalization() self.supports_masking = True def call(self, inputs, encoder_outputs, mask=None): causal_mask = self.get_causal_attention_mask(inputs) if mask is not None: padding_mask = ops.cast(mask[:, None, :], dtype="int32") padding_mask = ops.minimum(padding_mask, causal_mask) else: padding_mask = None attention_output_1 = self.attention_1( query=inputs, value=inputs, key=inputs, attention_mask=causal_mask ) out_1 = self.layernorm_1(inputs + attention_output_1) attention_output_2 = self.attention_2( query=out_1, value=encoder_outputs, key=encoder_outputs, attention_mask=padding_mask, ) out_2 = self.layernorm_2(out_1 + attention_output_2) proj_output = self.dense_proj(out_2) return self.layernorm_3(out_2 + proj_output) def get_causal_attention_mask(self, inputs): input_shape = ops.shape(inputs) batch_size, sequence_length = input_shape[0], input_shape[1] i = ops.arange(sequence_length)[:, None] j = ops.arange(sequence_length) mask = ops.cast(i >= j, dtype="int32") mask = ops.reshape(mask, (1, input_shape[1], input_shape[1])) mult = ops.concatenate( [ops.expand_dims(batch_size, -1), ops.convert_to_tensor([1, 1])], axis=0, ) return ops.tile(mask, mult) def get_config(self): config = super().get_config() config.update( { "embed_dim": self.embed_dim, "latent_dim": self.latent_dim, "num_heads": self.num_heads, } ) return config </code></pre></div> Next, we assemble the end-to-end model. <div class="codehilite"><pre><code>embed_dim = 256 latent_dim = 2048 num_heads = 8 encoder_inputs = keras.Input(shape=(None,), dtype="int64", name="encoder_inputs") x = PositionalEmbedding(sequence_length, vocab_size, embed_dim)(encoder_inputs) encoder_outputs = TransformerEncoder(embed_dim, latent_dim, num_heads)(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 = PositionalEmbedding(sequence_length, vocab_size, embed_dim)(decoder_inputs) x = TransformerDecoder(embed_dim, latent_dim, num_heads)(x, encoded_seq_inputs) x = layers.Dropout(0.5)(x) decoder_outputs = layers.Dense(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" ) </code></pre></div> <hr /> <h2 id="training-our-model">Training our model</h2> 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. Here we only train for 1 epoch, but to get the model to actually converge you should train for at least 30 epochs. <div class="codehilite"><pre><code>epochs = 1 # This should be at least 30 for convergence transformer.summary() transformer.compile( "rmsprop", loss="sparse_categorical_crossentropy", metrics=["accuracy"] ) transformer.fit(train_ds, epochs=epochs, validation_data=val_ds) </code></pre></div> <pre style="white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace">Model: "transformer" </pre> <pre style="white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace">┏━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━┓ ┃ Layer (type) ┃ Output Shape ┃ Param # ┃ Connected to ┃ ┡━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━┩ │ encoder_inputs │ (None, None) │ 0 │ - │ │ (InputLayer) │ │ │ │ ├─────────────────────┼───────────────────┼─────────┼──────────────────────┤ │ positional_embeddi… │ (None, None, 256) │ 3,845,… │ encoder_inputs[0][0] │ │ (PositionalEmbeddi… │ │ │ │ ├─────────────────────┼───────────────────┼─────────┼──────────────────────┤ │ decoder_inputs │ (None, None) │ 0 │ - │ │ (InputLayer) │ │ │ │ ├─────────────────────┼───────────────────┼─────────┼──────────────────────┤ │ transformer_encoder │ (None, None, 256) │ 3,155,… │ positional_embeddin… │ │ (TransformerEncode… │ │ │ │ ├─────────────────────┼───────────────────┼─────────┼──────────────────────┤ │ functional_5 │ (None, None, │ 12,959… │ decoder_inputs[0][0… │ │ (Functional) │ 15000) │ │ transformer_encoder… │ └─────────────────────┴───────────────────┴─────────┴──────────────────────┘ </pre> <pre style="white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace"> Total params: 19,960,216 (76.14 MB) </pre> <pre style="white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace"> Trainable params: 19,960,216 (76.14 MB) </pre> <pre style="white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace"> Non-trainable params: 0 (0.00 B) </pre> <div class="k-default-codeblock"> <div class="codehilite"><pre><code> 5/1302 [37m━━━━━━━━━━━━━━━━━━━━ 42s 33ms/step - accuracy: 0.3558 - loss: 8.3596 WARNING: All log messages before absl::InitializeLog() is called are written to STDERR I0000 00:00:1699484373.932513 76082 device_compiler.h:187] Compiled cluster using XLA! This line is logged at most once for the lifetime of the process. 1302/1302 ━━━━━━━━━━━━━━━━━━━━ 64s 39ms/step - accuracy: 0.7073 - loss: 2.2372 - val_accuracy: 0.7329 - val_loss: 1.6477 <keras.src.callbacks.history.History at 0x7ff611f21540> </code></pre></div> </div> <hr /> <h2 id="decoding-test-sentences">Decoding test sentences</h2> Finally, let's demonstrate how to translate brand new English sentences. We simply feed into the model the vectorized English sentence as well as the target token <code>"[start]"</code>, then we repeatedly generated the next token, until we hit the token <code>"[end]"</code>. <div class="codehilite"><pre><code>spa_vocab = spa_vectorization.get_vocabulary() spa_index_lookup = dict(zip(range(len(spa_vocab)), spa_vocab)) max_decoded_sentence_length = 20 def decode_sequence(input_sentence): tokenized_input_sentence = eng_vectorization([input_sentence]) decoded_sentence = "[start]" for i in range(max_decoded_sentence_length): tokenized_target_sentence = spa_vectorization([decoded_sentence])[:, :-1] predictions = transformer([tokenized_input_sentence, tokenized_target_sentence]) # ops.argmax(predictions[0, i, :]) is not a concrete value for jax here sampled_token_index = ops.convert_to_numpy( ops.argmax(predictions[0, i, :]) ).item(0) sampled_token = spa_index_lookup[sampled_token_index] decoded_sentence += " " + sampled_token if sampled_token == "[end]": break return decoded_sentence test_eng_texts = [pair[0] for pair in test_pairs] for _ in range(30): input_sentence = random.choice(test_eng_texts) translated = decode_sequence(input_sentence) </code></pre></div> After 30 epochs, we get results such as: <blockquote> She handed him the money. [start] ella le pasó el dinero [end] Tom has never heard Mary sing. [start] tom nunca ha oído cantar a mary [end] Perhaps she will come tomorrow. [start] tal vez ella vendrá mañana [end] I love to write. [start] me encanta escribir [end] His French is improving little by little. [start] su francés va a [UNK] sólo un poco [end] My hotel told me to call you. [start] mi hotel me dijo que te [UNK] [end] </blockquote> </div> <div class='k-outline'> <div class='k-outline-depth-1'> <a href='#englishtospanish-translation-with-a-sequencetosequence-transformer'>English-to-Spanish translation with a sequence-to-sequence Transformer</a> </div> <div class='k-outline-depth-2'> ◆ <a href='#introduction'>Introduction</a> </div> <div class='k-outline-depth-2'> ◆ <a href='#setup'>Setup</a> </div> <div class='k-outline-depth-2'> ◆ <a href='#downloading-the-data'>Downloading the data</a> </div> <div class='k-outline-depth-2'> ◆ <a href='#parsing-the-data'>Parsing the data</a> </div> <div class='k-outline-depth-2'> ◆ <a href='#vectorizing-the-text-data'>Vectorizing the text data</a> </div> <div class='k-outline-depth-2'> ◆ <a href='#building-the-model'>Building the model</a> </div> <div class='k-outline-depth-2'> ◆ <a href='#training-our-model'>Training our model</a> </div> <div class='k-outline-depth-2'> ◆ <a href='#decoding-test-sentences'>Decoding test sentences</a> </div> </div> </div> </div> </div> </body> <footer style="float: left; width: 100%; padding: 1em; border-top: solid 1px #bbb;"> <a href="https://policies.google.com/terms">Terms</a> | <a href="https://policies.google.com/privacy">Privacy</a> </footer> </html>