Bert训练小型RAG分块器 BertChunker: Efficient and Trained Chunking for Retrieval Augmented Generation

项目地址

本项目旨在为RAG训练一个分块器，作为语义分块的一个有效替代。共构造了 50 MB 数据集，在 NVIDIA P40 上训练了10分钟。模型大小 89 MB.

数据集构造

1. 下载英伟达短篇文章数据集。
2. 随机抽取里面的每条数据进行交叉混合，往里插入分块符，形成一条分块数据。
3. 抽取30000次，形成一个数据集。

模型结构

模型采用 sentence-transformers/all-MiniLM-L6-v2 作为基础模型，往上增加一个线性二元分类器，由于对每个token进行分类为分隔符或普通token。超过长度进行滑窗处理分割。

训练

具体训练参数可见项目地址。

性能

1. 随机混合两段不相关的英文片段，让BertChunker预测分隔符的位置。准确率70%，召回率接近100%。需要更大规模实验验证。
2. 对长文本进行处理。准确率69%，召回率70%。需要更大规模实验验证。