Granite-Embedding-125m-English

新闻： 具有8192上下文长度的Granite Embedding R2模型已发布。

• granite-embedding-english-r2（149M参数）：输出嵌入维度为768，取代granite-embedding-125m-english。
• granite-embedding-small-english-r2（47M参数）：首款精简版模型，层数更少，输出嵌入维度更小（384），取代granite-embedding-30m-english。

模型概述： Granite-Embedding-125m-English是Granite Embeddings系列中的一个125M参数的密集双编码器嵌入模型，可用于生成高质量的文本嵌入。该模型生成的嵌入向量维度为768。与其他大多数开源模型相比，该模型仅使用具有宽松、企业友好许可的开源相关性配对数据集以及IBM收集和生成的数据集进行训练。在保持BEIR等学术基准测试竞争力的同时，该模型在许多企业应用场景中也有出色表现。该模型采用检索导向的预训练、对比微调和知识蒸馏技术开发。

• 开发者： Granite Embedding Team, IBM
• GitHub仓库： ibm-granite/granite-embedding-models
• 网站：Granite Docs
• 论文： Technical Report
• 发布日期：2024年12月18日
• 许可证： Apache 2.0

支持的语言： 英语。

预期用途： 该模型旨在为给定文本生成固定长度的向量表示，可用于文本相似度、检索和搜索应用。

与Sentence Transformers配合使用： 该模型与SentenceTransformer库兼容，使用非常简单：

首先，安装sentence transformers库

pip install sentence_transformers

然后可以使用该模型对文本对进行编码并找出其表示之间的相似度

from sentence_transformers import SentenceTransformer, util

model_path = "ibm-granite/granite-embedding-125m-english"
# Load the Sentence Transformer model
model = SentenceTransformer(model_path)

input_queries = [
    ' Who made the song My achy breaky heart? ',
    'summit define'
    ]

input_passages = [
    "Achy Breaky Heart is a country song written by Don Von Tress. Originally titled Don't Tell My Heart and performed by The Marcy Brothers in 1991. ",
    "Definition of summit for English Language Learners. : 1 the highest point of a mountain : the top of a mountain. : 2 the highest level. : 3 a meeting or series of meetings between the leaders of two or more governments."
    ]

# encode queries and passages
query_embeddings = model.encode(input_queries)
passage_embeddings = model.encode(input_passages)

# calculate cosine similarity
print(util.cos_sim(query_embeddings, passage_embeddings))

与Huggingface Transformers配合使用： 以下是一个简单的示例，展示如何将Granite-Embedding-125m-English模型与Transformers库和PyTorch配合使用。

首先，安装所需的库

pip install transformers torch

然后可以使用该模型对文本对进行编码

import torch
from transformers import AutoModel, AutoTokenizer

model_path = "ibm-granite/granite-embedding-125m-english"

# Load the model and tokenizer
model = AutoModel.from_pretrained(model_path)
tokenizer = AutoTokenizer.from_pretrained(model_path)
model.eval()

input_queries = [
    ' Who made the song My achy breaky heart? ',
    'summit define'
    ]

# tokenize inputs
tokenized_queries = tokenizer(input_queries, padding=True, truncation=True, return_tensors='pt')

# encode queries
with torch.no_grad():
    # Queries
    model_output = model(**tokenized_queries)
    # Perform pooling. granite-embedding-125m-english uses CLS Pooling
    query_embeddings = model_output[0][:, 0]

# normalize the embeddings
query_embeddings = torch.nn.functional.normalize(query_embeddings, p=2, dim=1)

评估：

下表报告了Granite-Embedding-125M-English模型在MTEB Retrieval（即BEIR）和代码检索（CoIR）基准测试上的性能。

模型架构： Granite-Embedding-125m-English基于类似RoBERTa的纯编码器Transformer架构，在IBM Research内部训练。

训练数据： 训练数据主要包括四个来源：（1）从网页抓取的无监督标题-正文配对数据，（2）具有宽松、企业友好许可的公开配对数据，（3）针对特定技术领域的IBM内部配对数据，（4）IBM生成的合成数据。详细数据如下：

值得注意的是，由于其非商业许可协议，我们没有在训练语料库中使用流行的MS-MARCO检索数据集，而其他开源模型因该数据集的高质量而对其进行训练。

基础设施： 我们使用IBM的计算集群Cognitive Compute Cluster来训练Granite Embedding Models，该集群配备了NVIDIA A100 80gb GPU。该集群提供了可扩展且高效的基础设施，支持我们在多个GPU上训练模型。

伦理考量和局限性： 用于训练基础语言模型的数据经过过滤，以去除包含仇恨、辱骂和脏话的文本。Granite-Embedding-125m-English仅针对英语文本进行训练，上下文长度为512个token（更长的文本将被截断至此长度）。

资源

• ⭐️ 了解Granite的最新更新：https://www.ibm.com/granite
• 📄 开始使用教程、最佳实践和提示工程建议：https://www.ibm.com/granite/docs/
• 💡 了解最新的Granite学习资源：https://ibm.biz/granite-learning-resources

<!-- ## Citation

@misc{granite-embedding-models,
  author = {author 1, author2, ...},
  title = {},
  journal = {},
  volume = {},
  year = {2024},
  url = {https://arxiv.org/abs/0000.00000},
}
``` -->