🚀 DiSCo:用于对话搜索中高效稀疏检索的大语言模型知识蒸馏
DiSCo模型是对原始SPLADE++ (CoCondenser-EnsembleDistil)模型进行对话搜索适配后的产物。它保留了原始的文档编码器,并在专为多轮对话搜索设计的QReCC数据集上对查询编码器进行微调。
训练过程通过从多个教师模型(人类改写和Mistral改写)进行知识蒸馏来完成,这使得模型能够更好地捕捉对话查询的语义。更多详细信息,请参阅原始论文:
- DiSCo SPLADE - SIGIR 2025完整论文:https://arxiv.org/abs/2410.14609
⚠️ 重要提示
这是查询编码器。在进行推理时,你还需要对应的文档编码器,该文档编码器与原始的SPLADE++检查点保持一致。SPLADE可以使用非对称架构,即使用单独的模型来表示查询和文档。
🚀 快速开始
本模型是对原始 SPLADE++ (CoCondenser-EnsembleDistil) 模型进行对话搜索适配后的产物。它保留了原始的文档编码器,并在专为多轮对话搜索设计的 QReCC 数据集上对查询编码器进行微调。
训练过程通过 从多个教师模型(人类改写和 Mistral 改写)进行知识蒸馏 来完成,这使得模型能够更好地捕捉对话查询的语义。更多详细信息,请参阅原始论文:
- DiSCo SPLADE - SIGIR 2025 完整论文:https://arxiv.org/abs/2410.14609
✨ 主要特性
- 对话搜索适配:基于原始 SPLADE++ 模型,针对对话搜索场景进行了优化。
- 知识蒸馏:通过从人类和 Mistral 改写中进行蒸馏,更好地捕捉对话查询的语义。
- 非对称架构:支持使用单独的模型进行查询和文档表示。
💻 使用示例
基础用法
from transformers import AutoTokenizer, AutoModelForMaskedLM
import torch.nn.functional as F
import torch
model = AutoModelForMaskedLM.from_pretrained("slupart/splade-disco-human-mistral")
tokenizer = AutoTokenizer.from_pretrained("slupart/splade-disco-human-mistral")
model.eval()
conv = [
("what's the weather like today?", "it's sunny."),
("should I wear sunscreen?", "yes, UV index is high."),
("do I need sunglasses?", "definitely."),
("where can I buy sunglasses?", "try the optician nearby."),
("how much do they cost?", None)
]
parts = [conv[-1][0]] + [x for q, a in reversed(conv[:-1]) for x in (a, q) if x]
text = " [SEP] ".join(parts)
inputs = tokenizer(text, return_tensors="pt")
with torch.no_grad():
logits = model(**inputs).logits
sparse = F.relu(logits).max(1).values.squeeze(0)
scores = [(tokenizer.convert_ids_to_tokens([i.item()])[0], sparse[i].item())
for i in torch.nonzero(sparse).squeeze(1)]
for token, score in sorted(scores, key=lambda x: -x[1]):
print(f"Token: {token:15} | Score: {score:.4f}")
注意事项
请参考 DiSCo 的 GitHub 仓库以获取完整的使用说明 [github]。输入格式是对话历史的扁平化版本,即 q_n [SEP] a_{n-1} [SEP] q_{n-1} [SEP] ... [SEP] a_0 [SEP] q_0
。
📄 许可证
本模型采用 CC BY-NC-SA 4.0 许可证。
📚 引用
如果您使用了我们的检查点,请引用我们的工作:
@article{lupart2024disco,
title={DiSCo Meets LLMs: A Unified Approach for Sparse Retrieval and Contextual Distillation in Conversational Search},
author={Lupart, Simon and Aliannejadi, Mohammad and Kanoulas, Evangelos},
journal={arXiv preprint arXiv:2410.14609},
year={2024}
}