基于Tensorflow实现一个Transformer翻译器-CFANZ编程社区

Transformer是谷歌在2017年的一篇论文"Attention is all you need"提出的一个seq2seq的模型架构，其创造性的提出了自注意力的思想，可以很好的表达序列中各个单词之间的相互注意力关系。这个模型在NLP领域取得了巨大的成功。此外这个模型架构在最近几年也在CV领域取得了令人瞩目的进展，在图像识别，目标检测等方面都达到或超过CNN模型的性能。因此Transformer可以说是人工智能领域最近最值得关注和学习的一个架构。目前有网上已经有很多文章详细解读了Transformer的架构和其细节，这里我将不再重复这方面的内容，而是关注在实战方面，基于Tensorflow来搭建一个Transformer模型，实现法语和英语的翻译。

在Tensorflow的官网上有一个详细的教程，介绍了如何搭建Tranformer来实现葡萄牙语翻译为英语。我也是学习了这个教程之后，进行一些改造，以实现对法语-英语的翻译。

数据集的准备

在这个网站Tab-delimited Bilingual Sentence Pairs from the Tatoeba Project (Good for Anki and Similar Flashcard Applications)可以找到很多不同的语言与英语的翻译。这里我们下载法语-英语的数据作为训练集和验证集。下载http://www.manythings.org/anki/fra-eng.zip这个文件并解压之后，我们可以看到里面每一行对应一个英语句子和一个法语句子，以及句子的贡献者，中间以TAB分隔。

以下代码是读取文件的数据并查看法语和英语的句子：

fra = []
eng = []
with open('fra.txt', 'r') as f:
    content = f.readlines()
    for line in content:
        temp = line.split(sep='\t')
        eng.append(temp[0])
        fra.append(temp[1])

查看这些句子，可以看到有些句子包含特殊字符，例如'Cours\u202f!' 我们需要把这些特殊的不可见字符(\u202f, \xa0 ...)去除掉

new_fra = []
new_eng = []
for item in fra:
    new_fra.append(re.sub('\s', ' ', item).strip().lower())
for item in eng:
    new_eng.append(re.sub('\s', ' ', item).strip().lower())

单词处理为token

因为模型只能处理数字，需要把这些法语和英语的单词转为token。这里采用BERT tokenizer的方式来处理，具体可以参见tensorflow的教程Subword tokenizers | Text | TensorFlow

首先创建两个dataset，分别包含了法语和英语的句子。

ds_fra = tf.data.Dataset.from_tensor_slices(new_fra)
ds_eng = tf.data.Dataset.from_tensor_slices(new_eng)

调用tensorflow的bert_vocab库来创建词汇表，这里定义了一些保留token用于特殊目的，例如[START]标识句子的开始，[UNK]标识一个不在词汇表出现的新单词。

bert_tokenizer_params=dict(lower_case=True)
reserved_tokens=["[PAD]", "[UNK]", "[START]", "[END]"]

bert_vocab_args = dict(
    # The target vocabulary size
    vocab_size = 8000,
    # Reserved tokens that must be included in the vocabulary
    reserved_tokens=reserved_tokens,
    # Arguments for `text.BertTokenizer`
    bert_tokenizer_params=bert_tokenizer_params,
    # Arguments for `wordpiece_vocab.wordpiece_tokenizer_learner_lib.learn`
    learn_params={},
)

fr_vocab = bert_vocab.bert_vocab_from_dataset(
    ds_fra.batch(1000).prefetch(2),
    **bert_vocab_args
)

en_vocab = bert_vocab.bert_vocab_from_dataset(
    ds_eng.batch(1000).prefetch(2),
    **bert_vocab_args
)

词汇表处理完成之后，我们可以看看里面包含哪些内容：

print(en_vocab[:10])
print(en_vocab[100:110])
print(en_vocab[1000:1010])
print(en_vocab[-10:])

输出如下，可以看到词汇表不是严格按照每个英语单词来划分的，例如'##ers'表示某个单词如果以ers结尾，则会划分出一个'##ers'的token

['[PAD]', '[UNK]', '[START]', '[END]', '!', '"', '$', '%', '&', "'"]
['ll', 'there', 've', 'and', 'him', 'time', 'here', 'about', 'get', 'didn']
['##ers', 'chair', 'earth', 'honest', 'succeed', '##ted', 'animals', 'bill', 'drank', 'lend']
['##?', '##j', '##q', '##z', '##°', '##–', '##—', '##‘', '##’', '##€']

把词汇表保存为文件，然后我们就可以实例化两个tokenizer，以实现对法语和英语句子的token化处理。

def write_vocab_file(filepath, vocab):
    with open(filepath, 'w') as f:
        for token in vocab:
            print(token, file=f)
write_vocab_file('fr_vocab.txt', fr_vocab)
write_vocab_file('en_vocab.txt', en_vocab)

fr_tokenizer = text.BertTokenizer('fr_vocab.txt', **bert_tokenizer_params)
en_tokenizer = text.BertTokenizer('en_vocab.txt', **bert_tokenizer_params)

下面我们可以测试一下对一些英语句子进行token处理后的结果，这里我们需要给每个句子的开头和结尾分别加上[START]和[END]这两个特殊的token，这样可以方便以后模型的训练。

START = tf.argmax(tf.constant(reserved_tokens) == "[START]")
END = tf.argmax(tf.constant(reserved_tokens) == "[END]")

def add_start_end(ragged):
    count = ragged.bounding_shape()[0]
    starts = tf.fill([count,1], START)
    ends = tf.fill([count,1], END)
    return tf.concat([starts, ragged, ends], axis=1)

sentences = ["Hello Roy!", "The sky is blue.", "Nice to meet you!"]

add_start_end(en_tokenizer.tokenize(sentences).merge_dims(1,2)).to_tensor()

输出结果如下：

<tf.Tensor: shape=(3, 7), dtype=int64, numpy=
array([[   2, 1830,   45, 3450,    4,    3,    0],
       [   2,   62, 1132,   64,  996,   13,    3],
       [   2,  353,   61,  416,   60,    4,    3]])>

构建数据集

现在我们可以构建训练集和验证集了。这里需要把法语和英语的句子都包括在数据集中，其中法语句子作为Transformer编码器的输入，英语句子作为解码器的输入以及模型输出的Target。这里我们用Pandas构造一个Dataframe，随机划分其中80%的数据为训练集，其余为验证集。然后转换为Tensorflow的dataset

df = pd.DataFrame(data={'fra':new_fra, 'eng':new_eng})

# Shuffle the Dataframe
recordnum = df.count()['fra']
indexlist = list(range(recordnum-1))
random.shuffle(indexlist)
df_train = df.loc[indexlist[:int(recordnum*0.8)]]
df_val = df.loc[indexlist[int(recordnum*0.8):]]

ds_train = tf.data.Dataset.from_tensor_slices((df_train.fra.values, df_train.eng.values))
ds_val = tf.data.Dataset.from_tensor_slices((df_val.fra.values, df_val.eng.values))

查看训练集的句子最多包含多少个token

lengths = []

for fr_examples, en_examples in ds_train.batch(1024):
    fr_tokens = fr_tokenizer.tokenize(fr_examples)
    lengths.append(fr_tokens.row_lengths())

    en_tokens = en_tokenizer.tokenize(en_examples)
    lengths.append(en_tokens.row_lengths())
    print('.', end='', flush=True)

all_lengths = np.concatenate(lengths)

plt.hist(all_lengths, np.linspace(0, 100, 11))
plt.ylim(plt.ylim())
max_length = max(all_lengths)
plt.plot([max_length, max_length], plt.ylim())
plt.title(f'Max tokens per example: {max_length}');

从结果中可以看到训练集的句子转换为token后最多包含67个token:

之后就可以为数据集生成batch，如以下代码：

BUFFER_SIZE = 20000
BATCH_SIZE = 64
MAX_TOKENS = 67

def filter_max_tokens(fr, en):
    num_tokens = tf.maximum(tf.shape(fr)[1],tf.shape(en)[1])
    return num_tokens < MAX_TOKENS

def tokenize_pairs(fr, en):
    fr = add_start_end(fr_tokenizer.tokenize(fr).merge_dims(1,2))
    # Convert from ragged to dense, padding with zeros.
    fr = fr.to_tensor()

    en = add_start_end(en_tokenizer.tokenize(en).merge_dims(1,2))
    # Convert from ragged to dense, padding with zeros.
    en = en.to_tensor()
    return fr, en

def make_batches(ds):
    return (
        ds
        .cache()
        .shuffle(BUFFER_SIZE)
        .batch(BATCH_SIZE)
        .map(tokenize_pairs, num_parallel_calls=tf.data.AUTOTUNE)
        .filter(filter_max_tokens)
        .prefetch(tf.data.AUTOTUNE))

train_batches = make_batches(ds_train)
val_batches = make_batches(ds_val)