一个嵌入器，适用于任何任务：指令微调文本嵌入

本仓库包含我们论文《一个嵌入器，适用于任何任务：指令微调文本嵌入》（arXiv:2212.09741）的代码及预训练模型。请访问我们的项目主页以获取项目的快速概览。

我们推出了Instructor👨‍🏫，这是一种经过指令微调的文本嵌入模型，能够根据任务指令生成针对任意任务（如分类、检索、聚类、文本评估等）和领域的文本嵌入（如科学、金融等领域），只需提供任务指令即可，无需额外微调。Instructor👨‍在70个多样化的嵌入任务上均取得了当前最优性能！

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

• 01/21：我们更新了代码结构，支持便捷的包安装。
• 12/28：我们使用硬负样本更新了检查点。
• 12/20：我们发布了论文、代码、项目主页以及检查点。快来查看吧！

快速链接

• 一个嵌入器，适用于任何任务：指令微调文本嵌入
  • 快速链接
  • 安装
    • 环境设置
  • 入门指南
    • encode 函数
  • 模型列表
  • 使用场景
    • 为自定义文本计算嵌入
    • 计算文本之间的相似度
    • 将自定义嵌入用于信息检索
    • 将自定义嵌入用于聚类
  • 训练
    • 数据
    • 训练 INSTRUCTOR
  • 评估
    • MTEB
    • Billboard
    • 提示检索
  • 量化
  • 遇到问题或有疑问？
  • 引用
  • INSTRUCTOR 其他资源

安装

使用 INSTRUCTOR 进行任何文本嵌入都非常简单。您可以在Colab 笔记本中轻松试用。在本地机器上，我们建议首先创建一个虚拟环境：

conda env create -n instructor python=3.7
git clone https://github.com/HKUNLP/instructor-embedding
pip install -r requirements.txt

这将创建我们使用的 instructor 环境。要使用嵌入工具，首先从 PyPI 安装 InstructorEmbedding 包：

pip install InstructorEmbedding

或者直接从我们的代码中安装：

pip install -e .

环境设置

通过运行以下命令激活环境：

conda activate instructor

入门指南

首先下载一个预训练模型（完整可用模型列表见模型列表）：

from InstructorEmbedding import INSTRUCTOR
model = INSTRUCTOR('hkunlp/instructor-large')

然后将句子和自定义指令提供给模型：

# 准备带有指令的文本对
text_instruction_pairs = [
    {"instruction": "表示科学标题:", "text": "3D ActionSLAM：多层环境中可穿戴式人物跟踪"},
    {"instruction": "表示医学句子以便检索重复句子:", "text": "近期研究表明，他汀类药物——一种在心血管疾病死亡率预防中广泛应用的药物——可能延缓或预防乳腺癌复发，但其对疾病特异性死亡率的影响尚不明确。"}
]

# 后处理
texts_with_instructions = []
for pair in text_instruction_pairs:
    texts_with_instructions.append([pair["instruction"], pair["text"]])

# 计算嵌入
customized_embeddings = model.encode(texts_with_instructions)

就这样完成了。我们现在得到了一组包含嵌入的 NumPy 数组。

for pair, embedding in zip(text_instruction_pairs, customized_embeddings):
    print("指令: ", pair["instruction"])
    print("文本: ", pair["text"])
    print("嵌入: ", embedding)
    print("")

encode 函数

用户只需使用 encode 函数即可：

model.encode( sentences,
              batch_size: int = 32,
              show_progress_bar: bool = None,
              output_value: str = 'sentence_embedding',
              convert_to_numpy: bool = True,
              convert_to_tensor: bool = False,
              device: str = None,
              normalize_embeddings: bool = False)

• sentences：待嵌入的句子。格式应为 [["指令提示 0", "待嵌入文本 0], ["指令提示 1", "待嵌入文本 1], ...]。
• batch_size（默认：32）：用于计算的批次大小，决定每批同时处理的句子数量。
• show_progress_bar（默认：None）：若设置为 True，编码过程中会显示进度条，直观地展示编码进度。
• output_value（默认：'sentence_embedding'）：指定所需的输出类型。默认值 'sentence_embedding' 返回句子嵌入；设置为 'token_embeddings' 则返回词元嵌入；设置为 None 则返回所有输出值。
• convert_to_numpy（默认：True）：若设置为 True，输出为 NumPy 向量列表；若设置为 False，则输出为 PyTorch 张量列表。
• convert_to_tensor（默认：False）：若设置为 True，函数将以堆叠张量的形式作为单一输出返回。此参数会覆盖 convert_to_numpy 的设置。
• device（默认：None）：指定用于计算的 torch.device。若未指定，则使用默认设备。
• normalize_embeddings（默认：False）：若设置为 True，返回的向量长度将归一化为 1，此时相似度搜索将使用更快的点积方法（util.dot_score），而非余弦相似度。

模型列表

我们发布了一系列不同规模的 INSTRUCTOR 检查点。您可以使用 InstructorEmbedding 包轻松加载这些模型。

模型	平均得分
hkunlp/instructor-base	55.9
hkunlp/instructor-large	58.4
hkunlp/instructor-xl	58.8

使用场景

以下我们提供了一些具体的使用场景。更多示例和应用，请参阅我们的论文。

为自定义文本计算嵌入

如果您希望为特定句子计算自定义嵌入，可以按照统一模板编写指令：

                          表示用于 任务目标 的 领域 文本类型：

• 领域 是可选的，用于指定文本所属的领域，例如科学、金融、医学等。
• 文本类型 是必需的，用于指定编码单位，例如句子、文档、段落等。
• 任务目标 是可选的，用于指定嵌入的目的，例如检索文档、分类句子等。

计算文本之间的相似度

您可以使用 INSTRUCTOR 来计算两组句子之间的相似度，并生成自定义嵌入。

from sklearn.metrics.pairwise import cosine_similarity
sentences_a = [['表示科学领域的句子：','QCD物质中的部分子能量损失'], 
               ['表示金融领域的句子：','美联储周三上调了基准利率。']]
sentences_b = [['表示科学领域的句子：','耗散动力学中的手征相变'],
               ['表示金融领域的句子：','周五该基金上涨不足0.5%']]
embeddings_a = model.encode(sentences_a)
embeddings_b = model.encode(sentences_b)
similarities = cosine_similarity(embeddings_a,embeddings_b)

将自定义嵌入用于信息检索

import numpy as np
from sklearn.metrics.pairwise import cosine_similarity
query  = [['表示维基百科问题，用于检索支持性文档：','山药植物中食物储存在哪里']],
corpus = [['表示维基百科文档，用于检索：','资本主义自封建制度结束以来一直在西方世界占据主导地位，但大多数人认为[谁？]“混合经济”这一术语更准确地描述了大多数当代经济，因为它们同时包含私营企业和国有企业。在资本主义制度中，价格决定了供需关系。例如，对某些商品和服务的需求增加会导致价格上涨，而需求减少则会导致价格下降。'],
          ['表示维基百科文档，用于检索：',"公平住房法案下的差异影响理论尤其具有争议性，因为该法案规范了许多与住房、保险和抵押贷款相关的活动——一些学者认为，该理论在公平住房法案下的应用，加上社区再投资法案的扩展，导致了次级抵押贷款的兴起以及美国房地产市场的崩溃和随之而来的全球经济衰退"],
          ['表示维基百科文档，用于检索：','在美国劳动法中，差异影响是指在就业、住房等领域存在的做法，这些做法会以某种受保护特征为依据，对某一特定群体造成比其他群体更大的不利影响，即使雇主或房东所采用的规则表面上是中立的。尽管受保护的类别因法律而异，但大多数联邦民权法律都以种族、肤色、宗教、国籍和性别作为受保护的特征，有些法律还涵盖了残疾状况和其他特征。']]
query_embeddings = model.encode(query)
corpus_embeddings = model.encode(corpus)
similarities = cosine_similarity(query_embeddings,corpus_embeddings)
retrieved_doc_id = np.argmax(similarities)
print(retrieved_doc_id)

将自定义嵌入用于聚类

import sklearn.cluster
sentences = [['表示医学领域的句子，用于聚类：','霍瓦瓦-利夫希茨引力中的动力学标量自由度'],
             ['表示医学领域的句子，用于聚类：','低能区大气中微子通量计算的比较'],
             ['表示医学领域的句子，用于聚类：','大质量格罗斯-内韦模型中的费米子袋'],
             ['表示医学领域的句子，用于聚类：',"泰瓦特隆加速器上伴随t-tbar-H产生的QCD修正"],
             ['表示医学领域的句子，用于聚类：','R测量的新分析：粲夸克偶素更高能态矢量粒子的共振参数']]
embeddings = model.encode(sentences)
clustering_model = sklearn.cluster.MiniBatchKMeans(n_clusters=2)
clustering_model.fit(embeddings)
cluster_assignment = clustering_model.labels_
print(cluster_assignment)

训练

数据

我们构建了包含指令的多任务嵌入数据集（MEDI），该数据集由来自 Super-NI（Super-NaturalInstructions）、sentence-transformer 嵌入训练数据、KILT 和 MedMCQA 的 330 个数据集组成，涵盖了广泛的领域和任务。如果未提供正样本和负样本对，我们会自行构造，并以统一格式存储：

[
    {'query': ['表示用于检索相关文档的维基百科问题;', '大小谎言第二季有多少集'], 'pos': ['表示用于检索的维基百科文档;', '《大小谎言》（电视剧）获得了多项荣誉。它获得了16项艾美奖提名并赢得了8项，包括最佳限定剧集以及基德曼、斯卡斯加德和德恩的表演奖项。这三位演员还获得了金球奖，此外该剧还荣获了最佳迷你剧或电视电影金球奖。基德曼和斯卡斯加德也因其表演获得了美国影视演员协会奖。尽管最初被宣传为迷你剧，HBO还是续订了该剧的第二季。第二季的制作于2018年3月开始，预计将于2019年首播。全部七集均由凯利编剧'], 'neg': ['表示用于检索的维基百科文档;', '《小人物，大世界》第二季B部分的最后一集——“农场超载”。一群人围在躺在投石机旁地上的雅各布身边。第二季B的前两集聚焦于这起事故以及当地媒体对此事的反应。第一季《小人物，大世界》为TLC频道带来了稳定的收视率（尤其是在18至49岁这一重要人群中），从而促使该节目续订第二季。对该节目的评论总体上较为正面，称赞其对身材矮小人群的积极刻画。另一方面，也有评论认为该节目带有窥探性质'], 'task_id': 1}
    {'query': ['表示用于检索相关文档的维基百科问题;', '谁演唱了《等待像你这样的女孩》'], 'pos': ['表示用于检索的维基百科文档;', '《等待像你这样的女孩》\n《等待像你这样的女孩》是英国-美国摇滚乐队Foreigner于1981年发行的一首力量抒情歌曲。这首歌曲中标志性的合成器旋律由当时尚不知名的托马斯·多尔比演奏，这首歌也标志着乐队风格的重大转变：此前他们的单曲多为中快节奏的摇滚乐，而此曲则是一首更为柔和的情歌，同时兼具力量抒情曲的能量。它是专辑《4》（1981）中的第二支单曲，由卢·格拉姆和米克·琼斯共同创作。它已成为乐队最'], 'neg': ['表示用于检索的维基百科文档;', '《等待像你这样的女孩》曾连续九周力压奥莉维亚·牛顿-约翰的单曲《Physical》占据榜首，随后又在1982年1月30日被霍尔与欧茨的《I Can\'t Go for That (No Can Do)》挤下，保持了第十周的冠军位置。由于其在榜单上的持久表现，最终成为1982年百强单曲榜第19位的歌曲。这首歌是乐队当时最大的热门单曲，直到1985年《I Want to Know What Love Is》登顶为止。该歌曲在《公告牌》“史上百大歌曲”中排名第100位。《等待像你这样的女孩》\n《等待像一个女孩']], 'task_id': 1}
    ...
    {'query': ['表示用于检索相关文档的维基百科句子;', '我超爱甜味马提尼酒！'], 'pos': ['表示用于检索的维基百科文档;', '苹果马提尼\n苹果马提尼（简称Appletini）是一种含有伏特加以及一种或多种苹果汁、苹果酒、苹果利口酒或苹果白兰地的鸡尾酒。\n这种饮品最初被称为“亚当的苹果马提尼”，因为创造它的调酒师名叫亚当。它诞生于1996年位于西好莱坞的Lola's餐厅。\n这款名为“亚当的苹果”的饮品曾在1972年7月号的《花花公子》杂志内封面上由Smirnoff品牌进行广告宣传。配方要求使用大约一盎司的Smirnoff'], 'neg': ['表示用于检索的维基百科文档;', '本法案所描述的调味葡萄酒类饮料包括“调味葡萄酒基饮料”（非加强型）和“调味葡萄酒产品鸡尾酒”（混合型、酒精含量低于7%的低度饮品）。\n调味葡萄酒的种类。\n调味葡萄酒的种类——苦艾酒。\n苦艾酒是最常用的调味葡萄酒，因为它广泛用于调制鸡尾酒，并且拥有Martini和Cinzano等在全球范围内广为人知的知名品牌。苦艾酒有甜味和干味之分，颜色可以是红色、白色、粉色或橙色。传统上'], 'task_id': 300}
]

每个实例包含一个问题、一个正样本对、一个负样本对以及任务ID，任务ID用于确保同一训练批次中的数据来自同一任务。 MEDI 数据可通过此链接下载。

训练 INSTRUCTOR

我们提供了 INSTRUCTOR 的训练示例脚本。您可能需要先下载 MEDI 数据，解压缩文件夹并将 medi-data.json 放置在 --cache_dir 目录下。

python train.py --model_name_or_path sentence-transformers/gtr-t5-large --output_dir {输出目录 } --cache_dir { 缓存目录 } --max_source_length 512 --num_train_epochs 10 --save_steps 500 --cl_temperature 0.1 --warmup_ratio 0.1 --learning_rate 2e-5 --overwrite_output_dir

我们对各个参数说明如下：

• --model_name_or_path：用于开始训练的预训练检查点。我们支持模型标识符（例如 sentence-transformers/gtr-t5-large、sentence-transformers/sentence-t5-large）或检查点路径（例如由 transformers 训练器保存的检查点）。
• --cl_temperature：对比损失的温度系数。
• --cache_dir：用于缓存已下载模型和数据的目录。下载的 MEDI 数据（medi-data.json）应放置在 --cache_dir 目录下。
• --output_dir：用于存储训练好的模型（检查点）以便评估的目录。

其余参数均为标准的 Huggingface's transformers 训练参数，例如 --overwrite_output_dir、--num_train_epochs、--learning_rate。更多详细信息请参阅 Huggingface transformers。

评估

我们在70个多样化的任务上大规模评估INSTRUCTOR，这些任务涵盖了广泛的任务和领域。具体来说，我们的评估基于三个基准：MTEB、Billboard和Prompt Retrieval。以下我们将详细说明如何运行评估脚本。

MTEB

为了在MTEB基准数据集上评估模型性能，首先安装MTEB库：

cd evaluation/MTEB
pip install -e .

然后运行以下命令：

python examples/evaluate_model.py --model_name hkunlp/instructor-large --output_dir outputs --task_name ArguAna --result_file results

您可以通过指定--model_name来评估自己训练的模型检查点，并通过更改--task_name来运行所有的MTEB数据集。有关所有任务的评估指标，请参阅我们的论文或MTEB基准。

Billboard

要评估模型在Billboard上的性能，运行以下命令：

cd evaluation/text_evaluation
python main.py --model_name hkunlp/instructor-large --task mscoco --add_prompt

您可以通过指定--model_name来评估自己训练的模型检查点，并通过更改--task来运行所有的Billboard数据集。在Billboard的三个数据集中，我们都报告了皮尔逊相关系数。

Prompt Retrieval

要评估模型在Prompt Retrieval上的性能，运行以下命令：

cd evaluation/prompt_retrieval
python main.py --embedding_model hkunlp/instructor-large --task rte --model_cache_dir {cache_dir} --output_dir {output_dir} --add_prompt

您可以通过指定--model_name来评估自己训练的模型检查点，并通过更改--task来运行Prompt Retrieval数据集。为了保持一致的评价标准，我们将Prompt Retrieval中的所有任务统一转换为“文本到文本”的格式，并报告Rouge-L分数。

量化

要对Instructor嵌入模型进行量化，请运行以下代码：

# 导入模块
import torch
from InstructorEmbedding import INSTRUCTOR

# 加载模型
model = INSTRUCTOR('hkunlp/instructor-large', device='cpu')  # 您也可以使用GPU

# 对模型进行量化
qmodel = torch.quantization.quantize_dynamic(
    model, {torch.nn.Linear}, dtype=torch.qint8)

# 推理
sentence = "3D ActionSLAM: 可穿戴式多层环境人体跟踪"
instruction = "表示科学标题："

embeddings = qmodel.encode([[instruction,sentence]])  
# 您也可以对嵌入进行归一化：normalize_embeddings=True 

print(f"量化后的嵌入：\n {embeddings}")

这会将模型大小缩小10倍，并且推理时间也会比普通模型更短 :)

遇到问题或有疑问？

如果您对代码或论文有任何疑问，请随时发送邮件至Hongjin（hjsu@cs.hku.hk）和Weijia（swj0419@cs.washington.edu）。请尽量详细描述问题，以便我们能够更快更好地帮助您。

引用

如果您觉得我们的工作有所帮助，请引用我们：

@inproceedings{INSTRUCTOR,
  title={一个嵌入器，适用于任何任务：指令微调的文本嵌入},
  author={Su, Hongjin and Shi, Weijia and Kasai, Jungo and Wang, Yizhong and Hu, Yushi and  Ostendorf, Mari and Yih, Wen-tau and Smith, Noah A. and  Zettlemoyer, Luke and Yu, Tao},
  url={https://arxiv.org/abs/2212.09741},
  year={2022},
}

INSTRUCTOR的其他应用

我们感谢社区为扩展INSTRUCTOR所做的努力！

• LangChain 支持使用INSTRUCTOR模型的InstructEmbeddings。
• MosaicML 已将Instructor-Large和Instructor-XL纳入其支持列表。
• embaas 集成了Instructor-Large。
• Haystack 包含了InstructorTextEmbedder和InstructorDocumentEmbedder组件。

Instructor-Embedding 快速上手指南

Instructor 是一个指令微调的文本嵌入模型。它无需针对特定任务进行微调，只需在输入时提供简单的任务指令（Instruction），即可生成适用于分类、检索、聚类等任意任务和领域的文本向量。

环境准备

• 操作系统: Linux, macOS, Windows
• Python 版本: 推荐 Python 3.7 及以上
• 依赖管理: 推荐使用 conda 创建虚拟环境以避免依赖冲突
• 硬件: 支持 CPU 运行，建议使用 NVIDIA GPU 以加速推理

安装步骤

1. 创建并激活虚拟环境

conda create -n instructor python=3.7
conda activate instructor

2. 安装依赖包

你可以直接从 PyPI 安装官方发布的包（推荐）：

pip install InstructorEmbedding

或者克隆源码进行开发模式安装：

git clone https://github.com/HKUNLP/instructor-embedding
cd instructor-embedding
pip install -e .

国内加速提示：如果下载速度较慢，建议使用清华或阿里镜像源： pip install InstructorEmbedding -i https://pypi.tuna.tsinghua.edu.cn/simple

基本使用

1. 加载模型

首先导入库并加载预训练模型（如 instructor-large）。首次运行时会自动从 Hugging Face 下载模型。

from InstructorEmbedding import INSTRUCTOR

# 加载模型，可选模型包括：hkunlp/instructor-base, hkunlp/instructor-large, hkunlp/instructor-xl
model = INSTRUCTOR('hkunlp/instructor-large')

2. 准备数据与指令

Instructor 的核心在于指令。输入格式必须为列表嵌套列表：[[指令，文本], [指令，文本], ...]。

通用指令模板：

"Represent the 领域 文本类型 for 任务目标:"

• 领域 (可选): 如 science, finance, medicine
• 文本类型 (必填): 如 sentence, document, paragraph
• 任务目标 (可选): 如 retrieving a duplicate sentence, clustering, classification

# 定义包含指令和文本的数据对
text_instruction_pairs = [
    ["Represent the Science title:", "3D ActionSLAM: wearable person tracking in multi-floor environments"],
    ["Represent the Medicine sentence for retrieving a duplicate sentence:", "Recent studies have suggested that statins could delay breast cancer recurrence."]
]

3. 生成嵌入向量

调用 encode 方法即可获取向量表示。

# 计算嵌入向量
customized_embeddings = model.encode(text_instruction_pairs)

# 查看结果
for pair, embedding in zip(text_instruction_pairs, customized_embeddings):
    print(f"指令：{pair[0]}")
    print(f"文本：{pair[1]}")
    print(f"向量维度：{embedding.shape}")
    print("-" * 30)

进阶参数说明

model.encode() 支持以下常用参数：

• batch_size: 批处理大小，默认 32。
• show_progress_bar: 是否显示进度条，默认 None。
• convert_to_numpy: 是否转换为 numpy 数组，默认 True。
• normalize_embeddings: 是否归一化向量（长度为 1），默认 False。若设为 True，后续相似度计算可使用更快的点积。

# 示例：开启进度条并归一化输出
embeddings = model.encode(
    text_instruction_pairs, 
    show_progress_bar=True, 
    normalize_embeddings=True
)