sgpt

GitHub
873 51 中等 1 次阅读 2个月前MIT开发框架语言模型其他Agent
AI 解读 由 AI 自动生成,仅供参考

sgpt 是一个基于 GPT 大模型的开源项目,旨在将强大的语言生成能力转化为高效的语义搜索工具。它主要解决了传统小模型在理解复杂语境和长文本语义时精度不足的问题,让开发者能够利用 GPT 的强大表征能力进行高精度的文档检索、问答匹配及语义相似度计算。

该项目特别适合 AI 研究人员、后端开发者以及需要构建企业级知识库或智能搜索系统的技术团队使用。sgpt 的核心亮点在于其独特的技术路径:它提出了两种模式——SGPT-BE(双编码器)和 SGPT-CE(交叉编码器)。其中,SGPT-BE 通过仅微调偏置张量并结合位置加权平均池化,就能生成高质量的句子向量,极大降低了训练成本;而 SGPT-CE 则无需任何微调,直接利用 GPT 原有的对数概率机制即可实现卓越的排序效果。此外,sgpt 不仅支持英语,还发布了基于 BLOOM 的多语言版本,并能无缝集成到流行的 Sentence Transformers 库中,方便用户快速部署对称或非对称的搜索任务。虽然作者近期推荐了性能更统一的新一代模型 GRIT,但 sgpt 作为验证大模型嵌入能力的经典方案,依然具有重要的参考和应用价值。

使用场景

某大型电商公司的搜索团队正致力于优化其内部海量商品评论库的检索系统,以支持客服快速定位用户反馈。

没有 sgpt 时

  • 关键词匹配局限大:传统搜索引擎仅依赖关键词重合度,当用户搜索“电池不耐用”时,无法召回包含“续航太短”或“充电频繁”等语义相同但措辞不同的评论。
  • 长尾查询效果差:面对复杂的自然语言提问(如“适合送给老人的操作简单的手机”),系统难以理解意图,往往返回大量无关的高频商品列表。
  • 维护成本高昂:为了覆盖同义词和场景变体,工程师需要手动构建庞大的同义词库和规则引擎,耗时耗力且难以跟上用户表达习惯的变化。
  • 跨语言支持困难:面对多语种用户评论,需要为每种语言单独训练或配置模型,架构臃肿且推理延迟高。

使用 sgpt 后

  • 语义理解精准:sgpt 利用 GPT 生成的句子嵌入向量,能直接捕捉“电池不耐用”与“续航太短”之间的深层语义关联,显著提升召回率。
  • 自然语言交互流畅:借助 sgpt 的不对称搜索能力,系统能完美解析复杂的长句查询,直接定位到最符合意图的具体评论段落。
  • 开发效率飞跃:无需手动维护同义词库,仅需微调偏置张量即可让模型适应特定业务领域,大幅降低了迭代周期。
  • 多语言统一架构:利用 sgpt 的多语言 BLOOM 模型,一套架构即可同时处理中、英、法等多语种检索,简化了系统部署并降低了延迟。

sgpt 通过将大语言模型的生成能力转化为高效的语义搜索引擎,让机器真正“读懂”了用户的自然语言意图。

运行环境要求

操作系统
  • 未说明
GPU
  • 非绝对必需(代码支持 CPU 运行),但推荐使用 NVIDIA GPU
  • 显存需求取决于模型大小:125M 模型需少量显存,5.8B 模型建议 16GB+,7.1B 模型建议 24GB+
  • CUDA 版本未说明
内存

最低 8GB(小模型),推荐 32GB+(运行 5.8B 或 7.1B 大模型时)

依赖
notes该工具核心依赖 Hugging Face Transformers 和 PyTorch。README 中未提供具体的 requirements.txt 内容,但示例代码明确使用了 torch、transformers 和 scipy。若使用 Sentence Transformers 集成方案,需安装 sentence-transformers 库。模型文件较大(从 125M 到 7.1B 参数不等),首次运行会自动从 Hugging Face 下载,需确保网络连接畅通及磁盘空间充足。对于 5.8B 及以上的大模型,强烈建议使用 GPU 以获得可用速度。
python未说明
torch
transformers
scipy
sentence-transformers
sgpt hero image

快速开始

SGPT:用于语义搜索的GPT句子嵌入

本仓库包含论文《SGPT:用于语义搜索的GPT句子嵌入》(https://arxiv.org/abs/2202.08904)的相关代码、结果及预训练模型。

**************************** 更新 *****************************

  • 2024-02:我们发布了GRIT & GritLM——这些模型将SGPT双编码器、交叉编码器、对称式、非对称式以及常规GPT(即生成式)统一到一个单一模型中,并且在各项指标上都表现得更好。我们建议切换到这些新模型 :)
  • 2022-09:SGPT双编码器现在可以轻松地与Sentence Transformers一起使用,详见新脚本
  • 2022-08:多语言BLOOM SGPT模型发布:非对称式,71亿参数对称式,17亿参数。如果您需要其他模型,请随时提交issue。
  • 2022-06:OpenAI发布了其Search Endpoint的工作机制,我们在论文中将其与SGPT交叉编码器进行了对比。我们的方法非常相似。您可以尝试使用crossencoder/beir/openai_search_endpoint_functionality.py中所示的提示!
  • 2022-03:58亿参数的双编码器模型在USEB和BEIR上的表现分别提升了4%和1%。论文Hugging Face上的模型均已更新。这是通过使用GradCache增大批量大小实现的,更多信息请参阅论文。如果您之前下载过这些模型,我们建议用新版本替换。
  • 2022-02:我们发布了我们的论文。快来看看吧! :)

快速链接

概述

我们提出了SGPT-BE和SGPT-CE,用于将GPT模型作为双编码器或交叉编码器应用于对称或非对称搜索。SGPT-BE通过仅对偏置张量进行对比微调以及位置加权平均池化,生成具有语义意义的句子嵌入。SGPT-CE则直接使用GPT模型的日志概率,无需任何微调。以下是这些方法的示意图。

如果您有任何问题,欢迎随时提交issue~

结构

.
├── 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 # Sentence Transformers的改编版本——所有双编码器实验请安装此版本
│   │   │   └── ...
│   │   └── README.md
│   └── useb
│       ├── useb
│       │   └── ...
│       ├── *.bash # 用于运行多个实验的Bash脚本
│       ├── useb_dense_retriever.py
│       └── README.md
├── crossencoder  # 交叉编码器的推理
│   └── beir
│       ├── *.ipynb # 在README中解释的Notebook文件
│       └── README.md
├── other
│   ├── sgpt_graphic.png
│   └── sgpt_utils.ipynb # 用于生成论文中图表及其他内容的代码
├── requirements.txt
└── README.md

每个数据子目录都提供自己的README,其中包含其结构下载(数据集、模型)以及用于生成数据集、模型等的命令的概述。通常,您可以在https://huggingface.co/Muennighoff找到所有模型,并在https://www.kaggle.com/muennighoff/datasets中找到各种数据集的JSON结果。模型名称在其Hugging Face的README中有所说明。数据集名称则在本仓库的子文件夹中有所说明。

使用Hugging Face中的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):
    # 不进行填充的分词
    batch_tokens = tokenizer(texts, padding=False, truncation=True)   
    # 添加特殊括号并注意它们
    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)
    # 添加填充
    batch_tokens = tokenizer.pad(batch_tokens, padding=True,return_tensors="pt")
    return batch_tokens

def get_weightedmean_embedding(batch_tokens, model):
    # 获取嵌入
    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

    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)

# 计算余弦相似度
# 余弦相似度的取值范围为 [-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))

交叉编码器

非对称语义搜索 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)
        # 去掉最后一个标记,因为我们取的是其前一个标记的概率
        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:]
        # 收集延续部分标记的对数概率 -> [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 部分相同的代码,只需更改提示词即可。欢迎分享效果良好的提示词 :)

将 SGPT 与 Sentence Transformers 结合使用

双编码器 ST

对称语义搜索 BE ST

对称模型现在已通过 pip install git+https://github.com/UKPLab/sentence-transformers.git 与最新的 sentence-transformers 完全兼容。你应该会得到与 上述 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^14公里),位于宝瓶座内。",
    "塔图因是一颗位于银河系外环、环绕双子恒星运转的严酷沙漠星球,这里法律不存,由赫特族黑帮统治。许多定居者靠水分农场勉强维持生计,而莫斯艾斯利和莫斯埃斯帕等太空港城市则是走私犯、罪犯和其他亡命之徒的大本营。",
]

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特殊tokens=False)[0]
        self._first模块().bos_spec_token_d_rep = self._first模块().tokenizer.encode("{", add特殊tokens=False)[0]
        self._first模块().eos_spec_token_d = self._first模块().tokenizer.encode("}", add特殊tokens=False)[0]
        self._first模块().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)。 需要注意的是,这种方法产生的分数会略低于 SGPT Sentence Transformers,因为特殊括号在分词时可能会与其他标记混淆。在 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^14公里),位于宝瓶座内。",
    "塔图因是一颗位于银河系外环、环绕双子恒星的严酷沙漠星球,这里法律失效,由赫特族黑帮统治。许多定居者靠水分农场勉强维持生计,而莫斯艾斯利和莫斯埃斯帕等太空港城市则成为走私犯、罪犯和其他亡命之徒的大本营。",
]

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 学术访问计划的支持。 如果没有以下机构和资源,这项工作将无法完成:

引用

如果您觉得 SGPT 对您有所帮助,请随时引用我们的论文 :)

@article{muennighoff2022sgpt,
  title={SGPT:用于语义搜索的 GPT 句子嵌入},
  author={Muennighoff, Niklas},
  journal={arXiv 预印本 arXiv:2202.08904},
  year={2022}
}

常见问题

相似工具推荐

stable-diffusion-webui

stable-diffusion-webui 是一个基于 Gradio 构建的网页版操作界面,旨在让用户能够轻松地在本地运行和使用强大的 Stable Diffusion 图像生成模型。它解决了原始模型依赖命令行、操作门槛高且功能分散的痛点,将复杂的 AI 绘图流程整合进一个直观易用的图形化平台。 无论是希望快速上手的普通创作者、需要精细控制画面细节的设计师,还是想要深入探索模型潜力的开发者与研究人员,都能从中获益。其核心亮点在于极高的功能丰富度:不仅支持文生图、图生图、局部重绘(Inpainting)和外绘(Outpainting)等基础模式,还独创了注意力机制调整、提示词矩阵、负向提示词以及“高清修复”等高级功能。此外,它内置了 GFPGAN 和 CodeFormer 等人脸修复工具,支持多种神经网络放大算法,并允许用户通过插件系统无限扩展能力。即使是显存有限的设备,stable-diffusion-webui 也提供了相应的优化选项,让高质量的 AI 艺术创作变得触手可及。

162.1k|★★★☆☆|今天
开发框架图像Agent

everything-claude-code

everything-claude-code 是一套专为 AI 编程助手(如 Claude Code、Codex、Cursor 等)打造的高性能优化系统。它不仅仅是一组配置文件,而是一个经过长期实战打磨的完整框架,旨在解决 AI 代理在实际开发中面临的效率低下、记忆丢失、安全隐患及缺乏持续学习能力等核心痛点。 通过引入技能模块化、直觉增强、记忆持久化机制以及内置的安全扫描功能,everything-claude-code 能显著提升 AI 在复杂任务中的表现,帮助开发者构建更稳定、更智能的生产级 AI 代理。其独特的“研究优先”开发理念和针对 Token 消耗的优化策略,使得模型响应更快、成本更低,同时有效防御潜在的攻击向量。 这套工具特别适合软件开发者、AI 研究人员以及希望深度定制 AI 工作流的技术团队使用。无论您是在构建大型代码库,还是需要 AI 协助进行安全审计与自动化测试,everything-claude-code 都能提供强大的底层支持。作为一个曾荣获 Anthropic 黑客大奖的开源项目,它融合了多语言支持与丰富的实战钩子(hooks),让 AI 真正成长为懂上

139k|★★☆☆☆|今天
开发框架Agent语言模型

ComfyUI

ComfyUI 是一款功能强大且高度模块化的视觉 AI 引擎,专为设计和执行复杂的 Stable Diffusion 图像生成流程而打造。它摒弃了传统的代码编写模式,采用直观的节点式流程图界面,让用户通过连接不同的功能模块即可构建个性化的生成管线。 这一设计巧妙解决了高级 AI 绘图工作流配置复杂、灵活性不足的痛点。用户无需具备编程背景,也能自由组合模型、调整参数并实时预览效果,轻松实现从基础文生图到多步骤高清修复等各类复杂任务。ComfyUI 拥有极佳的兼容性,不仅支持 Windows、macOS 和 Linux 全平台,还广泛适配 NVIDIA、AMD、Intel 及苹果 Silicon 等多种硬件架构,并率先支持 SDXL、Flux、SD3 等前沿模型。 无论是希望深入探索算法潜力的研究人员和开发者,还是追求极致创作自由度的设计师与资深 AI 绘画爱好者,ComfyUI 都能提供强大的支持。其独特的模块化架构允许社区不断扩展新功能,使其成为当前最灵活、生态最丰富的开源扩散模型工具之一,帮助用户将创意高效转化为现实。

107.7k|★★☆☆☆|2天前
开发框架图像Agent

NextChat

NextChat 是一款轻量且极速的 AI 助手,旨在为用户提供流畅、跨平台的大模型交互体验。它完美解决了用户在多设备间切换时难以保持对话连续性,以及面对众多 AI 模型不知如何统一管理的痛点。无论是日常办公、学习辅助还是创意激发,NextChat 都能让用户随时随地通过网页、iOS、Android、Windows、MacOS 或 Linux 端无缝接入智能服务。 这款工具非常适合普通用户、学生、职场人士以及需要私有化部署的企业团队使用。对于开发者而言,它也提供了便捷的自托管方案,支持一键部署到 Vercel 或 Zeabur 等平台。 NextChat 的核心亮点在于其广泛的模型兼容性,原生支持 Claude、DeepSeek、GPT-4 及 Gemini Pro 等主流大模型,让用户在一个界面即可自由切换不同 AI 能力。此外,它还率先支持 MCP(Model Context Protocol)协议,增强了上下文处理能力。针对企业用户,NextChat 提供专业版解决方案,具备品牌定制、细粒度权限控制、内部知识库整合及安全审计等功能,满足公司对数据隐私和个性化管理的高标准要求。

87.6k|★★☆☆☆|今天
开发框架语言模型

ML-For-Beginners

ML-For-Beginners 是由微软推出的一套系统化机器学习入门课程,旨在帮助零基础用户轻松掌握经典机器学习知识。这套课程将学习路径规划为 12 周,包含 26 节精炼课程和 52 道配套测验,内容涵盖从基础概念到实际应用的完整流程,有效解决了初学者面对庞大知识体系时无从下手、缺乏结构化指导的痛点。 无论是希望转型的开发者、需要补充算法背景的研究人员,还是对人工智能充满好奇的普通爱好者,都能从中受益。课程不仅提供了清晰的理论讲解,还强调动手实践,让用户在循序渐进中建立扎实的技能基础。其独特的亮点在于强大的多语言支持,通过自动化机制提供了包括简体中文在内的 50 多种语言版本,极大地降低了全球不同背景用户的学习门槛。此外,项目采用开源协作模式,社区活跃且内容持续更新,确保学习者能获取前沿且准确的技术资讯。如果你正寻找一条清晰、友好且专业的机器学习入门之路,ML-For-Beginners 将是理想的起点。

85k|★★☆☆☆|今天
图像数据工具视频

ragflow

RAGFlow 是一款领先的开源检索增强生成(RAG)引擎,旨在为大语言模型构建更精准、可靠的上下文层。它巧妙地将前沿的 RAG 技术与智能体(Agent)能力相结合,不仅支持从各类文档中高效提取知识,还能让模型基于这些知识进行逻辑推理和任务执行。 在大模型应用中,幻觉问题和知识滞后是常见痛点。RAGFlow 通过深度解析复杂文档结构(如表格、图表及混合排版),显著提升了信息检索的准确度,从而有效减少模型“胡编乱造”的现象,确保回答既有据可依又具备时效性。其内置的智能体机制更进一步,使系统不仅能回答问题,还能自主规划步骤解决复杂问题。 这款工具特别适合开发者、企业技术团队以及 AI 研究人员使用。无论是希望快速搭建私有知识库问答系统,还是致力于探索大模型在垂直领域落地的创新者,都能从中受益。RAGFlow 提供了可视化的工作流编排界面和灵活的 API 接口,既降低了非算法背景用户的上手门槛,也满足了专业开发者对系统深度定制的需求。作为基于 Apache 2.0 协议开源的项目,它正成为连接通用大模型与行业专有知识之间的重要桥梁。

77.1k|★★★☆☆|2天前
Agent图像开发框架