代码示例 / 自然语言处理 / 英文到西班牙文翻译的序列到序列变换器

英文到西班牙文翻译的序列到序列变换器

作者: fchollet
创建日期: 2021/05/26
最后修改: 2023/02/25
描述: 实现一个序列到序列的变换器,并在机器翻译任务上进行训练。

在Colab中查看 GitHub源代码


介绍

在这个示例中,我们将构建一个序列到序列的变换器模型,我们将对其进行训练以完成从英语到西班牙语的机器翻译任务。

您将学习如何:

  • 使用Keras的TextVectorization层对文本进行向量化。
  • 实现TransformerEncoder层、TransformerDecoder层和PositionalEmbedding层。
  • 准备数据以训练序列到序列模型。
  • 使用训练好的模型生成从未见过的输入句子的翻译(序列到序列推理)。

这里的代码改编自书籍 用Python进行深度学习(第二版) (第11章:文本的深度学习)。 本示例相当简陋,关于每个构建模块的工作原理以及变换器的理论,建议阅读该书。


设置

# 我们将后端设置为TensorFlow。该代码适用于
# `tensorflow`和`torch`。它不适用于JAX
# 由于`jax.numpy.tile`在jit作用域中的行为
# (用于`TransformerDecoder.get_causal_attention_mask()`:
# JAX中的`tile`不支持动态`reps`参数。
# 您可以通过在
# `get_causal_attention_mask`方法的内部包装一个
# 装饰器来使代码在JAX中运行,以防止jit编译:
# `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))
("周六晚上,这里很难找到停车位。", '[start] Los sábados por la noche es difícil encontrar aparcamiento por aquí. [end]')
('我在班里是最差的学生。', '[start] Fui el peor estudiante en la clase. [end]')
('今天没有什么可做的。', '[start] No hay nada que hacer hoy. [end]')
('这对双胞胎确实很像。', '[start] Los gemelos se parecen mutuamente. [end]')
('他们在拥挤的人群中找到了汤姆。', '[start] Encontraron a Tom entre la multitud. [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)} 总对")
print(f"{len(train_pairs)} 训练对")
print(f"{len(val_pairs)} 验证对")
print(f"{len(test_pairs)} 测试对")
118964 总对
83276 训练对
17844 验证对
17844 测试对

向量化文本数据

我们将使用两个TextVectorization层的实例对文本数据进行向量化(一个用于英语,一个用于西班牙语), 也就是说,将原始字符串转换为整数序列,其中每个整数代表词汇表中一个单词的索引。

英语层将使用默认的字符串标准化(去除标点字符)和分割方案(按空格分割),而西班牙语层将使用自定义标准化,我们将字符 "¿" 添加到要去除的标点字符集合中。

注意:在生产级机器翻译模型中,我不建议在任何语言中去除标点字符。相反,我建议将每个标点字符变成它自己的标记,通过为 TextVectorization 层提供一个自定义的 split 函数来实现。

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_inputsdecoder_inputsencoder_inputs 是向量化的源句子,decoder_inputs 是截至目前的目标句子,也就是说,单词 0 到 N 用于预测目标句子中的单词 N+1(及之后)。
  • 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,它将生成一个新的表示。然后,这个新的表示将传递给 TransformerDecoder,以及截至目前的目标序列(目标单词 0 到 N)。接着,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):
        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

Next, we assemble the end-to-end model.

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"
)

训练我们的模型

我们将使用准确率作为监控验证数据训练进展的快速方法。 请注意,机器翻译通常使用BLEU分数以及其他指标,而不是准确率。

这里我们只训练1个周期,但为了使模型真正收敛, 您应该至少训练30个周期。

epochs = 1  # 这应该至少为30以实现收敛

transformer.summary()
transformer.compile(
    "rmsprop", loss="sparse_categorical_crossentropy", 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)        │                   │         │                      │
├─────────────────────┼───────────────────┼─────────┼──────────────────────┤
│ 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        │ (, ,      │ 12,959… │ decoder_inputs[0][0… │
│ (功能性)        │ 15000)            │         │ transformer_encoder… │
└─────────────────────┴───────────────────┴─────────┴──────────────────────┘
 总参数: 19,960,216 (76.14 MB)
 可训练参数: 19,960,216 (76.14 MB)
 非可训练参数: 0 (0.00 B)
    5/1302 ━━━━━━━━━━━━━━━━━━━━  42s 33ms/step - 准确率: 0.3558 - 损失: 8.3596

警告: 所有日志消息在 absl::InitializeLog() 被调用之前都写入 STDERR
I0000 00:00:1699484373.932513   76082 device_compiler.h:187] 使用 XLA 编译了集群! 该行在进程的生命周期内最多记录一次。

 1302/1302 ━━━━━━━━━━━━━━━━━━━━ 64s 39ms/step - 准确率: 0.7073 - 损失: 2.2372 - 验证准确率: 0.7329 - 验证损失: 1.6477

<keras.src.callbacks.history.History at 0x7ff611f21540>

解码测试句子

最后,让我们演示如何翻译全新的英语句子。 我们只需将向量化的英语句子 以及目标标记 "[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([tokenized_input_sentence, tokenized_target_sentence])

        # ops.argmax(predictions[0, i, :]) 对于 jax 来说不是具体值
        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 轮训练,我们得到了如下结果:

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]