SGPT: GPT 句子嵌入用于语义搜索
此代码库包含论文 SGPT: GPT Sentence Embeddings for Semantic Search 的代码、结果和预训练模型。
**************************** 更新 ****************************
- 2024-02: 我们发布了 GRIT & GritLM - 这些模型将 SGPT 双编码器、交叉编码器、对称、不对称和常规 GPT(即生成)统一在单一模型中,且在所有方面表现更好。我们建议切换到这些新模型 :)
- 2022-09: 使用 Sentence Transformers 现在更容易使用 SGPT 双编码器,参见新脚本
- 2022-08: 现在发布了多语言 BLOOM SGPT 模型:不对称,7.1B 参数 和 对称,1.7B 参数。如果需要不同的模型,请随时提出问题。
- 2022-06: OpenAI 发布了我们在论文中与 SGPT 交叉编码器进行比较的其搜索端点的机制。我们的方法非常相似。可以在
crossencoder/beir/openai_search_endpoint_functionality.py
中测试他们的提示! - 2022-03: 5.8B 双编码器模型在USEB和BEIR上分别提高了4%和1%。论文 和 模型 已在 HF 上进行了更新。这是通过使用更大的批量大小和 GradCache 完成的,详情请参见论文。如果您之前下载过它们,建议替换为新版本。
- 2022-02: 我们发布了我们的论文。来看看吧! :)
快速链接
概览
我们展示了 SGPT-BE 和 SGPT-CE,用于将 GPT 模型应用为双编码器或交叉编码器以进行对称或不对称搜索。SGPT-BE 通过仅对偏置张量进行对比微调和位置加权平均池化来生成语义上有意义的句子嵌入。SGPT-CE 使用 GPT 模型的对数概率而不进行任何微调。以下是方法的示意图。
如果有任何问题,请随时提交问题~
结构
.
├── biencoder # 双编码器的训练与推理
│ ├── beir
│ │ ├── custommodels # 提供 BEIR 兼容性用于不对称模型和具有特殊标记的模型
│ │ │ └── ...
│ │ ├── io_utils # 专用于 beir_openai_embeddings_batched_parallel.py
│ │ │ └── ...
│ │ ├── parallelizer # 专用于 beir_openai_embeddings_batched_parallel.py
│ │ │ └── ...
│ │ ├── beir_dense_retriever.py
│ │ ├── beir_openai_embeddings_batched_parallel.py
│ │ ├── requirements.txt
│ │ ├── *.bash # 运行多个实验的 Bash 脚本
│ │ └── README.md
│ ├── nli_msmarco
│ │ ├── sentence-transformers # 句子转换器的改编版本 - 安装此版本用于所有双编码器实验
│ │ │ └── ...
│ │ └── README.md
│ └── useb
│ ├── useb
│ │ └── ...
│ ├── *.bash # 运行多个实验的 Bash 脚本
│ ├── useb_dense_retriever.py
│ └── README.md
├── crossencoder # 交叉编码器的推理
│ └── beir
│ ├── *.ipynb # README 中解释的笔记
│ └── README.md
├── other
│ ├── sgpt_graphic.png
│ └── sgpt_utils.ipynb # 用于创建论文中的图表及其他代码
├── requirements.txt
└── README.md
每个数据子目录提供其结构概览以及用于生成数据集、模型和其他内容的结构、下载(数据集、模型)和命令。通常可以在 https://huggingface.co/Muennighoff 找到所有模型,在 https://www.kaggle.com/muennighoff/datasets 中找到各种数据集的 json 结果。模型名称在 Huggingface 的 README 中进行了说明。数据集名称在此代码库的子文件夹中进行了说明。
使用 Huggingface 的 SGPT
以下我们提供了使用预训练模型进行语义搜索的 Python 示例。
我们强烈推荐将模型名称替换为更大的模型,例如用于对称双编码器的 Muennighoff/SGPT-5.8B-weightedmean-nli-bitfit
。
双编码器
对称语义搜索 BE
import torch
from transformers import AutoModel, AutoTokenizer
from scipy.spatial.distance import cosine
# 获取我们的模型 - 该包将自动下载模型
# 为获得最佳性能:Muennighoff/SGPT-5.8B-weightedmean-nli-bitfit
tokenizer = AutoTokenizer.from_pretrained("Muennighoff/SGPT-125M-weightedmean-nli-bitfit")
model = AutoModel.from_pretrained("Muennighoff/SGPT-125M-weightedmean-nli-bitfit")
# 禁用 Dropout(上述模型中没有 Dropout,所以这里没有差别,但其他 SGPT 模型可能有 Dropout)
model.eval()
# 标记输入文本
texts = [
"深度学习",
"人工智能",
"深度潜水",
"人工造雪"
]
batch_tokens = tokenizer(texts, padding=True, truncation=True, return_tensors="pt")
# 获取嵌入
with torch.no_grad():
# 获取形状为 [bs, seq_len, hid_dim] 的隐藏状态
last_hidden_state = model(**batch_tokens, output_hidden_states=True, return_dict=True).last_hidden_state
# 获取形状为 [bs, seq_len, hid_dim] 的权重
weights = (
torch.arange(start=1, end=last_hidden_state.shape[1] + 1)
.unsqueeze(0)
.unsqueeze(-1)
.expand(last_hidden_state.size())
.float().to(last_hidden_state.device)
)
# 获取形状为 [bs, seq_len, hid_dim] 的注意力掩码
input_mask_expanded = (
batch_tokens["attention_mask"]
.unsqueeze(-1)
.expand(last_hidden_state.size())
.float()
)
# 在 seq_len 上执行加权平均池化:bs, seq_len, hidden_dim -> bs, hidden_dim
sum_embeddings = torch.sum(last_hidden_state * input_mask_expanded * weights, dim=1)
sum_mask = torch.sum(input_mask_expanded * weights, dim=1)
embeddings = sum_embeddings / sum_mask
# 计算余弦相似度
# 余弦相似度在 [-1, 1] 之间。值越高表示越相似
cosine_sim_0_1 = 1 - cosine(embeddings[0], embeddings[1])
cosine_sim_0_2 = 1 - cosine(embeddings[0], embeddings[2])
cosine_sim_0_3 = 1 - cosine(embeddings[0], embeddings[3])
print("“%s” 和 “%s” 之间的余弦相似度为:%.3f" % (texts[0], texts[1], cosine_sim_0_1))
print("“%s” 和 “%s” 之间的余弦相似度为:%.3f" % (texts[0], texts[2], cosine_sim_0_2))
print("“%s” 和 “%s” 之间的余弦相似度为:%.3f" % (texts[0], texts[3], cosine_sim_0_3))
不对称语义搜索 BE
import torch
from transformers import AutoModel, AutoTokenizer
from scipy.spatial.distance import cosine
# 获取我们的模型 - 该包将自动下载模型
# 为获得最佳性能:Muennighoff/SGPT-5.8B-weightedmean-msmarco-specb-bitfit
tokenizer = AutoTokenizer.from_pretrained("Muennighoff/SGPT-125M-weightedmean-msmarco-specb-bitfit")
model = AutoModel.from_pretrained("Muennighoff/SGPT-125M-weightedmean-msmarco-specb-bitfit")
# 禁用 Dropout(上述模型中没有 Dropout,所以这里没有差别,但其他 SGPT 模型可能有 Dropout)
model.eval()
queries = [
"我正在寻找一个离地球不太远的星球。",
]
docs = [
"海王星是离太阳最远的第八颗行星。在太阳系中,它是按直径计算的第四大行星,也是质量第三大的行星,且是密度最大的巨行星。它的质量是地球的17倍,略微大于其近似双胞胎天王星。",
"TRAPPIST-1d,也称为2MASS J23062928-0502285 d,是一颗小型系外行星(质量约为地球的30%),它在超冷矮星TRAPPIST-1的宜居带内轨道上运行,距离地球约40光年(12.1秒差距,或约3.7336×10^14公里),位于水瓶座星座中。",
"塔图因星是一个围绕银河系外环的双太阳运行的严酷沙漠世界,这是一个由赫特帮派统治的无法律之地。许多定居者在潮湿农场上勉强度日,而像莫斯艾斯利和莫斯艾斯帕这样的太空港城市则成为走私者、罪犯和其他流氓的根据地。",
]
SPECB_QUE_BOS = tokenizer.encode("[", add_special_tokens=False)[0]
SPECB_QUE_EOS = tokenizer.encode("]", add_special_tokens=False)[0]
SPECB_DOC_BOS = tokenizer.encode("{", add_special_tokens=False)[0]
SPECB_DOC_EOS = tokenizer.encode("}", add_special_tokens=False)[0]
def tokenize_with_specb(texts, is_query):
# Tokenize without padding
batch_tokens = tokenizer(texts, padding=False, truncation=True)
# Add special brackets & pay attention to them
for seq, att in zip(batch_tokens["input_ids"], batch_tokens["attention_mask"]):
if is_query:
seq.insert(0, SPECB_QUE_BOS)
seq.append(SPECB_QUE_EOS)
else:
seq.insert(0, SPECB_DOC_BOS)
seq.append(SPECB_DOC_EOS)
att.insert(0, 1)
att.append(1)
# Add padding
batch_tokens = tokenizer.pad(batch_tokens, padding=True, return_tensors="pt")
return batch_tokens
def get_weightedmean_embedding(batch_tokens, model):
# Get the embeddings
with torch.no_grad():
# Get hidden state of shape [bs, seq_len, hid_dim]
last_hidden_state = model(**batch_tokens, output_hidden_states=True, return_dict=True).last_hidden_state
# Get weights of shape [bs, seq_len, hid_dim]
weights = (
torch.arange(start=1, end=last_hidden_state.shape[1] + 1)
.unsqueeze(0)
.unsqueeze(-1)
.expand(last_hidden_state.size())
.float().to(last_hidden_state.device)
)
# Get attn mask of shape [bs, seq_len, hid_dim]
input_mask_expanded = (
batch_tokens["attention_mask"]
.unsqueeze(-1)
.expand(last_hidden_state.size())
.float()
)
# Perform weighted mean pooling across seq_len: bs, seq_len, hidden_dim -> bs, hidden_dim
sum_embeddings = torch.sum(last_hidden_state * input_mask_expanded * weights, dim=1)
sum_mask = torch.sum(input_mask_expanded * weights, dim=1)
embeddings = sum_embeddings / sum_mask
return embeddings
query_embeddings = get_weightedmean_embedding(tokenize_with_specb(queries, is_query=True), model)
doc_embeddings = get_weightedmean_embedding(tokenize_with_specb(docs, is_query=False), model)
# Calculate cosine similarities
# Cosine similarities are in [-1, 1]. Higher means more similar
cosine_sim_0_1 = 1 - cosine(query_embeddings[0], doc_embeddings[0])
cosine_sim_0_2 = 1 - cosine(query_embeddings[0], doc_embeddings[1])
cosine_sim_0_3 = 1 - cosine(query_embeddings[0], doc_embeddings[2])
print("查询\"%s\" 与文档 \"%s\" 的余弦相似度为: %.3f" % (queries[0], docs[0][:20] + "...", cosine_sim_0_1))
print("查询\"%s\" 与文档 \"%s\" 的余弦相似度为: %.3f" % (queries[0], docs[1][:20] + "...", cosine_sim_0_2))
print("查询\"%s\" 与文档 \"%s\" 的余弦相似度为: %.3f" % (queries[0], docs[2][:20] + "...", cosine_sim_0_3))
交叉编码器
非对称语义搜索CE
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
from scipy.spatial.distance import cosine
# 获得模型 - 该包会自动下载模型
# 为了获得最佳性能:EleutherAI/gpt-j-6B
tokenizer = AutoTokenizer.from_pretrained("EleutherAI/gpt-neo-125M")
model = AutoModelForCausalLM.from_pretrained("EleutherAI/gpt-neo-125M")
# 禁用Dropout(上述模型中没有dropout,所以这里没有任何区别,但其他SGPT模型可能有dropout)
model.eval()
prompt = '文档正在搜索内容相同的匹配项。\n文档 "{}" 是搜索 "'
queries = [
"我正在寻找一颗离地球不太远的行星。",
]
docs = [
"海王星是离太阳最远的第八颗行星。在太阳系中,它是按直径计算的第四大行星,也是质量第三大的行星,且是密度最大的巨行星。它的质量是地球的17倍,略微大于其近似双胞胎天王星。",
"TRAPPIST-1d,也称为2MASS J23062928-0502285 d,是一颗小型系外行星(质量约为地球的30%),它在超冷矮星TRAPPIST-1的宜居带内轨道上运行,距离地球约40光年(12.1秒差距,或约3.7336×10^14公里),位于水瓶座星座中。",
"塔图因星是一个围绕银河系外环的双太阳运行的严酷沙漠世界,这是一个由赫特帮派统治的无法律之地。许多定居者在潮湿农场上勉强度日,而像莫斯艾斯利和莫斯艾斯帕这样的太空港城市则成为走私者、罪犯和其他流氓的根据地。",
]
for query in queries:
print(f"查询: {query}")
for doc in docs:
context = prompt.format(doc)
context_enc = tokenizer.encode(context, add_special_tokens=False)
continuation_enc = tokenizer.encode(query, add_special_tokens=False)
# 切掉最后一个token,因为我们使用之前的token的概率
model_input = torch.tensor(context_enc+continuation_enc[:-1])
continuation_len = len(continuation_enc)
input_len, = model_input.shape
# [seq_len] -> [seq_len, vocab]
logprobs = torch.nn.functional.log_softmax(model(model_input)[0], dim=-1).cpu()
# [seq_len, vocab] -> [continuation_len, vocab]
logprobs = logprobs[input_len-continuation_len:]
# 收集续词token的log概率 -> [continuation_len]
logprobs = torch.gather(logprobs, 1, torch.tensor(continuation_enc).unsqueeze(-1)).squeeze(-1)
score = torch.sum(logprobs)
# 值越高(越接近0),则相似度越高
print(f"文档: {doc[:20] + '...'} 分数: {score}")
对称语义搜索CE
你可以使用与上面CE-Asym部分相同的代码,但需要更改提示词。欢迎分享效果好的提示词 :)
使用 Sentence Transformers 的 SGPT
双编码 ST
对称语义搜索 BE ST
对称模型现在可以通过 pip install git+https://github.com/UKPLab/sentence-transformers.git
与最新的 sentence-transformers 100% 兼容。你应该能获得与 上面的HuggingFace脚本 相同的结果。
from scipy.spatial.distance import cosine
from sentence_transformers import SentenceTransformer
texts = [
"深度学习",
"人工智能",
"深潜",
"人工雪",
]
model = SentenceTransformer("Muennighoff/SGPT-125M-weightedmean-nli-bitfit")
embeddings = model.encode(texts)
cosine_sim_0_1 = 1 - cosine(embeddings[0], embeddings[1])
cosine_sim_0_2 = 1 - cosine(embeddings[0], embeddings[2])
cosine_sim_0_3 = 1 - cosine(embeddings[0], embeddings[3])
print("深度学习\"%s\" 与 人工智能 \"%s\" 的余弦相似度为: %.3f" % (texts[0], texts[1], cosine_sim_0_1))
print("深度学习\"%s\" 与 深潜 \"%s\" 的余弦相似度为: %.3f" % (texts[0], texts[2], cosine_sim_0_2))
print("深度学习\"%s\" 与 人工雪 \"%s\" 的余弦相似度为: %.3f" % (texts[0], texts[3], cosine_sim_0_3))
非对称语义搜索 BE ST
SGPT Sentence Transformers
安装:pip install --upgrade git+https://github.com/Muennighoff/sentence-transformers.git@sgpt_poolings_specb
使用以下代码,这将产生与 上面的HuggingFace解决方案 相同的结果。
from scipy.spatial.distance import cosine
from sentence_transformers import SentenceTransformer
queries = [
"我正在寻找一颗离地球不太远的行星。",
]
docs = [
"海王星是距离太阳最远的第八颗已知太阳系行星。在太阳系中,它的直径是第四大的行星,质量是第三大的行星,也是密度最高的巨行星。它的质量是地球的17倍,比其近似的双胞胎天王星稍大。",
"TRAPPIST-1d,也被称为2MASS J23062928-0502285 d,是一颗小型系外行星(约为地球质量的30%),它围绕超冷矮星TRAPPIST-1在适居带的内边缘运行,距离地球大约40光年(12.1秒差距,或接近3.7336×10¹⁴公里),位于宝瓶座。",
"塔图因是一个围绕银河系外缘双星运行的严酷沙漠世界,是一个由赫特帮派统治的无法无天之地。许多定居者在湿气农场勉强维生,而像莫斯艾斯利和莫斯埃斯帕这样的太空港城市则成为走私者、罪犯和其他流氓的基地。",
]
class SentenceTransformerSpecb(SentenceTransformer):
# 需要:
# pip install git+https://github.com/Muennighoff/sentence-transformers.git@sgpt_poolings_specb
def __init__(self, *args, **kwargs):
super().__init__(*args, **kwargs)
tokens = ["[SOS]", "{SOS}"]
self._first_module().tokenizer.add_tokens(tokens, special_tokens=True)
self._first_module().auto_model.resize_token_embeddings(len(self._first_module().tokenizer))
# 将在模型中被替换为表示标记
# 问题是我们在Transformer.py模块中进行标记化时不知道文本是查询还是文档,
# 因此我们使用SOS标记作为标识符来确定手头是查询还是文档,然后再替换它们
# 如果我们直接在这里使用括号,它们可能会成为另一个标记的一部分
self._first_module().bos_spec_token_q = self._first_module().tokenizer.encode("[SOS]", add_special_tokens=False)[0]
self._first_module().bos_spec_token_d = self._first_module().tokenizer.encode("{SOS}", add_special_tokens=False)[0]
self._first_module().bos_spec_token_q_rep = self._first_module().tokenizer.encode("[", add_special_tokens=False)[0]
self._first_module().eos_spec_token_q = self._first_module().tokenizer.encode("]", add_special_tokens=False)[0]
self._first_module().bos_spec_token_d_rep = self._first_module().tokenizer.encode("{", add_special_tokens=False)[0]
self._first_module().eos_spec_token_d = self._first_module().tokenizer.encode("}", add_special_tokens=False)[0]
self._first_module().replace_bos = True
def encode(self, sentences, **kwargs):
is_query = kwargs.pop("is_query", True)
if is_query:
sentences = "[SOS]" + sentences if isinstance(sentences, str) else ["[SOS]" + sent for sent in sentences]
else:
sentences = "{SOS}" + sentences if isinstance(sentences, str) else ["{SOS}" + sent for sent in sentences]
return super().encode(sentences, **kwargs)
model = SentenceTransformerSpecb("Muennighoff/SGPT-125M-weightedmean-msmarco-specb-bitfit")
query_embeddings = model.encode(queries, is_query=True)
doc_embeddings = model.encode(docs, is_query=False)
# 计算余弦相似度
# 余弦相似度在[-1, 1]之间,值越高表示越相似
cosine_sim_0_1 = 1 - cosine(query_embeddings[0], doc_embeddings[0])
cosine_sim_0_2 = 1 - cosine(query_embeddings[0], doc_embeddings[1])
cosine_sim_0_3 = 1 - cosine(query_embeddings[0], doc_embeddings[2])
print("查询\"%s\"与文档\"%s\"之间的余弦相似度为: %.3f" % (queries[0], docs[0][:20] + "...", cosine_sim_0_1))
print("查询\"%s\"与文档\"%s\"之间的余弦相似度为: %.3f" % (queries[0], docs[1][:20] + "...", cosine_sim_0_2))
print("查询\"%s\"与文档\"%s\"之间的余弦相似度为: %.3f" % (queries[0], docs[2][:20] + "...", cosine_sim_0_3))
原始的Sentence Transformers
如果你想使用 https://github.com/UKPLab/sentence-transformers
中的Sentence Transformers,你可以使用以下代码。确保使用最新版本(pip install --upgrade git+https://github.com/UKPLab/sentence-transformers.git
)。
注意这会产生略微差的评分,因为特殊括号在标记化时可能会与其他标记交织。在SciFact (BEIR) NDCG@10上的评分从SGPT-125M-weightedmean-msmarco-specb-bitfit
的0.569下降到0.566。
from scipy.spatial.distance import cosine
from sentence_transformers import SentenceTransformer
queries = [
"我正在寻找一颗离地球不太远的行星。",
]
docs = [
"海王星是距离太阳最远的第八颗已知太阳系行星。在太阳系中,它的直径是第四大的行星,质量是第三大的行星,也是密度最高的巨行星。它的质量是地球的17倍,比其近似的双胞胎天王星稍大。",
"TRAPPIST-1d,也被称为2MASS J23062928-0502285 d,是一颗小型系外行星(约为地球质量的30%),它围绕超冷矮星TRAPPIST-1在适居带的内边缘运行,距离地球大约40光年(12.1秒差距,或接近3.7336×10¹⁴公里),位于宝瓶座。",
"塔图因是一个围绕银河系外缘双星运行的严酷沙漠世界,是一个由赫特帮派统治的无法无天之地。许多定居者在湿气农场勉强维生,而像莫斯艾斯利和莫斯埃斯帕这样的太空港城市则成为走私者、罪犯和其他流氓的基地。",
]
class SentenceTransformerSpecb(SentenceTransformer):
def encode(self, sentences, **kwargs):
is_query = kwargs.pop("is_query", True)
if is_query:
sentences = "[" + sentences + "]" if isinstance(sentences, str) else ["[" + sent + "]" for sent in sentences]
else:
sentences = "{" + sentences + "}" if isinstance(sentences, str) else ["{" + sent + "}" for sent in sentences]
return super().encode(sentences, **kwargs)
model = SentenceTransformerSpecb("Muennighoff/SGPT-125M-weightedmean-msmarco-specb-bitfit")
query_embeddings = model.encode(queries, is_query=True)
doc_embeddings = model.encode(docs, is_query=False)
# 计算余弦相似度
# 余弦相似度在[-1, 1]之间,值越高表示越相似
cosine_sim_0_1 = 1 - cosine(query_embeddings[0], doc_embeddings[0])
cosine_sim_0_2 = 1 - cosine(query_embeddings[0], doc_embeddings[1])
cosine_sim_0_3 = 1 - cosine(query_embeddings[0], doc_embeddings[2])
print("查询\"%s\"与文档\"%s\"之间的余弦相似度为: %.3f" % (queries[0], docs[0][:20] + "...", cosine_sim_0_1))
print("查询\"%s\"与文档\"%s\"之间的余弦相似度为: %.3f" % (queries[0], docs[1][:20] + "...", cosine_sim_0_2))
print("查询\"%s\"与文档\"%s\"之间的余弦相似度为: %.3f" % (queries[0], docs[2][:20] + "...", cosine_sim_0_3))
致谢
感谢Constantin Eichenberg和Samuel Weinbach在整个项目中提供的富有洞察力的讨论和宝贵的反馈。感谢Robert Baldock、Marco Bellagente和Koen Oostermeijer阅读论文草稿。感谢OpenAI在学术访问计划下的支持。 此项目如果没有以下单位的支持是不可能完成的:
- UKPLab: SBERT, BEIR, USEB
- Eleuther AI 模型
- Huggingface Transformers
引用
如果SGPT对你有帮助,请随意引用我们的论文 :)
@article{muennighoff2022sgpt,
title={SGPT: GPT Sentence Embeddings for Semantic Search},
author={Muennighoff, Niklas},
journal={arXiv preprint arXiv:2202.08904},
year={2022}
}