如何在TensorFlow中对无法载入内存的超大文本数据集分词？

针对无法载入内存的超大文本数据集，完全可以不用依赖外部库，借助TensorFlow原生工具实现分批加载+实时分词的流程，和图像领域的ImageDataGenerator逻辑类似：

核心思路 #

用tf.data API分批读取原始文本数据（每次仅加载一批到内存），再结合TensorFlow内置的分词工具对每一批数据实时处理，无需提前完成全量分词。

具体实现步骤 #

1. 用tf.data分批读取文本数据 #

通过tf.data.TextLineDataset直接从文件读取文本，按设定批次大小加载，避免一次性载入全部数据：

# 读取单个或多个超大文本文件
dataset = tf.data.TextLineDataset(["large_text_1.txt", "large_text_2.txt"])
# 根据内存情况设置批次大小
batch_size = 32
# 按批次划分数据集
dataset = dataset.batch(batch_size)

2. 用内置TextVectorization层实现批量分词 #

TextVectorization是TensorFlow原生的文本处理层，支持批量分词，还能自动构建词汇表：

# 初始化分词层，可自定义规则
vectorize_layer = tf.keras.layers.TextVectorization(
    max_tokens=10000,  # 限制最大词汇数量
    output_mode='int',  # 输出单词对应的整数索引
    split='whitespace',  # 按空格分词，支持自定义分词逻辑
    standardize='lower_and_strip_punctuation'  # 预处理：转小写+去除标点
)

# 用部分数据构建词汇表，无需加载全部数据集
# 取100批数据来适配，平衡词汇表准确性和内存占用
vectorize_layer.adapt(dataset.take(100))

# 将分词层应用到整个数据集，每一批数据都会被实时分词
vectorized_dataset = dataset.map(lambda text_batch: vectorize_layer(text_batch))

3. 自定义批量分词逻辑（可选） #

如果需要更灵活的分词规则（比如中文分词、自定义过滤规则），可以用tf.strings原生函数实现批量处理：

@tf.function
def custom_batch_tokenizer(text_batch):
    # 转小写
    text_batch = tf.strings.lower(text_batch)
    # 按自定义分隔符分词（这里以中文全角空格为例，可替换为其他规则）
    tokens = tf.strings.split(text_batch, sep="　")
    # 过滤长度小于2的词
    tokens = tf.ragged.map_flat_values(lambda x: x if tf.strings.length(x)>=2 else "", tokens)
    return tokens

# 应用自定义分词到数据集
tokenized_dataset = dataset.map(custom_batch_tokenizer)

4. 直接用于模型训练 #

处理后的数据集可以直接输入模型训练，全程无需加载全部数据到内存：

# 构建简单模型示例
model = tf.keras.Sequential([
    tf.keras.layers.Embedding(10000, 128),
    tf.keras.layers.GlobalAveragePooling1D(),
    tf.keras.layers.Dense(1, activation='sigmoid')
])

model.compile(optimizer='adam', loss='binary_crossentropy', metrics=['accuracy'])
# 用分批分词后的数据集训练
model.fit(vectorized_dataset, epochs=10)

关键说明 #

• TensorFlow内置的TextVectorization支持整合到模型中，部署时可以直接将文本输入模型，无需额外处理分词逻辑；
• 所有操作都是基于TensorFlow的图模式实现，效率高且支持GPU加速；
• 完全无需依赖外部库，仅用TensorFlow原生工具即可完成超大文本数据集的分批分词。

内容的提问来源于stack exchange，提问作者jgoh