一个嵌入器,适用于任何任务:指令微调的文本嵌入
这个仓库包含了我们的论文《一个嵌入器,适用于任何任务:指令微调的文本嵌入模型》(地址)的代码和预训练模型。请参阅我们的项目页面以获得项目概述。
我们介绍了 Instructor👨🏫,一个指令微调的文本嵌入模型,它可以根据要求生成适用于任何任务(例如分类、检索、聚类、文本评估等)和领域(例如科学、金融等)的文本嵌入仅通过提供任务指令,无需任何微调。Instructor👨在70个不同的嵌入任务中达到了最先进的水平!
**************************** 更新信息 ****************************
- 01/21:我们更新了代码结构,支持方便的包安装。
- 12/28:我们更新了带有硬否定示例的模型检查点。
- 12/20:我们发布了论文、代码、项目页面和模型检查点。快来看看吧!
快速链接
安装
使用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 = [['表征医学句子以进行聚类:', 'Horava-Lifshitz引力中的动态标量自由度'],
['表征医学句子以进行聚类:', '低能量的大气中微子通量计算比较'],
['表征医学句子以进行聚类:', '质量Gross-Neveu模型中的费米子袋'],
['表征医学句子以进行聚类:', '在费米实验室的关联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)、句子变换器嵌入训练数据(地址)、KILT(地址)和MedMCQA(地址)的330个数据集组成,覆盖了广泛的领域和任务。如果未提供正负对,我们将构建正负对,并将它们存储在统一的格式中:
[
{'query': ['表征维基百科问题以检索相关文档;', '大谎言第二季有多少集'], 'pos': ['表征维基百科文档以供检索;', '《大谎言》(电视剧)系列获得了几个奖项。它获得了16项艾美奖提名并赢得了八项,包括杰出限制剧集和Kidman、Skarsgård和Dern的表演奖。三人还获得了金球奖,此外,Kidman和Skarsgård还获得了演员工会奖。尽管最初被定为限制剧集,但HBO续订了第二季。第二季的制作始于2018年3月,预计将于2019年首播。所有七集由Kelley编写'], 'neg': ['表征维基百科文档以供检索;', '小人物,大世界 第二季A结局的最后几分钟,在投石机附近,一群人围着Jacob,他躺在地上。第二季B的前两集集中在事故和当地媒体对它的反应上。《小人物,大世界》的第一季为TLC带来了不错的收视率(尤其是在18-49岁重要人口统计中),导致该剧续订了第二季。该系列被普遍认为是对小人物的正面描绘。相反,一些评论认为该剧有着窥视的倾向'], 'task_id': 1}
{'query': ['表征维基百科问题以检索相关文档;', '谁唱了《期待一个女孩像你》'], 'pos': ['表征维基百科文档以供检索;', '期待一个女孩像你 期待一个女孩像你 这是一首1981年的力量情歌,由英国-美国摇滚乐队Foreigner演唱。独特的合成器主题由那时尚不知名的Thomas Dolby演奏,这首歌也标志着他们早期单曲的大转变,因为他们之前的单曲是中快节奏摇滚歌曲,而这首歌是一首柔和的爱情曲,带有力量情歌的能量。它是从专辑“4”(1981)中发布的第二首单曲,由Lou Gramm和Mick Jones共同创作。这成为了乐队最大的热门歌曲之一'], 'neg': ['表征维基百科文档以供检索;', '期待一个女孩像你 被Olivia Newton-John的单曲“Physical”连续九周推开了榜首位置,然后被Hall & Oates的“I Can’t Go for That (No Can Do)”在1982年1月30日第十周的榜首位置推开。由于其榜单长寿,它成为了1982年百大单曲中的第19首。这首歌成为了乐队最大的热门歌曲,直到1985年“I Want to Know What Love Is”达到了榜首。这首歌在“Billboard's Greatest Songs of All Time”中排第100'], 'task_id': 1}
...
{'query': ['表征维基百科句子以检索相关文档;', '我喜欢甜马提尼饮料!'], 'pos': ['表征维基百科文档以供检索;', "苹果马提尼 苹果马提尼\n一种苹果马提尼(简称Appletini)是一种含有伏特加和一种或多种苹果汁、苹果酒、苹果利口酒或苹果白兰地的鸡尾酒。\n这种饮料,最初被称为Adam's Apple Martini,因为创造它的酒保名叫Adam,于1996年在西好莱坞的Lola餐厅创立。\n这款饮料,Adam's Apple,于1972年7月的《花花公子》杂志内封面广告中推广。尽管食谱中包括一盎司或更多的"], 'neg': ['表征维基百科文档以供检索;', "芳香型葡萄酒 类似饮料在该法律中描述为“芳香型基酒饮品”(未强化)和“芳香型基酒鸡尾酒”(混合,酒精含量低于7%)。\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 {output_directory} --cache_dir {cache_directory} --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
:预训练的检查点。我们支持模型ID(例如,sentence-transformers/gtr-t5-large
,sentence-transformers/sentence-t5-large
)或检查点路径(例如,通过 transformers trainer 保存的检查点)。--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
并运行所有 MTEB 数据集来评估训练好的模型检查点,并通过更改 --task_name
检查 我们的论文 或 MTEB 基准 获取所有任务的评估指标。
Billboard
要评估模型在 Billboard 上的表现,运行以下命令:
cd evaluation/text_evaluation
python main.py --model_name hkunlp/instructor-large --task mscoco --add_prompt
您可以通过指定 --model_name
并运行所有 Billboard 数据集来评估训练好的模型检查点,并通过更改 --task
在 Billboard 的所有三个数据集中,我们报告 Pearson 相关性。
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 中转换为“文本到文本”格式,并报告 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: wearable person tracking in multi-floor environments"
instruction = "Represent the Science title:"
embeddings = qmodel.encode([[instruction,sentence]])
# 您还可以归一化嵌入: normalize_embeddings=True
print(f"Quantized Embeddings:\n {embeddings}")
这将使模型大小减少10倍,并且推理时间将比普通模型少 :)
错误或问题
如果您有任何关于代码或论文的问题,请随时邮件联系 Hongjin(hjsu@cs.hku.hk')和 Weijia(
swj0419@cs.washington.edu`)。请尽量详细说明问题,以便我们能更好更快地帮助您。
引用
如果您认为我们的工作对您有帮助,请引用我们:
@inproceedings{INSTRUCTOR,
title={One Embedder, Any Task: Instruction-Finetuned Text Embeddings},
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 支持 InstructEmbeddings,使用 INSTRUCTOR 模型。
- MosaicML 已经包含了 Instructor-Large 和 Instructor-XL
- embaas 集成了 Instructor-Large
- Haystack 包含
InstructorTextEmbedder
和InstructorDocumentEmbedder
组件。