如何基于TensorFlow TextLineDataset实现双词（Token）迭代？

当然可以实现！TextLineDataset本身是逐行返回文本，但我们可以通过一系列数据集转换操作，把它变成每次返回两个单词的形式。下面是具体的实现步骤和代码：

步骤说明 #

1. 读取原始文本行：用TextLineDataset加载你的文本文件，这一步和你原来的代码一致。
2. 分词（拆分句子为单词）：对每一行文本进行拆分，把句子转换成单个单词的序列。这里可以用tf.strings.split来完成分词。
3. 展平单词序列：通过flat_map操作，把每行拆分出的单词“展开”，让整个数据集的元素从“整行句子”变成“单个单词”。
4. 批量获取两个单词：最后用batch(2)操作，让迭代器每次返回两个单词的批次。

完整代码示例 #

import tensorflow as tf

# 1. 加载文本文件
sentences = tf.data.TextLineDataset("data/train.src")

# 2. 分词并展平成单个单词的数据集
def split_sentence(sentence):
    # 拆分句子为单词，返回单词的数据集
    return tf.data.Dataset.from_tensor_slices(tf.strings.split(sentence))

words_dataset = sentences.flat_map(split_sentence)

# 3. 每次获取两个单词
batch_dataset = words_dataset.batch(2)

# 4. 创建迭代器并遍历
iterator = batch_dataset.make_initializable_iterator()
next_pair = iterator.get_next()

with tf.Session() as sess:
    sess.run(tf.tables_initializer())
    sess.run(iterator.initializer)
    try:
        while True:
            pair = sess.run(next_pair)
            print(pair)
    except tf.errors.OutOfRangeError:
        print("遍历结束")

补充说明 #

• 如果你的需求是获取连续的滑动窗口对（比如"hello world foo"变成("hello", "world")、("world", "foo")这种二元组），可以用window操作替代batch：

  # 创建滑动窗口，每个窗口包含2个单词，步长为1，丢弃不足2个单词的窗口
  window_dataset = words_dataset.window(size=2, shift=1, drop_remainder=True)
  # 把窗口转换成可迭代的批次
  pair_dataset = window_dataset.flat_map(lambda window: window.batch(2))
  

  这样迭代器每次返回的就是连续的两个单词对。

• 注意：如果你的文本有特殊分词需求（比如处理标点、统一大小写），可以在split_sentence函数里添加预处理逻辑，比如用tf.strings.regex_replace先清理文本内容。

内容的提问来源于stack exchange，提问作者Valentin Macé