使用序列到序列 Transformer 的英翻西翻譯
作者： fchollet
建立日期 2021/05/26
上次修改日期 2024/11/18
描述： 實作一個序列到序列 Transformer 並在機器翻譯任務上訓練它。
簡介
在這個範例中，我們將建構一個序列到序列 Transformer 模型，並在英翻西機器翻譯任務上訓練它。
您將學習如何
- 使用 Keras
TextVectorization層來向量化文字。 - 實作
TransformerEncoder層、TransformerDecoder層和PositionalEmbedding層。 - 準備資料以訓練序列到序列模型。
- 使用訓練好的模型來生成從未見過的輸入句子的翻譯（序列到序列推論）。
這裡的功能程式碼改編自書籍 Deep Learning with Python, Second Edition （第 11 章：用於文字的深度學習）。目前範例相當簡陋，因此若要詳細說明每個建構區塊如何運作，以及 Transformer 背後的理論，我建議閱讀本書。
設定
# 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
下載資料
我們將使用 Anki 提供的英翻西翻譯資料集。讓我們下載它
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"
剖析資料
每行包含一個英文句子及其對應的西班牙文句子。英文句子是來源序列，西班牙文句子是目標序列。我們在西班牙文句子前面加上符號 "[start]"，並在後面加上符號 "[end]"。
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))
以下是我們的句子對範例
for _ in range(5):
print(random.choice(text_pairs))
('The trouble is that we have nowhere to stay tonight.', '[start] El problema es que no tenemos donde quedarnos esta noche. [end]')
("I want to help you, but I can't.", '[start] Quiero ayudarte, pero no puedo. [end]')
('I can help.', '[start] Yo puedo ayudar. [end]')
('Tom fed his dog table scraps.', '[start] Tom alimentó a su perro con sobras de la mesa. [end]')
('Tom never eats junk food.', '[start] Tom nunca come comida chatarra. [end]')
現在，讓我們將句子對分割為訓練集、驗證集和測試集。
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")
118964 total pairs
83276 training pairs
17844 validation pairs
17844 test pairs
向量化文字資料
我們將使用兩個 TextVectorization 層的實例來向量化文字資料（一個用於英文，一個用於西班牙文），也就是說，將原始字串轉換為整數序列，其中每個整數代表詞彙中單字的索引。
英文層將使用預設的字串標準化（去除標點符號字元）和分割方案（以空白分割），而西班牙文層將使用自訂標準化，其中我們將字元 "¿" 新增至要去除的標點符號字元集。
注意：在生產級機器翻譯模型中，我不建議去除任何語言中的標點符號字元。相反地，我建議將每個標點符號字元轉換為自己的符號，您可以使用自訂的 split 函數提供給 TextVectorization 層來達成此目的。
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)
接下來，我們將格式化我們的資料集。
在每個訓練步驟中，模型會尋求使用來源句子和目標單字 0 到 N 來預測目標單字 N+1（以及之後的單字）。
因此，訓練資料集會產生一個元組 (inputs, targets)，其中
inputs是一個具有鍵encoder_inputs和decoder_inputs的字典。encoder_inputs是向量化的來源句子，而encoder_inputs是「目前為止」的目標句子，也就是說，用於預測目標句子中單字 N+1（以及之後的單字）的單字 0 到 N。target是目標句子偏移一步：它提供目標句子中的下一個單字 — 模型會嘗試預測的單字。
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)
讓我們快速看一下序列形狀（我們有 64 對的批次，且所有序列長度為 20 步）
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}")
inputs["encoder_inputs"].shape: (64, 20)
inputs["decoder_inputs"].shape: (64, 20)
targets.shape: (64, 20)
建構模型
我們的序列到序列 Transformer 由鏈接在一起的 TransformerEncoder 和 TransformerDecoder 組成。為了讓模型了解單字順序，我們也使用 PositionalEmbedding 層。
來源序列將傳遞到 TransformerEncoder，它將產生新的表示法。然後，此新的表示法將與目前為止的目標序列（目標單字 0 到 N）一起傳遞到 TransformerDecoder。然後，TransformerDecoder 會尋求預測目標序列中的下一個單字（N+1 及之後的單字）。
使這成為可能的關鍵細節是因果遮罩（請參閱 TransformerDecoder 上的方法 get_causal_attention_mask()）。TransformerDecoder 會一次查看整個序列，因此我們必須確保它在預測符號 N+1 時僅使用來自目標符號 0 到 N 的資訊（否則，它可能會使用來自未來的資訊，這會導致模型無法在推論時使用）。
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):
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, mask=None):
inputs, encoder_outputs = inputs
causal_mask = self.get_causal_attention_mask(inputs)
if mask is None:
inputs_padding_mask, encoder_outputs_padding_mask = None, None
else:
inputs_padding_mask, encoder_outputs_padding_mask = mask
attention_output_1 = self.attention_1(
query=inputs,
value=inputs,
key=inputs,
attention_mask=causal_mask,
query_mask=inputs_padding_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,
query_mask=inputs_padding_mask,
key_mask=encoder_outputs_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
接下來，我們組裝端到端模型。
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, encoder_outputs])
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)
transformer = keras.Model(
{"encoder_inputs": encoder_inputs, "decoder_inputs": decoder_inputs},
decoder_outputs,
name="transformer",
)
訓練模型
我們將使用準確度作為快速監控驗證資料訓練進度的方式。請注意，機器翻譯通常使用 BLEU 分數以及其他度量，而不是準確度。
這裡我們只訓練 1 個週期，但為了讓模型實際收斂，您應該至少訓練 30 個週期。
epochs = 1 # This should be at least 30 for convergence
transformer.summary()
transformer.compile(
"rmsprop",
loss=keras.losses.SparseCategoricalCrossentropy(ignore_class=0),
metrics=["accuracy"],
)
transformer.fit(train_ds, epochs=epochs, validation_data=val_ds)
Model: "transformer"
┏━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━┓ ┃ Layer (type) ┃ Output Shape ┃ Param # ┃ Connected to ┃ ┡━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━┩ │ encoder_inputs │ (None, None) │ 0 │ - │ │ (InputLayer) │ │ │ │ ├─────────────────────┼───────────────────┼────────────┼───────────────────┤ │ decoder_inputs │ (None, None) │ 0 │ - │ │ (InputLayer) │ │ │ │ ├─────────────────────┼───────────────────┼────────────┼───────────────────┤ │ positional_embeddi… │ (None, None, 256) │ 3,845,120 │ encoder_inputs[0… │ │ (PositionalEmbeddi… │ │ │ │ ├─────────────────────┼───────────────────┼────────────┼───────────────────┤ │ not_equal │ (None, None) │ 0 │ encoder_inputs[0… │ │ (NotEqual) │ │ │ │ ├─────────────────────┼───────────────────┼────────────┼───────────────────┤ │ positional_embeddi… │ (None, None, 256) │ 3,845,120 │ decoder_inputs[0… │ │ (PositionalEmbeddi… │ │ │ │ ├─────────────────────┼───────────────────┼────────────┼───────────────────┤ │ transformer_encoder │ (None, None, 256) │ 3,155,456 │ positional_embed… │ │ (TransformerEncode… │ │ │ not_equal[0][0] │ ├─────────────────────┼───────────────────┼────────────┼───────────────────┤ │ not_equal_1 │ (None, None) │ 0 │ decoder_inputs[0… │ │ (NotEqual) │ │ │ │ ├─────────────────────┼───────────────────┼────────────┼───────────────────┤ │ transformer_decoder │ (None, None, 256) │ 5,259,520 │ positional_embed… │ │ (TransformerDecode… │ │ │ transformer_enco… │ │ │ │ │ not_equal_1[0][0… │ │ │ │ │ not_equal[0][0] │ ├─────────────────────┼───────────────────┼────────────┼───────────────────┤ │ dropout_3 (Dropout) │ (None, None, 256) │ 0 │ transformer_deco… │ ├─────────────────────┼───────────────────┼────────────┼───────────────────┤ │ dense_4 (Dense) │ (None, None, │ 3,855,000 │ dropout_3[0][0] │ │ │ 15000) │ │ │ └─────────────────────┴───────────────────┴────────────┴───────────────────┘
Total params: 19,960,216 (76.14 MB)
Trainable params: 19,960,216 (76.14 MB)
Non-trainable params: 0 (0.00 B)
<keras.src.callbacks.history.History at 0x7ffae0753a60>
解碼測試句子
最後，讓我們示範如何翻譯全新的英文句子。我們只需將向量化的英文句子以及目標符號 "[start]" 輸入到模型中，然後重複產生下一個符號，直到我們命中符號 "[end]"。
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(
{
"encoder_inputs": tokenized_input_sentence,
"decoder_inputs": 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)
在 30 個週期之後，我們會得到如下結果
她把錢遞給他。[start] ella le pasó el dinero [end]
湯姆從未聽過瑪麗唱歌。[start] tom nunca ha oído cantar a mary [end]
也許她明天會來。[start] tal vez ella vendrá mañana [end]
我喜歡寫作。[start] me encanta escribir [end]
他的法語一點一點地在進步。[start] su francés va a [UNK] sólo un poco [end]
我的飯店叫我打電話給你。[start] mi hotel me dijo que te [UNK] [end]