本案例展示了如何使用NUDGE方法微调语料库嵌入以提高使用预训练嵌入模型进行语义相似性检索文本的准确性。1. 案例目标本案例的主要目标是介绍NUDGENon-parametric Updates for Data-Driven Gradient Enhancement微调方法这是一种新颖、简单且轻量级的微调方法展示如何在常用的信息检索基准Scifact上应用NUDGE方法比较NUDGE微调、适配器微调与基线模型在NDCGk指标上的性能差异演示如何向现有数据集插入新记录并进行增量微调2. 技术栈与核心依赖llama-index-experimental- 提供NUDGE实现llama-index-embeddings-huggingface- 用于加载Hugging Face嵌入模型nudge-ft- NUDGE微调方法的核心库torch- PyTorch深度学习框架datasets- Hugging Face数据集库3. 环境配置安装必要的依赖%pip install llama-index-experimental llama-index-embeddings-huggingface nudge-ft torch datasets4. 案例实现4.1 加载数据集首先我们定义一个函数来加载Hugging Face数据集from llama_index.finetuning import EmbeddingQAFinetuneDataset from datasets import load_dataset def load_hf_dataset(dataset_name): hf_dataset_name fsepz/{dataset_name}_ft corpus load_dataset(hf_dataset_name, data_records, splittrain) queries_train load_dataset(hf_dataset_name, qs, splittrain) queries_validation load_dataset(hf_dataset_name, qs, splitdev) queries_test load_dataset(hf_dataset_name, qs, splittest) qrels_train load_dataset(hf_dataset_name, qs_rel, splittrain) qrels_validation load_dataset(hf_dataset_name, qs_rel, splitdev) qrels_test load_dataset(hf_dataset_name, qs_rel, splittest) # 数据处理和转换 corpus { str(corpus[i][record_id]): corpus[i][text] for i in range(len(corpus)) } # 查询和相关性处理 queries_train { str(queries_train[i][q_id]): queries_train[i][input] for i in range(len(queries_train)) } # ... 其他查询集处理 # 相关文档处理 qrels_train ( qrels_train.to_pandas() .groupby(q_id)[record_id] .apply(list) .to_dict() ) # ... 其他相关文档处理 # 创建数据集对象 train_dataset EmbeddingQAFinetuneDataset( corpuscorpus, queriesqueries_train, relevant_docsqrels_train ) validation_dataset EmbeddingQAFinetuneDataset( corpuscorpus, queriesqueries_validation, relevant_docsqrels_validation, ) test_dataset EmbeddingQAFinetuneDataset( corpuscorpus, queriesqueries_test, relevant_docsqrels_test ) return train_dataset, validation_dataset, test_dataset4.2 加载基础嵌入模型加载BAAI/bge-small-en-v1.5作为基础嵌入模型from llama_index.core.embeddings import resolve_embed_model train_dataset, val_dataset, test_dataset load_hf_dataset(scifact) base_embed_model resolve_embed_model(local:BAAI/bge-small-en-v1.5)4.3 评估指标实现实现NDCGk评估指标def ndcg_at_k( dataset: EmbeddingQAFinetuneDataset, retriever: BaseRetriever, k: int 10 ): queries dataset.queries relevant_docs dataset.relevant_docs ndcg_scores [] for query_id, query in tqdm(queries.items()): retrieved_nodes retriever.retrieve(query) retrieved_ids [node.node.node_id for node in retrieved_nodes] expected_ids relevant_docs[query_id] # 计算NDCG ideal_dcg np.sum( [1 / np.log2(i 2) for i in range(min(k, len(expected_ids)))] ) rel_scores np.zeros(k) for j in range(min(k, len(retrieved_ids))): if retrieved_ids[j] in expected_ids: rel_scores[j] 1 dcg np.sum( [rel_scores[i] / np.log2(i 2) for i in range(len(rel_scores))] ) ndcg dcg / ideal_dcg if ideal_dcg 0 else 0 ndcg_scores.append(ndcg) mean_ndcg np.mean(ndcg_scores) return mean_ndcg4.4 NUDGE微调使用NUDGE方法微调语料库嵌入from llama_index.experimental import Nudge k 10 nudge Nudge( train_datasettrain_dataset, val_datasetval_dataset, embed_modelbase_embed_model, use_nudge_nTrue, ) nudge.finetune() nudge_corpus_embeddings nudge.get_finetuned_corpus_embeddings() nudge_retriever build_retriever( train_dataset.corpus, base_embed_model, nudge_corpus_embeddings, kk ) nudge_ndcg_test ndcg_at_k(test_dataset, nudge_retriever, k)4.5 适配器微调使用适配器方法微调嵌入模型from llama_index.finetuning import EmbeddingAdapterFinetuneEngine embedding_adapater_finetune_engine EmbeddingAdapterFinetuneEngine( train_dataset, base_embed_model, epochs4, batch_size10, ) embedding_adapater_finetune_engine.finetune() embedding_adapter_model ( embedding_adapater_finetune_engine.get_finetuned_model() ) ft_retriever build_retriever( train_dataset.corpus, embedding_adapter_model, kk ) ft_ndcg_test ndcg_at_k(test_dataset, ft_retriever, k)4.6 基线评估评估基础模型性能base_retriever build_retriever(train_dataset.corpus, base_embed_model, kk) bge_ndcg_test ndcg_at_k(test_dataset, base_retriever, k)5. 案例效果在Scifact数据集上的测试结果bge test - ndcg10: 0.71 adaptor finetune test - ndcg10: 0.72 NUDGE-N test - ndcg10: 0.87从结果可以看出NUDGE-N在NDCG10指标上达到了0.87显著优于基础模型(0.71)和适配器微调(0.72)。6. 案例实现思路6.1 NUDGE方法原理NUDGE是一种非参数微调方法它直接修改数据记录的嵌入而不是修改模型参数。具体来说NUDGE直接修改嵌入本身而不是通过修改模型参数或训练适配器来生成更好的嵌入与微调预训练模型和训练适配器相比NUDGE提供了3.3倍和4.3倍的更高精度提升并且运行速度快200倍和3倍NUDGE-N是NUDGE的变体使用归一化技术进一步提高性能6.2 数据集扩展案例还演示了如何向现有数据集插入新记录并进行增量微调new_train_dataset, new_val_dataset, new_test_dataset load_hf_dataset( nfcorpus ) # 为键添加前缀以避免与scifact ID冲突 new_train_dataset.queries { fnfcorpus-{k}: v for k, v in new_train_dataset.queries.items() }使用NUDGE可以轻松地迭代扩展数据集只需关注最新批次的数据而不用担心灾难性遗忘。这仅在插入的新数据与现有数据不冲突时才有效。7. 扩展建议自定义数据集可以使用自己的文档和查询数据集只需按照案例中的格式组织数据其他评估指标除了NDCGk还可以尝试其他信息检索指标如MRR、MAP等不同嵌入模型尝试不同的基础嵌入模型如更大的bge模型或其他开源嵌入模型参数调优调整NUDGE和适配器微调的参数如学习率、批大小等多数据集融合尝试将多个数据集融合构建更全面的检索系统增量学习策略研究更复杂的数据集扩展策略处理新旧数据间的潜在冲突8. 总结本案例展示了NUDGE方法在微调语料库嵌入方面的强大能力。与传统的模型微调或适配器微调相比NUDGE直接修改嵌入向量实现了更高的检索精度和更快的训练速度。在Scifact数据集上NUDGE-N在NDCG10指标上达到了0.87显著优于基础模型(0.71)和适配器微调(0.72)。此外NUDGE还支持增量学习可以轻松地向现有数据集添加新记录而无需重新训练整个模型这为构建可扩展的检索系统提供了便利。