应使用Hugging Face的TFAutoModelForMaskedLM加载预训练模型（如bert-base-uncased）进行微调，而非纯TensorFlow从零实现MLM头，因其自动处理词表对齐、掩码位置损失计算、权重共享及label=-100忽略逻辑，避免索引越界、loss=nan等常见错误。
直接用 TensorFlow 实现一个完整、可训练的掩码语言模型（MLM）不现实——TF 没有内置 RobertaTokenizer 或 fill_mask 接口，也不自带预训练 MLM 头结构；你真正该做的是：用 transformers 库加载 Hugging Face 的预训练模型（如 bert-base-uncased），再用 TensorFlow 后端进行微调或推理。
为什么不能从零用纯 TensorFlow 写 MLM 头
MLM 的核心是“预测被 [MASK] 替换的 token”，这要求：① 词表对齐（tokenizer.vocab 必须和模型 embedding 层严格一致）；② 损失只计算被掩码位置（非所有 token）；③ 输出 logits 需经 tf.gather_nd 或布尔索引提取对应位置。纯 TF 实现极易在索引对齐、padding 处理、label 平滑上出错。
常见错误现象：InvalidArgumentError: indices[0] = -1 is not in [0, 30522)（label 值越界）、loss=nan（未屏蔽 padding 位置的 loss）、output shape mismatch（logits 维度和词表大小不一致）。
- 预训练 MLM 头本质是线性层 + 词表映射，但权重需与 embedding 层共享（或 tied），
transformers自动处理该逻辑，手写易漏 - TensorFlow 2.x 的
tf.keras.Model不原生支持动态掩码位置采样，需手动构建tf.where(masked_positions)，调试成本高 - Hugging Face 的
TFAutoModelForMaskedLM已封装好损失计算、label 处理、梯度截断，且兼容tf.data.Dataset
用 TFAutoModelForMaskedLM 加载并微调
这是最稳妥的路径：利用 transformers 的 TF 接口，保留全部预训练结构，只替换数据 pipeline 为 TensorFlow 原生流程。
立即学习“Python免费学习笔记（深入）”；
使用场景：你在已有 TF 生产环境（如 TF Serving）、或必须用 tf.function 装饰训练 step、或需和 TF 数据增强链路（如 tf.image 扩展）耦合。
- 安装依赖：
pip install transformers[tf-cpu]（或[tf-gpu]） - 加载模型：
model = TFAutoModelForMaskedLM.from_pretrained("bert-base-uncased") - tokenizer 必须匹配：
tokenizer = AutoTokenizer.from_pretrained("bert-base-uncased")，否则input_ids和模型词表错位 - 构建 dataset 时，用
tokenizer(..., return_tensors="tf")确保输出是tf.Tensor，而非 Python list - label 构造必须显式：对每个样本，先用
tokenizer.mask_token_id替换随机 15% token，再把原 token id 填入labels张量，未掩码位置填-100（TF 默认忽略该值的 loss 计算）
自定义 MLM 数据集的 tf.data pipeline 关键点
TF 原生 pipeline 容易在 padding、mask 生成、label 对齐三处翻车。重点不是“怎么写”，而是“怎么避坑”。
性能影响：若用 tf.py_function 调用 Python tokenizer，会严重拖慢 pipeline；必须用 tokenizer.encode_plus 预处理后存为 TFRecord，或用 tf.lookup.StaticVocabularyTable 实现纯 TF 分词（仅限简单分词器）。
- padding 长度必须统一：
max_length=512，且truncation=True，否则 batch 内 shape 不一致导致ValueError: All input tensors must have the same shape - mask 不能用
np.random：TF graph 模式下需用tf.random.uniform+tf.where动态生成掩码位置 -
labels张量必须和input_ids同 shape，未掩码位置设为-100（不是0或tokenizer.pad_token_id），否则 loss 错误累积 - batch size 建议 ≤ 16：BERT 类模型在单卡 GPU 上，batch_size=16 已接近显存极限，更大值触发
ResourceExhaustedError
推理时如何获取被掩码位置的 top-k 预测
不同于 PyTorch 的 model(**inputs).logits，TF 模型输出是 TFMaskedLMOutput 字典，需手动提取 logits 并索引。
容易踩的坑：直接对整个 logits 取 tf.nn.top_k 会返回全词表预测，而你只关心被 [MASK] 占位的位置。
- 先定位 mask 位置：
mask_positions = tf.where(input_ids == tokenizer.mask_token_id) - 用
tf.gather_nd(logits, mask_positions)提取对应 logits（shape=[num_masks, vocab_size]） - 再做
tf.nn.top_k(mask_logits, k=5)，得到values和indices - 最后用
tokenizer.convert_ids_to_tokens(indices.numpy().tolist())解码——注意这里要退出 graph 模式，因为convert_ids_to_tokens是 Python 方法
复杂点在于 mask_positions 的 batch 维度处理：若 batch 中各句 mask 数量不同，tf.gather_nd 会报错，必须用 tf.RaggedTensor 或 padding 对齐，这点文档极少提，但实际必遇。
