[feat] Support Minimind retrieval-augmented generation (RAG)

[ztzhu1]

感谢作者开源这么好的项目，让我收获良多！我基于Minimind实现了检索增强生成 (Retrieval-Augmented Generation, RAG) 功能，让Minimind拥有查阅外部资料后再回答的能力，基本是在Minimind上对 DPR (Karpukhin et al., 2020) 的简易复现。

RAG大致分三步，首先让retriever在大量文档中找到与用户问题相关的top-k条内容，再用reranker对k条内容精细排序，最后将挑选出的外部资料以某种方式喂给语言模型（比如把rank1的文档与用户问题拼接组成新的prompt）。这个PR涉及第一步与第三步，第二步的reranker采用了jina-reranker-v2-base-multilingual。

最后的效果肯定比不上DeepSeek, Qwen等模型，因为参数量和数据规模的限制，即便将正确资料直接提供给Minimind2，它很多时候也不能提取出答案，这我们完全可以理解。这个PR旨在自己动手实现经典算法，并用其为自己训练出的模型插上（稚嫩的）翅膀，非常符合Minimind的哲学。

主要特性

1. 以Minimind2-Pretrain或BERT-Mini为基座模型，训练retriever。 作为decoder-only的模型，Minimind并非实现retriever的最佳选择，因为retriever的主要任务是提取语义而非生成，在此情况下encoder-only的表现更好，Karpukhin et al. (2020) 就使用了BERT。但从下面的结果可以看到即便是decoder-only，Minimind也可以胜任一部分retriever的任务。
2. 实现了 Lucene-BM25 检索算法。 BM25是传统检索方法，通常作为 baseline ，这里自己动手实现的 BM25 没有考虑性能问题，主要是简单清晰地展示其工作方式。
3. 用 RAG 增强 Minimind 的问答能力。 如上所述，当 RAG 模式开启后，会先根据用户的问题检索相关资料，再将资料和问题拼接成新的 prompt 输入模型，这样模型就可以利用外部资料回答。

效果展示

1. Retriever benchmark 。我在 natural-questions (英文) 和 ChineseSquad (中文) 数据集上测试了不同模型的检索效果，SBERT 是 sentence-transformers 的多语言模型（约 118 M 参数），Minimind2-Small-DPR 是我在 Minimind2-Small-Pretrain 的基础上训练出的模型，BERT-Mini-DPR 是我在 BERT-Mini 的基础上训练出的模型（约 11 M 参数）。对比指标为检索准确度，这里的做法并不是很正规，因为精力和资源有限，我没有从互联网获取大量信息作为资料库，而是直接将 test 数据集中的资料拼接成了资料库，只要 query 对应的 context 出现在检索结果的 top-k 中就算成功。Minimind2-Small-Pretrain 基本不能理解英文数据集，在用 DPR 训练后还是能检索出一些数据，而它在理解中文数据方面就比较好了，训练前后区别不大，因此也可以将 Minimind 的预训练模型直接作为简易 retriever。

2. 问答能力。

RAG prompt: 你是一个智能问答助手，请严格按照以下要求回答问题： 如果资料中包含问题答案，请直接使用资料信息并注明"根据资料"，如果资料不相关、信息不足或未包含答案，请明确说明"资料中未包含相关信息"，然后可以基于常识进行补充\n资料：{doc}\n问题：{prompt}\n现在请开始回答：

[A] paraphrase-multilingual-MiniLM-L12-v2 (retriever) + Minimind2-full-sft-768
[B] paraphrase-multilingual-MiniLM-L12-v2 (retriever) + Minimind2-rag-768 (微调 Minimind2-full-sft-768 得到，输出格式更统一)
[C] Minimind2-Small-Pretrain (retriever) + Minimind2-full-sft-768
[D] Minimind2-Small-DPR (retriever) + Minimind2-full-sft-768
[E] without RAG + Minimind2-full-sft-768

  [Q]: 估计有4.88亿至5.35亿人信奉什么宗教
✅[A]: 根据估计，4.88亿至5.35亿人信奉佛教。
✅[B]: 根据资料，4.88亿至5.35亿人信奉佛教，因此估计信奉佛教的宗教是佛教，信奉佛教为多元宗教，信奉佛教为三大宗教。
✅[C]: 根据估计，4.88亿至5.35亿人信奉佛教。
✅[D]: 根据估计，4.88亿至5.35亿人信奉佛教。
❌[E]: 信奉基督教教义，通常基于上帝的恩典，是人类文明进程中不可或缺的一部分。信奉基督教教义，主要基于...

  [Q]: 路德什么时候死的
✅[A]: 路德在1546年凌晨2点45分，他死于中风，享年62岁。
✅[B]: 根据资料，资料中提到"1546年2月18日凌晨2点45分，他死于中风，享年62岁，出生地艾斯勒本。"，因此路德的死亡时间是1546年2月18日凌晨2点45分。
❌[C]: 路德在1983年10月2日，在美国的阿波罗11号任务中，发现了一台具有“时间扭曲”能力的装置，并将其命名为“星云”，这台装置在1984年10月2日被发射成功。这台装置可以将时间扭曲，使时间流逝得更快，因此“星云”被命名为“星云”...
❌[D]: 路德派神学家弗兰兹·皮珀观察到，路德派关于基督徒死后灵魂状态的教导不同于后来路德派神学家约翰·格哈德。莱辛（1755）在这个问题上对路德正统学说的分析中，早就得出了同样的结论。
❌[E]: 这个问题的答案可能因个人或社会的不同而有所差异，但通常来说，路德是19世纪末到20世纪初的一位法国反对宗教改革、支持黑人自由运动的领袖...

  [Q]: 超级碗开幕之夜于何时在何地举行
✅[A]: 超级碗开幕之夜于2016年2月1日在圣何塞SAP中心举行。
✅[B]: 根据资料，资料中提到"2016年2月1日"在圣何塞SAP中心举行，因此在2016年2月1日举行的超级碗开幕之夜，活动于2016年2月1日在圣何塞SAP中心举行。
❌[C]: 超级碗开幕之夜于2023年10月1日在哪里举行？
✅[D]: 超级碗开幕之夜于2016年2月1日在圣何塞SAP中心举行。
❌[E]: 超级碗的开幕之夜通常在每年的12月的第三个星期日举行，具体时间根据地区和活动安排有所不同。例如...

以上问答有一个共同点：问题和资料都比较短，且相关度很高。如果问题比较隐晦，在资料中没有与答案强相关，那么 Minimind 就表现得不太理想，它更像在复述资料中的话，把答案猜出来，而不是真正理解了资料。这使得它难以实现另一个功能（也可能是我能力有限）：发现资料与问题无关。因为复述是简单的，而知道资料与问题无关的前提是要理解两者之间的关系。

代码用法

为了少改动原仓库代码，我将train和eval放在了一个文件，这导致参数比较多，建议用notebook、ipython或在另一个py文件里 import 本文件的方式跑，以下是一些示例

# ===== train =====
train_args = train_dpr.TrainArgs(
    save_weight="dpr",  # 保存时的前缀
    epochs=1,
    batch_size=32,
    accumulation_steps=1,
    learning_rate=4e-6,
    hidden_size=512,
    num_hidden_layers=8,
    max_query_len=64,  # 问题最大长度
    max_passage_len=128,  # 单个document的最大长度，对于`natural-questions`数据集，建议设置为512
    from_weight="pretrain",  # 从MiniMind pretrain 512的权重开始训练
    train_bert=False,  # 是否训练BERT-Mini，若是则加载BERT，则MiniMind模型和权重均失效
    use_wandb=True,
    use_swanlab=False,  # 若use_swanlab=True,需要同时设置use_wandb=True, wandb会被替换为swanlab
    dtype="float16",
    dataset="ChineseSquad",  # 训练数据集（同时也是检索数据集），`ChineseSquad`或`natural-questions`
    wandb_project="MiniMind-DPR",
)
model, tokenizer, train_ds, test_ds = train_dpr.train(train_args, early_return=1) # 加载数据
train_dpr.train(train_args, model, tokenizer, train_ds) # 训练retriever

# ===== eval =====
eval_args = train_dpr.EvalArgs(
    weight="full_sft",  # 加载MiniMind full_sft_768权重
    hidden_size=768,  # chat模型隐藏层维度(与retriever无关)
    num_hidden_layers=16,  # chat模型隐藏层数量(与retriever无关)
    top_k=20,  # 检索时返回的文档数量，reranker会从中选择最终答案
    rag=True,  # 是否开启RAG
    historys=0,
    use_sbert_retriever=True,  # 使用sentence-transformer作为retriever，检索成功率会大幅提升
)

# ----- 下面的代码只有在rag=True时才需要，否则可令retriever=reranker=docs=doc_embeddings=None -----

# 若想使用自行训练的Minimind retriever，可将retriever和tokenizer传入init_rag函数
# 例如: retriever, tokenizer, _, _ = train_dpr.train(train_args, early_return=1)
# 这里设为None是因为用了SBERT retriever
retriever = tokenizer = None
dataset = train_dpr.load_dataset(
    "ChineseSquad",
    max_query_len=eval_args.max_query_len,
    max_passage_len=eval_args.max_passage_len,
)
docs = []
for key in ["train", "validation"]:
    for doc in dataset[key]["context"]:
        if doc not in docs:  # 文档去重
            docs.append(doc)
retriever, reranker, doc_embeddings = train_dpr.init_rag(
    device, docs, retriever, tokenizer, eval_args
)

print(doc_embeddings.shape)
print(docs[0])

eval_args.rag = True # or False
train_dpr.evaluate(eval_args, retriever, reranker, docs, doc_embeddings)