ReadTwice:面向超长文档阅读的“读两遍“BERT 模型实战指南
发布时间:2026/9/21 15:29:45 锦皓数字建站

人工智能深度学习NLP计算机视觉强化学习【免费下载链接】google-researchGoogle Research项目地址https://gitcode.com/gh_mirrors/go/google-research点击查看免费下载导读ReadTwice 是 Google Research 提出的带记忆的超长文档阅读Reading Very Large Documents with Memories模型实现核心思想是让模型对文档进行两次阅读第一次读取生成全局摘要记忆第二次阅读时携带这些摘要信息重新处理原文从而突破 Transformer 输入长度限制、更好地处理跨段长距离依赖。本文以readtwice/README.md为主干结合仓库源码完整讲解该模型的架构原理、依赖安装、预训练权重加载以及 HotpotQA、TriviaQA、NarrativeQA 三个阅读理解任务从数据预处理、TFRecord 生成到 TPU 微调的全流程实操并附上配置项与关键源码的对照解析。模型概述与仓库结构ReadTwice 仓库对应论文ReadTwice: Reading Very Large Documents with Memories部分函数最初为 ETCEnhanced Transformer Capture同样位于本仓库etcmodel目录项目实现在此基础上演进出二次阅读 全局记忆的机制。仓库核心目录结构如下readtwice/data_utils数据工具包含 Beam 流水线工具、TF 示例构造data_utils.py、SentencePiece 分词tokenization.py及对应测试readtwice/layers模型层实现如侧输入注意力 attention.py、Transformer 主体 transformer.py、嵌入层、重计算梯度、TPU 工具tpu_utils.py等readtwice/models模型主体modeling.py、config.py、优化器、损失函数以及 HotpotQA / TriviaQA / NarrativeQA 三个下游任务的预处理与微调脚本readtwice/run.sh一键脚本创建虚拟环境、安装依赖并运行全部单元测试。依赖安装与环境要求安装命令README 推荐的获取与安装方式如下svn export https://github.com/google-research/google-research/trunk/readtwice pip install -r readtwice/requirements.txt仓库为只读镜像实际使用中也可直接以本仓库readtwice目录为工作目录其依赖清单见 requirements.txt。依赖清单readtwice/requirements.txt主要包括tensorflow1.15.0使用 TF1 兼容 API、absl-py、apache-beam预处理阶段用于并行生成 TFRecord、nltk句子切分与答案标注、numpy、sentencepieceRoBERTa 同款分词、intervaltree、sortedcontainers以及bert-tensorflow。重要硬件前提WARNING当前代码依赖部分自定义算子特别是cross_replica_concat要求TPU或 CPU环境。从源码看该算子位于 readtwice/layers/tpu_utils.py其实现通过xla.replica_id()获取当前 TPU 核编号用tf.scatter_nd把本地张量散落到复制副本维度再经tpu.cross_replica_sum跨核求和、tf.reshape展平从而把分布在num_replicas个 TPU 核上的张量拼接成首维扩大num_replicas倍的完整张量详见 losses.py 中cross_batch_softmax等跨批操作的使用。README 也提示可以通过调整代码使其适配 GPU但需要自行改造。此外还需安装pip install cloud-tpu-client运行单元测试单元测试通过 unittest 发现机制运行前提是当前工作目录必须为readtwice文件夹的父目录即仓库根目录否则模块导入会失败python -m unittest discover -s readtwice -p *_test.py该命令会扫描readtwice下所有*_test.py覆盖数据工具、各网络层、模型、损失、优化器与三个任务的评估代码。readtwice/run.sh则等价地封装了虚拟环境 安装依赖 跑测试的完整流程。预训练模型与分词器官方发布了一个预训练 checkpointreadtwice.tar.gz存放于公共存储桶可用于复现论文实验。模型词汇表与 RoBERTa 完全一致SentencePiece 分词器由本仓库的 bertseq2seq 项目提供vocab_gpt.model。微调前需要准备export PRETRAINED_MODEL_DIRgs://path/to/directory/with/pretrained/model export CONFIG_PATH${PRETRAINED_MODEL_DIR}/read_it_twice_bert_config.json export PRETRAINED_MODEL_CHECKPOINIT${PRETRAINED_MODEL_DIR}/model.ckpt-1000000 export SPM_MODEL_PATH/path/to/vocab_gpt.model export NLTK_DATA_PATH/tmp/nltk_dir其中read_it_twice_bert_config.json是模型配置文件model.ckpt-1000000是预训练 checkpointREADME 原文如此拼写CHECKPOINIT实际为 checkpoint 之意NLTK_DATA_PATH用于 nltk 句子切分所需数据的存放位置。配置文件的完整字段配置类ReadItTwiceBertConfig定义在 readtwice/models/config.py由from_json_file从 JSON 加载、to_json_string序列化。完整字段及含义如下字段默认值说明vocab_size必填token_ids的词汇表大小use_sparse_memory_attention必填是否允许非实体 token 关注基于实体的摘要max_seq_length512token_ids的最大长度max_num_blocks_per_document256单个文档内最大块数即block_pos的最大值cross_attention_pos_emb_modeNone是否基于block_pos添加位置嵌入embedding_sizeNonetoken 嵌入维度None时等于hidden_size同原始 BERT可设小值如 128类似 ALBERThidden_size768编码器与池化层维度num_hidden_layers12Transformer 编码器层数num_attention_heads12注意力头数intermediate_size3072FFN 中间层维度hidden_actgelu激活函数share_kv_projectionsFalse主-主与主-侧注意力是否共享 K/V 投影每层由 2 组 K/V 变为 1 组hidden_dropout_prob0.1全连接层 dropoutattention_probs_dropout_prob0.1注意力概率 dropoutinitializer_range0.02权重初始化截断正态标准差grad_checkpointing_period0激活重计算间隔0表示全部保存大于 0 时以重计算换显存second_read_typefrom_scratch第二次读取方式见下文两次读取详解second_read_num_new_layersNonesecond_read_typenew_layers时新增的 Transformer 层数second_read_num_cross_attention_headsNone第二次读取跨注意力的头数second_read_enable_default_side_inputFalse是否加入默认侧输入类似 no-op 注意力允许注意力权重之和小于 1summary_modecls摘要提取方式如cls、text_block、entity等summary_postprocessing_typenone摘要后处理none/linear/transformersummary_postprocessing_num_layersNone摘要后处理 Transformer 层数cross_attention_top_kNone计算摘要注意力前是否做 Top-K 截断仅支持cross_attend_oncetext_block_extract_every_xNone文本块摘要抽取间隔模型配置加载逻辑get_model_configconfig.py支持三种来源优先读取model_dir下的read_it_twice_bert_config.json若不存在可从source_file或 Base64 编码的source_base64读取并默认把源配置写入模型目录以便后续复用。核心机制两次读取Read-It-Twice架构从 readtwice/models/modeling.py 的实现看ReadItTwiceBertModel的处理流程为第一次读取transformer_with_side_inputsTransformerWithSideInputLayers对块内 token 做标准自注意力同时通过FusedSideAttentionlayers/attention.py把侧输入第一次读取时可无作为额外的 K/V 参与注意力。该层不使用相对注意力是标准 Transformer 层的直接推广因此可以方便地从预训练 BERT/RoBERTa 直接迁移权重。摘要提取SummaryExtraction从第一次读取的隐藏状态中提取块级摘要Summary结构体含states、processed_states、block_ids、block_pos、labels并通过get_cross_block_attmodeling.py依据文档 ID 计算块间注意力掩码——cross_block_attention_mode决定了不同块摘要之间的交互范围block/doc/batch/other_blocks。必要时借助cross_replica_concat在 TPU 多核间汇聚全局摘要实现全局记忆。第二次读取依据second_read_type决定如何利用摘要from_scratch默认把第一次读取的输出从零重新处理将全局摘要的processed_states作为侧输入att_mask_with_side_input同时包含 token-token 掩码与 token-摘要映射掩码modeling.pynew_layers/new_layers_cross_attention在第一次读取结果之上堆叠second_read_num_new_layers个新 Transformer 层cross_attend_once先用SideAttention残差块让 token 一次性关注全局摘要支持cross_attention_top_kTop-K 截断与块位置嵌入再送入新层modeling.py。这种两遍 记忆设计使模型无需把整个长文档硬塞进单一上下文窗口而是以摘要为媒介实现跨块、跨文档的信息传递这正是应对超长文档阅读的核心思路。模型测试 modeling_test.py 中的test_model、test_get_cross_block_att等用例覆盖了不同second_read_type、cross_block_attention_mode与摘要后处理的组合可作为理解各配置行为的参考。下游任务微调实战三个任务的通用流程均为① 下载数据并设置路径 → ② 用preprocess脚本生成 TFRecord 示例并拷贝到 GCS → ③ 用run_finetuning在 TPU 上微调先训练后评估。HotpotQA多跳抽取式问答下载hotpot_train_v1.1.json、hotpot_dev_distractor_v1.json到${HOTPOTQA_DATA_DIR}后export HOTPOTQA_DATA_DIR/path/to/HotpotQA export HOTPOTQA_EXAMPLE_DIR${HOTPOTQA_DATA_DIR}/examples export HOTPOTQA_EXAMPLE_GCP_BUCKETgs://path/to/gcp/bucket export HOTPOTQA_OUTPUT_FOLDERgs://path/to/HotpotQA/output/folder mkdir -p ${HOTPOTQA_EXAMPLE_DIR}生成 TFRecord验证集不生成答案标注训练集加--generate_answerspython -m readtwice.models.hotpot_qa.preprocess \ --spm_model_path${SPM_MODEL_PATH} \ --input_file${HOTPOTQA_DATA_DIR}/hotpot_dev_distractor_v1.json \ --output_prefix${HOTPOTQA_EXAMPLE_DIR}/valid \ --nltk_data_path${NLTK_DATA_PATH} python -m readtwice.models.hotpot_qa.preprocess \ --spm_model_path${SPM_MODEL_PATH} \ --input_file${HOTPOTQA_DATA_DIR}/hotpot_train_v1.1.json \ --output_prefix${HOTPOTQA_EXAMPLE_DIR}/train \ --generate_answers \ --nltk_data_path${NLTK_DATA_PATH} gcloud storage cp ${HOTPOTQA_EXAMPLE_DIR}/* ${HOTPOTQA_EXAMPLE_GCP_BUCKET} gcloud storage cp ${HOTPOTQA_DATA_DIR}/hotpot_dev_distractor_v1.json ${HOTPOTQA_EXAMPLE_GCP_BUCKET}TPU 微调训练阶段python -m readtwice.models.hotpot_qa.run_finetuning \ --read_it_twice_bert_config_file${CONFIG_PATH} \ --input_file${HOTPOTQA_EXAMPLE_GCP_BUCKET}/train.tfrecord-* \ --output_dir${HOTPOTQA_OUTPUT_FOLDER} \ --init_checkpoint${PRETRAINED_MODEL_CHECKPOINIT} \ --enable_side_inputs \ --cross_block_attention_modedoc \ --do_train \ --nodo_eval \ --optimizeradamw \ --learning_rate3e-05 \ --num_train_epochs6 \ --warmup_proportion0.1 \ --learning_rate_scheduleinverse_sqrt \ --poly_power1 \ --start_warmup_step0 \ --save_checkpoints_steps5000 \ --iterations_per_loop1000 \ --nouse_one_hot_embeddings \ --use_tpu \ --tpu_job_name??? \ --num_tpu_cores16 \ --num_tpu_tasks1 \ --decode_top_k40 \ --decode_max_size10 \ --tpu_name??? \ --cross_attention_top_k100评估阶段仅把--do_train改为--nodo_train --do_eval、--nodo_eval改为--do_eval其余参数不变注意 README 中nodo_eval/nodo_train为原文写法其语义即关闭对应开关。WARNING对输出结果的正式评估还需要论文附录中的额外步骤HotpotQA 的完整打分包含 yes/no、支持事实等仓库 hotpot_qa/evaluation.py 与 hotpot_qa/losses.py 分别实现了评估指标与含cross_replica_concat的跨核损失。TriviaQA开放域长文档问答下载 TriviaQA 官方数据wikipedia/web 证据 QA json后export TRIVIAQA_DATA_DIR/path/to/TriviaQA export TRIVIAQA_EXAMPLE_DIR${TRIVIAQA_DATA_DIR}/examples export TRIVIAQA_EXAMPLE_GCP_BUCKETgs://path/to/gcp/bucket export TRIVIAQA_OUTPUT_FOLDERgs://path/to/TriviaQA/output/folder mkdir -p ${TRIVIAQA_EXAMPLE_DIR}生成 TFRecord 时需额外指定证据语料目录python -m readtwice.models.trivia_qa.preprocess \ --spm_model_path${SPM_MODEL_PATH} \ --input_file${TRIVIAQA_DATA_DIR}/qa/wikipedia-dev.json \ --wikipedia_dir${TRIVIAQA_DATA_DIR}/evidence/wikipedia \ --web_dir${TRIVIAQA_DATA_DIR}/evidence/web \ --output_prefix${TRIVIAQA_EXAMPLE_DIR}/valid \ --nltk_data_path${NLTK_DATA_PATH} python -m readtwice.models.trivia_qa.preprocess \ --spm_model_path${SPM_MODEL_PATH} \ --input_file${TRIVIAQA_DATA_DIR}/qa/wikipedia-train.json \ --wikipedia_dir${TRIVIAQA_DATA_DIR}/evidence/wikipedia \ --web_dir${TRIVIAQA_DATA_DIR}/evidence/web \ --output_prefix${TRIVIAQA_EXAMPLE_DIR}/train \ --generate_answers \ --nltk_data_path${NLTK_DATA_PATH} gcloud storage cp ${TRIVIAQA_EXAMPLE_DIR}/* ${TRIVIAQA_EXAMPLE_GCP_BUCKET} gcloud storage cp ${TRIVIAQA_DATA_DIR}/qa/wikipedia-dev.json ${TRIVIAQA_EXAMPLE_GCP_BUCKET}微调命令与 HotpotQA 的差异点学习率更低--learning_rate1e-05、学习率调度为--learning_rate_schedulepoly_decay、--save_checkpoints_steps3000、--iterations_per_loop200、--decode_top_k8 --decode_max_size20并新增--eval_json_path与--eval_data_splitvalidpython -m readtwice.models.trivia_qa.run_finetuning \ --read_it_twice_bert_config_file${CONFIG_PATH} \ --input_file${TRIVIAQA_EXAMPLE_GCP_BUCKET}/train.tfrecord-* \ --eval_json_path${TRIVIAQA_EXAMPLE_GCP_BUCKET}/wikipedia-dev.json \ --output_dir${TRIVIAQA_OUTPUT_FOLDER} \ --init_checkpoint${PRETRAINED_MODEL_CHECKPOINIT} \ --enable_side_inputs \ --cross_block_attention_modedoc \ --do_train \ --nodo_eval \ --optimizeradamw \ --learning_rate1e-05 \ --num_train_epochs6 \ --warmup_proportion0.1 \ --learning_rate_schedulepoly_decay \ --poly_power1 \ --start_warmup_step0 \ --save_checkpoints_steps3000 \ --iterations_per_loop200 \ --nouse_one_hot_embeddings \ --use_tpu \ --tpu_job_name??? \ --num_tpu_cores16 \ --num_tpu_tasks1 \ --decode_top_k8 \ --decode_max_size20 \ --eval_data_splitvalid \ --spm_model_path${SPM_MODEL_PATH} \ --tpu_name??? \ --cross_attention_top_k100评估阶段同样切换为--nodo_train --do_eval。TriviaQA 的评估实现evaluate_triviaqa位于 trivia_qa/evaluation.py包含答案归一化去冠词、标点、下划线、统一大小写与空白以及基于 ground-truth 集合的 EM/F1 计算。NarrativeQA整本故事书阅读从 NarrativeQA 官网下载后NarrativeQA 特殊之处在于复用trivia_qa.run_finetuning入口且预处理需要 qaps问答对与 documents故事文本两份 CSVexport NARRATIVEQA_DATA_DIR/path/to/NarrativeQA export NARRATIVEQA_EXAMPLE_DIR${NARRATIVEQA_DATA_DIR}/examples export NARRATIVEQA_EXAMPLE_GCP_BUCKETgs://path/to/gcp/bucket export NARRATIVEQA_OUTPUT_FOLDERgs://path/to/NarrativeQA/output/folder mkdir -p ${NARRATIVEQA_EXAMPLE_DIR}python -m readtwice.models.narrative_qa.preprocess \ --spm_model_path${SPM_MODEL_PATH} \ --input_qaps${NARRATIVEQA_DATA_DIR}/qaps.csv \ --input_documents${NARRATIVEQA_DATA_DIR}/documents.csv \ --data_splitvalid \ --stories_dir${NARRATIVEQA_DATA_DIR}/tmp/ \ --output_prefix${NARRATIVEQA_EXAMPLE_DIR}/valid \ --nltk_data_path${NLTK_DATA_PATH} python -m readtwice.models.narrative_qa.preprocess \ --spm_model_path${SPM_MODEL_PATH} \ --input_qaps${NARRATIVEQA_DATA_DIR}/qaps.csv \ --input_documents${NARRATIVEQA_DATA_DIR}/documents.csv \ --data_splittrain \ --stories_dir${NARRATIVEQA_DATA_DIR}/tmp/ \ --output_prefix${OUTPUT_NARRATIVE_QA}/train \ --generate_answers \ --nltk_data_path${NLTK_DATA_PATH} gcloud storage cp ${NARRATIVEQA_DATA_DIR}qaps.csv ${NARRATIVEQA_EXAMPLE_DIR}NarrativeQA 的预处理实现了基于 ROUGE-L oracle 的抽取式答案搜索extractive_oracle.py与故事文本解析Gutenberg/电影剧本格式评估逻辑见 narrative_qa/evaluation.py。微调时使用trivia_qa.run_finetuning学习率进一步降至5e-06并关闭默认侧输入python -m readtwice.models.trivia_qa.run_finetuning \ --read_it_twice_bert_config_file${CONFIG_PATH} \ --input_file${NARRATIVEQA_EXAMPLE_DIR}/train.tfrecord-* \ --input_qaps${NARRATIVEQA_EXAMPLE_DIR}/qaps.csv \ --eval_data_splitvalid \ --output_dir${NARRATIVEQA_OUTPUT_FOLDER} \ --init_checkpoint${PRETRAINED_MODEL_CHECKPOINIT} \ --enable_side_inputs \ --cross_block_attention_modedoc \ --do_train \ --nodo_eval \ --optimizeradamw \ --learning_rate5e-06 \ --nosummary_enable_default_side_input \ --num_train_epochs6 \ --warmup_proportion0.1 \ --learning_rate_scheduleinverse_sqrt \ --poly_power1 \ --start_warmup_step0 \ --save_checkpoints_steps15000 \ --spm_model_path${SPM_MODEL_PATH} \ --iterations_per_loop1000 \ --nouse_one_hot_embeddings \ --use_tpu \ --tpu_job_name??? \ --num_tpu_cores16 \ --num_tpu_tasks1 \ --decode_top_k40 \ --decode_max_size10 \ --tpu_name??? \ --cross_attention_top_k100评估阶段切换为--nodo_train --do_eval参数与训练基本一致。微调命令行参数速查综合三个任务的run_finetuning命令核心参数及其作用如下实现见各任务的 run_finetuning.py 与 trivia_qa/run_finetuning.py参数示例值作用--read_it_twice_bert_config_file${CONFIG_PATH}模型架构配置 JSON--input_filegs://.../train.tfrecord-*预处理产物支持 glob--init_checkpoint${PRETRAINED_MODEL_CHECKPOINIT}预训练 checkpoint 初始化--output_dirgs://...checkpoint 与结果输出目录--enable_side_inputs布尔开启 Read-It-Twice 侧输入机制关闭则退化为标准 Transformer--cross_block_attention_modedoc摘要跨块交互范围block/doc/batch/other_blocks--do_train/--do_eval布尔训练 / 评估开关--optimizeradamw优化器adamw或lamblamb_optimizer.py--learning_rate3e-05/1e-05/5e-06初始学习率各任务不同--num_train_epochs6训练轮数--warmup_proportion0.1预热步数占比--learning_rate_scheduleinverse_sqrt/poly_decay学习率调度optimization.py 中实现inverse_sqrt_learning_rate_schedule与多项式衰减--poly_power1多项式衰减幂次--start_warmup_step0预热起始步--save_checkpoints_steps5000/3000/15000保存 checkpoint 间隔--iterations_per_loop1000/200TPU 每次 loop 迭代数--nouse_one_hot_embeddings布尔关闭 one-hot 嵌入查找--use_tpu布尔启用 TPU--tpu_name/--tpu_job_name???TPU 名称与作业名按实际环境填写--num_tpu_cores/--num_tpu_tasks16/1TPU 核数与任务数--decode_top_k40/8解码候选 top-k--decode_max_size10/20解码候选最大数量--cross_attention_top_k100摘要注意力 Top-K 截断--eval_json_path/--eval_data_splitTriviaQA/NarrativeQA评估数据指定训练目标方面微调采用 masked language model 与跨度预测等损失models/losses.py 提供LanguageModelLoss、批量共指消解损失等HotpotQA 另有 yes/no 与 supporting-fact 损失。预训练阶段的核心配置如mlm_fraction_to_mask0.15、mention_mask_mode、mlm_use_whole_word、num_replicas_concat等可在预训练 demo 脚本 run_pretraining_demo.py 中查看。预训练说明与 demoREADME 明确说明预训练代码尚未完整发布主要有两个待解决事项预训练依赖自定义 TF 算子用于在训练过程中动态执行词与实体掩码对应run_pretraining_demo.py中的mention_mask_mode、mlm_use_whole_word等掩码策略以及input_utils.py的mask_same_entity_mentions等函数数据预处理目前依赖内部专有基础设施无法直接开源。作为替代仓库发布了预训练 demorun_pretraining_demo.py该脚本虽不可直接执行但完整展示了核心实现细节包括MLM 与共指消解coreference resolution损失的组合、source_model_config_file/source_model_config_base64两种模型配置注入方式、num_replicas_concat跨副本摘要汇聚、以及cross_block_attention_mode的块交互策略等是理解 ReadTwice 预训练目标与数据流的关键入口。小结ReadTwice 通过两次阅读 块级全局记忆的机制把超长文档阅读理解转化为可并行的块级处理第一次读取生成摘要记忆第二次读取携带记忆精读原文。本文从架构原理ReadItTwiceBertModel与FusedSideAttention、SummaryExtraction的配合、配置字段ReadItTwiceBertConfig全参数、依赖与测试到 HotpotQA / TriviaQA / NarrativeQA 三大任务的数据预处理与 TPU 微调命令完整复现了官方 README 的实操路径并补充了对应的源码级证据。实际复现时请注意环境需 TPU或 CPU、tpu_name等参数需按集群实际填写、评估还需论文附录的额外步骤。赞分享人工智能深度学习NLP计算机视觉强化学习【免费下载链接】google-researchGoogle Research项目地址https://gitcode.com/gh_mirrors/go/google-research点击查看免费下载相关推荐专利通俗解读阅读模式实战指南patent-disclosure-skill 反向阅读专利的完整流水线专利通俗解读阅读模式实战指南patent disclosure skill 反向阅读专利的完整流水线 本指南面向需要 读懂已有专利 而非撰写交底书的用AI 技能AI 插件OpenCore Legacy Patcher终极指南三步让老旧Mac免费吃上最新macOSOpenCore Legacy Patcher终极指南三步让老旧Mac免费吃上最新macOS 想象这样一个画面你书桌角落那台2012年的MacBook Pr操作系统固件驱动开发告别PDF阅读痛点Qwen-Agent的VirtualMemoryAgent如何高效处理超长文档告别PDF阅读痛点Qwen Agent的VirtualMemoryAgent如何高效处理超长文档 你是否还在为处理数百页的PDF文档而烦恼当需要从学术论文、人工智能大模型AI AgentAgent 框架工具调用RAG创作声明:本文部分内容由AI辅助生成(AIGC),仅供参考
锦
锦皓数字建站
深耕本土企业品牌数字化升级,专注原创端正雅致商务官网,从视觉设计到稳定运维全程保驾护航。