Splade Disco Human Mistral

🚀 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}
}