multi-qa-MiniLM-L6-dot-v1项目介绍
multi-qa-MiniLM-L6-dot-v1是一个基于sentence-transformers库开发的模型,它将句子和段落映射为384维的稠密向量空间,专为语义搜索设计。该模型在来自多种来源的2.15亿对(问题,答案)上进行了训练。想要了解语义搜索的朋友可以点击这里:SBERT.net - 语义搜索。
使用方法(sentence-transformers)
如果系统中已安装sentence-transformers,使用这个模型变得非常简单:
pip install -U sentence-transformers
然后您可以这样使用模型:
from sentence_transformers import SentenceTransformer, util
query = "伦敦有多少人居住?"
docs = ["大约900万人住在伦敦", "伦敦以其金融区闻名"]
# 加载模型
model = SentenceTransformer('sentence-transformers/multi-qa-MiniLM-L6-dot-v1')
# 编码查询和文档
query_emb = model.encode(query)
doc_emb = model.encode(docs)
# 计算查询和所有文档的点积分数
scores = util.dot_score(query_emb, doc_emb)[0].cpu().tolist()
# 组合文档和分数
doc_score_pairs = list(zip(docs, scores))
# 依据分数降序排序
doc_score_pairs = sorted(doc_score_pairs, key=lambda x: x[1], reverse=True)
# 输出段落和分数
for doc, score in doc_score_pairs:
print(score, doc)
使用方法(HuggingFace Transformers)
若不使用sentence-transformers,可以通过以下方式使用模型:首先通过transformer模型处理输入,然后在词嵌套上应用正确的池化操作。
from transformers import AutoTokenizer, AutoModel
import torch
# CLS Pooling - 从第一个token中提取输出
def cls_pooling(model_output):
return model_output.last_hidden_state[:,0]
# 编码文本
def encode(texts):
# 句子分词
encoded_input = tokenizer(texts, padding=True, truncation=True, return_tensors='pt')
# 计算token嵌入
with torch.no_grad():
model_output = model(**encoded_input, return_dict=True)
# 执行池化
embeddings = cls_pooling(model_output)
return embeddings
# 我们希望对其进行句子嵌入的句子
query = "伦敦有多少人居住?"
docs = ["大约900万人住在伦敦", "伦敦以其金融区闻名"]
# 从HuggingFace Hub中加载模型
tokenizer = AutoTokenizer.from_pretrained("sentence-transformers/multi-qa-MiniLM-L6-dot-v1")
model = AutoModel.from_pretrained("sentence-transformers/multi-qa-MiniLM-L6-dot-v1")
# 编码查询和文档
query_emb = encode(query)
doc_emb = encode(docs)
# 计算查询和所有文档的点积分数
scores = torch.mm(query_emb, doc_emb.transpose(0, 1))[0].cpu().tolist()
# 组合文档和分数
doc_score_pairs = list(zip(docs, scores))
# 依据分数降序排序
doc_score_pairs = sorted(doc_score_pairs, key=lambda x: x[1], reverse=True)
# 输出段落和分数
for doc, score in doc_score_pairs:
print(score, doc)
技术细节
在以下是如何使用此模型的一些技术细节:
设置 | 值 |
---|---|
维度 | 384 |
生成归一化的嵌入 | 否 |
池化方法 | CLS pooling |
适合的分数函数 | 点积(例如 util.dot_score ) |
背景
这个项目旨在利用自监督对比学习目标对大规模句子级别的数据集进行句子嵌入模型的训练。其目标是在一组随机采样的其他句子中预测实际配对的句子。
模型是在Hugging Face组织的JAX/Flax词汇语音处理(NLP)和计算机视觉(CV)的社区周期间开发的。项目也受益于高效的硬件基础设施:7个TPU v3-8,以及Google Flax、JAX和Cloud团队成员关于高效深度学习框架的支持。
预期用途
multi-qa-MiniLM-L6-dot-v1模型预期用于语义搜索:它在密集向量空间中对查询/问题和文本段落进行编码,可用于找到给定段落的相关文档。
请注意有512个词片的输入限制:超过该长度的文本将被截断。此外,模型仅在250个词片及以下的输入文本上进行训练,对于更长的文本可能效果不佳。
训练过程
完整训练脚本可以在当前存储库train_script.py
中找到。
预训练
我们使用预训练的nreimers/MiniLM-L6-H384-uncased
模型进行预训练。欲知更多关于预训练的详细信息,请查阅模型卡片。
训练
我们通过多个数据集的拼接来微调模型。总体上,我们有大约2.15亿(问题,答案)对。我们根据data_config.json
文件中详细配置的加权概率抽样每个数据集。
模型使用了MultipleNegativesRankingLoss进行训练,使用CLS池化,点积作为相似度函数,尺度为1。
训练中使用的数据集包括WikiAnswers、PAQ、Stack Exchange、MS MARCO等,其子集如问题与答案的对组成不同数量的训练数据。数据来源非常广泛,确保了模型的鲁棒性和通用性。