使用 TF 文本进行 BERT 预处理

在 TensorFlow.org 上查看

在 Google Colab 中运行

在 GitHub 上查看

下载笔记本

概述

文本预处理是将原始文本端到端转换为模型的整数输入。NLP 模型通常伴随着数百行（如果不是数千行）的 Python 代码用于预处理文本。文本预处理通常对模型来说是一个挑战，因为

训练-服务偏差。 确保模型输入的预处理逻辑在模型开发的所有阶段（例如预训练、微调、评估、推理）保持一致变得越来越困难。在不同阶段使用不同的超参数、标记化、字符串预处理算法，或者只是以不一致的方式打包模型输入，可能会导致难以调试和灾难性的模型效果。
效率和灵活性。 虽然预处理可以在离线进行（例如，通过将处理后的输出写入磁盘上的文件，然后在输入管道中重新使用这些预处理数据），但这种方法会产生额外的文件读写成本。如果需要动态进行预处理决策，离线预处理也很不方便。尝试不同的选项将需要重新生成数据集。
复杂的模型接口。 当文本模型的输入是纯文本时，它们更容易理解。当模型的输入需要额外的间接编码步骤时，很难理解模型。降低预处理复杂性对于模型调试、服务和评估特别有用。

此外，更简单的模型接口也使尝试模型（例如推理或训练）变得更加方便，可以尝试不同的、未探索的数据集。

使用 TF.Text 进行文本预处理

使用 TF.Text 的文本预处理 API，我们可以构建一个预处理函数，该函数可以将用户的文本数据集转换为模型的整数输入。用户可以将预处理直接打包为其模型的一部分，以缓解上述问题。

本教程将展示如何使用 TF.Text 预处理操作将文本数据转换为 BERT 模型的输入，以及转换为“Masked LM and Masking Procedure”中描述的语言掩码预训练任务的输入，该任务在 BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding 中进行了描述。该过程涉及将文本标记化为子词单元、组合句子、将内容修剪为固定大小以及提取用于掩码语言建模任务的标签。

设置

首先，让我们导入所需的包和库。

pip install -q -U "tensorflow-text==2.11.*"

import tensorflow as tf
import tensorflow_text as text
import functools
print("TensorFlow version: ", tf.__version__)

2024-06-25 11:46:30.898009: W tensorflow/compiler/xla/stream_executor/platform/default/dso_loader.cc:64] Could not load dynamic library 'libcudart.so.11.0'; dlerror: libcudart.so.11.0: cannot open shared object file: No such file or directory
2024-06-25 11:46:31.716557: W tensorflow/compiler/xla/stream_executor/platform/default/dso_loader.cc:64] Could not load dynamic library 'libnvinfer.so.7'; dlerror: libnvinfer.so.7: cannot open shared object file: No such file or directory
2024-06-25 11:46:31.716646: W tensorflow/compiler/xla/stream_executor/platform/default/dso_loader.cc:64] Could not load dynamic library 'libnvinfer_plugin.so.7'; dlerror: libnvinfer_plugin.so.7: cannot open shared object file: No such file or directory
2024-06-25 11:46:31.716655: W tensorflow/compiler/tf2tensorrt/utils/py_utils.cc:38] TF-TRT Warning: Cannot dlopen some TensorRT libraries. If you would like to use Nvidia GPU with TensorRT, please make sure the missing libraries mentioned above are installed properly.
TensorFlow version:  2.11.1

我们的数据包含两个文本特征，我们可以创建一个示例 tf.data.Dataset。我们的目标是创建一个函数，我们可以向 Dataset.map() 提供该函数，以便在训练中使用。

examples = {
    "text_a": [
      "Sponge bob Squarepants is an Avenger",
      "Marvel Avengers"
    ],
    "text_b": [
     "Barack Obama is the President.",
     "President is the highest office"
  ],
}

dataset = tf.data.Dataset.from_tensor_slices(examples)
next(iter(dataset))

2024-06-25 11:46:33.419785: W tensorflow/compiler/xla/stream_executor/platform/default/dso_loader.cc:64] Could not load dynamic library 'libcudart.so.11.0'; dlerror: libcudart.so.11.0: cannot open shared object file: No such file or directory
2024-06-25 11:46:33.419883: W tensorflow/compiler/xla/stream_executor/platform/default/dso_loader.cc:64] Could not load dynamic library 'libcublas.so.11'; dlerror: libcublas.so.11: cannot open shared object file: No such file or directory
2024-06-25 11:46:33.419946: W tensorflow/compiler/xla/stream_executor/platform/default/dso_loader.cc:64] Could not load dynamic library 'libcublasLt.so.11'; dlerror: libcublasLt.so.11: cannot open shared object file: No such file or directory
2024-06-25 11:46:33.420005: W tensorflow/compiler/xla/stream_executor/platform/default/dso_loader.cc:64] Could not load dynamic library 'libcufft.so.10'; dlerror: libcufft.so.10: cannot open shared object file: No such file or directory
2024-06-25 11:46:33.475841: W tensorflow/compiler/xla/stream_executor/platform/default/dso_loader.cc:64] Could not load dynamic library 'libcusparse.so.11'; dlerror: libcusparse.so.11: cannot open shared object file: No such file or directory
2024-06-25 11:46:33.476035: W tensorflow/core/common_runtime/gpu/gpu_device.cc:1934] Cannot dlopen some GPU libraries. Please make sure the missing libraries mentioned above are installed properly if you would like to use GPU. Follow the guide at https://tensorflowcn.cn/install/gpu for how to download and setup the required libraries for your platform.
Skipping registering GPU devices...
{'text_a': <tf.Tensor: shape=(), dtype=string, numpy=b'Sponge bob Squarepants is an Avenger'>,
 'text_b': <tf.Tensor: shape=(), dtype=string, numpy=b'Barack Obama is the President.'>}

标记化

我们的第一步是运行任何字符串预处理并标记化我们的数据集。这可以使用 text.BertTokenizer 完成，它是一个 text.Splitter，可以根据从 Wordpiece 算法生成的词汇表，将句子标记化为子词或词片，用于 BERT 模型。您可以从这里了解有关 TF.Text 中提供的其他子词标记器的更多信息。

词汇表可以来自先前生成的 BERT 检查点，或者您可以在自己的数据上自行生成一个。在本例中，让我们创建一个玩具词汇表

_VOCAB = [
    # Special tokens
    b"[UNK]", b"[MASK]", b"[RANDOM]", b"[CLS]", b"[SEP]",
    # Suffixes
    b"##ack", b"##ama", b"##ger", b"##gers", b"##onge", b"##pants",  b"##uare",
    b"##vel", b"##ven", b"an", b"A", b"Bar", b"Hates", b"Mar", b"Ob",
    b"Patrick", b"President", b"Sp", b"Sq", b"bob", b"box", b"has", b"highest",
    b"is", b"office", b"the",
]

_START_TOKEN = _VOCAB.index(b"[CLS]")
_END_TOKEN = _VOCAB.index(b"[SEP]")
_MASK_TOKEN = _VOCAB.index(b"[MASK]")
_RANDOM_TOKEN = _VOCAB.index(b"[RANDOM]")
_UNK_TOKEN = _VOCAB.index(b"[UNK]")
_MAX_SEQ_LEN = 8
_MAX_PREDICTIONS_PER_BATCH = 5

_VOCAB_SIZE = len(_VOCAB)

lookup_table = tf.lookup.StaticVocabularyTable(
    tf.lookup.KeyValueTensorInitializer(
      keys=_VOCAB,
      key_dtype=tf.string,
      values=tf.range(
          tf.size(_VOCAB, out_type=tf.int64), dtype=tf.int64),
          value_dtype=tf.int64
        ),
      num_oov_buckets=1
)

让我们使用上面的词汇表构建一个 text.BertTokenizer，并将文本输入标记化为 RaggedTensor。`.

bert_tokenizer = text.BertTokenizer(lookup_table, token_out_type=tf.string)
bert_tokenizer.tokenize(examples["text_a"])

<tf.RaggedTensor [[[b'Sp', b'##onge'], [b'bob'], [b'Sq', b'##uare', b'##pants'], [b'is'],
  [b'an'], [b'A', b'##ven', b'##ger']]                                  ,
 [[b'Mar', b'##vel'], [b'A', b'##ven', b'##gers']]]>

bert_tokenizer.tokenize(examples["text_b"])

<tf.RaggedTensor [[[b'Bar', b'##ack'], [b'Ob', b'##ama'], [b'is'], [b'the'], [b'President'],
  [b'[UNK]']]                                                              ,
 [[b'President'], [b'is'], [b'the'], [b'highest'], [b'office']]]>

来自 text.BertTokenizer 的文本输出使我们能够看到文本是如何被标记化的，但模型需要整数 ID。我们可以将 token_out_type 参数设置为 tf.int64 以获取整数 ID（它们是词汇表中的索引）。

bert_tokenizer = text.BertTokenizer(lookup_table, token_out_type=tf.int64)
segment_a = bert_tokenizer.tokenize(examples["text_a"])
segment_a

<tf.RaggedTensor [[[22, 9], [24], [23, 11, 10], [28], [14], [15, 13, 7]],
 [[18, 12], [15, 13, 8]]]>

segment_b = bert_tokenizer.tokenize(examples["text_b"])
segment_b

<tf.RaggedTensor [[[16, 5], [19, 6], [28], [30], [21], [0]], [[21], [28], [30], [27], [29]]]>

text.BertTokenizer 返回一个形状为 [batch, num_tokens, num_wordpieces] 的 RaggedTensor。因为我们当前的用例不需要额外的 num_tokens 维度，所以我们可以合并最后两个维度，以获得一个形状为 [batch, num_wordpieces] 的 RaggedTensor

segment_a = segment_a.merge_dims(-2, -1)
segment_a

<tf.RaggedTensor [[22, 9, 24, 23, 11, 10, 28, 14, 15, 13, 7], [18, 12, 15, 13, 8]]>

segment_b = segment_b.merge_dims(-2, -1)
segment_b

<tf.RaggedTensor [[16, 5, 19, 6, 28, 30, 21, 0], [21, 28, 30, 27, 29]]>

内容修剪

BERT 的主要输入是两个句子的串联。但是，BERT 要求输入具有固定的大小和形状，而我们的内容可能超出了我们的预算。

我们可以使用 text.Trimmer 来将我们的内容修剪到预定的尺寸（在最后一个轴上连接后）。有不同的 text.Trimmer 类型，它们使用不同的算法选择要保留的内容。例如，text.RoundRobinTrimmer 将为每个段分配相等的配额，但可能会修剪句子的末尾。 text.WaterfallTrimmer 将从最后一个句子的末尾开始修剪。

在我们的示例中，我们将使用 RoundRobinTrimmer，它以从左到右的方式从每个段中选择项目。

trimmer = text.RoundRobinTrimmer(max_seq_length=_MAX_SEQ_LEN)
trimmed = trimmer.trim([segment_a, segment_b])
trimmed

[<tf.RaggedTensor [[22, 9, 24, 23],
  [18, 12, 15, 13]]>,
 <tf.RaggedTensor [[16, 5, 19, 6],
  [21, 28, 30, 27]]>]

trimmed 现在包含了在批次中元素数量为 8 个元素（在轴=-1 上连接后）的段。

组合段

现在我们已经修剪了段，我们可以将它们组合在一起以获得单个 RaggedTensor。BERT 使用特殊标记来指示段的开头 ([CLS]) 和结尾 ([SEP])。我们还需要一个 RaggedTensor 来指示组合的 Tensor 中哪些项目属于哪个段。我们可以使用 text.combine_segments() 来获取这两个带有插入特殊标记的 Tensor。

segments_combined, segments_ids = text.combine_segments(
  trimmed,
  start_of_sequence_id=_START_TOKEN, end_of_segment_id=_END_TOKEN)
segments_combined, segments_ids

(<tf.RaggedTensor [[3, 22, 9, 24, 23, 4, 16, 5, 19, 6, 4],
  [3, 18, 12, 15, 13, 4, 21, 28, 30, 27, 4]]>,
 <tf.RaggedTensor [[0, 0, 0, 0, 0, 0, 1, 1, 1, 1, 1],
  [0, 0, 0, 0, 0, 0, 1, 1, 1, 1, 1]]>)

掩码语言模型任务

现在我们有了基本输入，我们可以开始提取在 BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding 中描述的“掩码 LM 和掩码过程”任务所需的输入。

掩码语言模型任务有两个子问题需要我们考虑：（1）选择哪些项目进行掩码，以及（2）它们被分配什么值？

项目选择

因为我们将选择随机选择项目进行掩码，所以我们将使用 text.RandomItemSelector。 RandomItemSelector 随机选择批次中的项目，受给定限制 (max_selections_per_batch、selection_rate 和 unselectable_ids) 的约束，并返回一个布尔掩码，指示哪些项目被选中。

random_selector = text.RandomItemSelector(
    max_selections_per_batch=_MAX_PREDICTIONS_PER_BATCH,
    selection_rate=0.2,
    unselectable_ids=[_START_TOKEN, _END_TOKEN, _UNK_TOKEN]
)
selected = random_selector.get_selection_mask(
    segments_combined, axis=1)
selected

<tf.RaggedTensor [[False, False, False, True, False, False, True, False, False, False,
  False],
 [False, False, False, True, False, False, False, True, False, False,
  False]]>

选择掩码值

原始 BERT 论文中描述的用于选择掩码值的方法如下

对于 mask_token_rate 的时间，用 [MASK] 标记替换项目

"my dog is hairy" -> "my dog is [MASK]"

对于 random_token_rate 的时间，用随机词替换项目

"my dog is hairy" -> "my dog is apple"

对于 1 - mask_token_rate - random_token_rate 的时间，保持项目不变

"my dog is hairy" -> "my dog is hairy."

text.MaskedValuesChooser 封装了这种逻辑，可以用于我们的预处理函数。以下是一个 MaskValuesChooser 在给定 80% 的 mask_token_rate 和默认 random_token_rate 时返回的结果示例

mask_values_chooser = text.MaskValuesChooser(_VOCAB_SIZE, _MASK_TOKEN, 0.8)
mask_values_chooser.get_mask_values(segments_combined)

<tf.RaggedTensor [[1, 1, 1, 1, 1, 1, 25, 1, 19, 1, 1],
 [1, 1, 21, 1, 1, 1, 16, 1, 1, 1, 20]]>

当提供 RaggedTensor 输入时，text.MaskValuesChooser 返回一个形状相同的 RaggedTensor，其中包含 _MASK_VALUE (0)、随机 ID 或相同的未更改 ID。

生成掩码语言模型任务的输入

现在我们有了 RandomItemSelector 来帮助我们选择要掩码的项目，以及 text.MaskValuesChooser 来分配值，我们可以使用 text.mask_language_model() 来组装此任务的所有输入，以供我们的 BERT 模型使用。

masked_token_ids, masked_pos, masked_lm_ids = text.mask_language_model(
  segments_combined,
  item_selector=random_selector, mask_values_chooser=mask_values_chooser)

让我们深入研究并检查 mask_language_model() 的输出。 masked_token_ids 的输出是

masked_token_ids

<tf.RaggedTensor [[3, 22, 9, 24, 23, 4, 1, 5, 19, 6, 4],
 [3, 13, 12, 15, 13, 4, 21, 28, 30, 27, 4]]>

请记住，我们的输入使用词汇表进行编码。如果我们使用词汇表解码 masked_token_ids，我们将得到

tf.gather(_VOCAB, masked_token_ids)

<tf.RaggedTensor [[b'[CLS]', b'Sp', b'##onge', b'bob', b'Sq', b'[SEP]', b'[MASK]',
  b'##ack', b'Ob', b'##ama', b'[SEP]'],
 [b'[CLS]', b'##ven', b'##vel', b'A', b'##ven', b'[SEP]', b'President',
  b'is', b'the', b'highest', b'[SEP]']]>

请注意，一些词片标记已被 [MASK]、[RANDOM] 或不同的 ID 值替换。 masked_pos 输出为我们提供了已替换的标记的索引（在相应的批次中）。

masked_pos

<tf.RaggedTensor [[2, 6],
 [1, 7]]>

masked_lm_ids 为我们提供了标记的原始值。

masked_lm_ids

<tf.RaggedTensor [[9, 16],
 [18, 28]]>

我们再次可以解码这些 ID 以获得人类可读的值。

tf.gather(_VOCAB, masked_lm_ids)

<tf.RaggedTensor [[b'##onge', b'Bar'],
 [b'Mar', b'is']]>

填充模型输入

现在我们有了模型的所有输入，预处理的最后一步是将它们打包成固定二维 Tensor，并进行填充，以及生成一个掩码 Tensor，指示哪些值是填充值。我们可以使用 text.pad_model_inputs() 来帮助我们完成此任务。

# Prepare and pad combined segment inputs
input_word_ids, input_mask = text.pad_model_inputs(
  masked_token_ids, max_seq_length=_MAX_SEQ_LEN)
input_type_ids, _ = text.pad_model_inputs(
  segments_ids, max_seq_length=_MAX_SEQ_LEN)

# Prepare and pad masking task inputs
masked_lm_positions, masked_lm_weights = text.pad_model_inputs(
  masked_pos, max_seq_length=_MAX_PREDICTIONS_PER_BATCH)
masked_lm_ids, _ = text.pad_model_inputs(
  masked_lm_ids, max_seq_length=_MAX_PREDICTIONS_PER_BATCH)

model_inputs = {
    "input_word_ids": input_word_ids,
    "input_mask": input_mask,
    "input_type_ids": input_type_ids,
    "masked_lm_ids": masked_lm_ids,
    "masked_lm_positions": masked_lm_positions,
    "masked_lm_weights": masked_lm_weights,
}
model_inputs

{'input_word_ids': <tf.Tensor: shape=(2, 8), dtype=int64, numpy=
 array([[ 3, 22,  9, 24, 23,  4,  1,  5],
        [ 3, 13, 12, 15, 13,  4, 21, 28]])>,
 'input_mask': <tf.Tensor: shape=(2, 8), dtype=int64, numpy=
 array([[1, 1, 1, 1, 1, 1, 1, 1],
        [1, 1, 1, 1, 1, 1, 1, 1]])>,
 'input_type_ids': <tf.Tensor: shape=(2, 8), dtype=int64, numpy=
 array([[0, 0, 0, 0, 0, 0, 1, 1],
        [0, 0, 0, 0, 0, 0, 1, 1]])>,
 'masked_lm_ids': <tf.Tensor: shape=(2, 5), dtype=int64, numpy=
 array([[ 9, 16,  0,  0,  0],
        [18, 28,  0,  0,  0]])>,
 'masked_lm_positions': <tf.Tensor: shape=(2, 5), dtype=int64, numpy=
 array([[2, 6, 0, 0, 0],
        [1, 7, 0, 0, 0]])>,
 'masked_lm_weights': <tf.Tensor: shape=(2, 5), dtype=int64, numpy=
 array([[1, 1, 0, 0, 0],
        [1, 1, 0, 0, 0]])>}

回顾

让我们回顾一下到目前为止我们所做的，并组装我们的预处理函数。以下是我们所做的

def bert_pretrain_preprocess(vocab_table, features):
  # Input is a string Tensor of documents, shape [batch, 1].
  text_a = features["text_a"]
  text_b = features["text_b"]

  # Tokenize segments to shape [num_sentences, (num_words)] each.
  tokenizer = text.BertTokenizer(
      vocab_table,
      token_out_type=tf.int64)
  segments = [tokenizer.tokenize(text).merge_dims(
      1, -1) for text in (text_a, text_b)]

  # Truncate inputs to a maximum length.
  trimmer = text.RoundRobinTrimmer(max_seq_length=6)
  trimmed_segments = trimmer.trim(segments)

  # Combine segments, get segment ids and add special tokens.
  segments_combined, segment_ids = text.combine_segments(
      trimmed_segments,
      start_of_sequence_id=_START_TOKEN,
      end_of_segment_id=_END_TOKEN)

  # Apply dynamic masking task.
  masked_input_ids, masked_lm_positions, masked_lm_ids = (
      text.mask_language_model(
        segments_combined,
        random_selector,
        mask_values_chooser,
      )
  )

  # Prepare and pad combined segment inputs
  input_word_ids, input_mask = text.pad_model_inputs(
    masked_input_ids, max_seq_length=_MAX_SEQ_LEN)
  input_type_ids, _ = text.pad_model_inputs(
    segment_ids, max_seq_length=_MAX_SEQ_LEN)

  # Prepare and pad masking task inputs
  masked_lm_positions, masked_lm_weights = text.pad_model_inputs(
    masked_lm_positions, max_seq_length=_MAX_PREDICTIONS_PER_BATCH)
  masked_lm_ids, _ = text.pad_model_inputs(
    masked_lm_ids, max_seq_length=_MAX_PREDICTIONS_PER_BATCH)

  model_inputs = {
      "input_word_ids": input_word_ids,
      "input_mask": input_mask,
      "input_type_ids": input_type_ids,
      "masked_lm_ids": masked_lm_ids,
      "masked_lm_positions": masked_lm_positions,
      "masked_lm_weights": masked_lm_weights,
  }
  return model_inputs

我们之前构建了一个 tf.data.Dataset，现在我们可以使用我们组装的预处理函数 bert_pretrain_preprocess() 在 Dataset.map() 中使用它。这使我们能够创建一个输入管道，将原始字符串数据转换为整数输入，并直接馈送到我们的模型中。

dataset = (
    tf.data.Dataset.from_tensors(examples)
    .map(functools.partial(bert_pretrain_preprocess, lookup_table))
)

next(iter(dataset))

WARNING:tensorflow:From /tmpfs/src/tf_docs_env/lib/python3.9/site-packages/tensorflow/python/autograph/pyct/static_analysis/liveness.py:83: Analyzer.lamba_check (from tensorflow.python.autograph.pyct.static_analysis.liveness) is deprecated and will be removed after 2023-09-23.
Instructions for updating:
Lambda fuctions will be no more assumed to be used in the statement where they are used, or at least in the same block. https://github.com/tensorflow/tensorflow/issues/56089
{'input_word_ids': <tf.Tensor: shape=(2, 8), dtype=int64, numpy=
 array([[ 3, 22,  9, 24,  4,  1,  1, 19],
        [ 3, 18, 12, 15,  4, 21, 28, 30]])>,
 'input_mask': <tf.Tensor: shape=(2, 8), dtype=int64, numpy=
 array([[1, 1, 1, 1, 1, 1, 1, 1],
        [1, 1, 1, 1, 1, 1, 1, 1]])>,
 'input_type_ids': <tf.Tensor: shape=(2, 8), dtype=int64, numpy=
 array([[0, 0, 0, 0, 0, 1, 1, 1],
        [0, 0, 0, 0, 0, 1, 1, 1]])>,
 'masked_lm_ids': <tf.Tensor: shape=(2, 5), dtype=int64, numpy=
 array([[16,  5,  0,  0,  0],
        [12, 15,  0,  0,  0]])>,
 'masked_lm_positions': <tf.Tensor: shape=(2, 5), dtype=int64, numpy=
 array([[5, 6, 0, 0, 0],
        [2, 3, 0, 0, 0]])>,
 'masked_lm_weights': <tf.Tensor: shape=(2, 5), dtype=int64, numpy=
 array([[1, 1, 0, 0, 0],
        [1, 1, 0, 0, 0]])>}