🚀 基于Snowflake/snowflake-arctic-embed-l-v2.0的句子转换器

这是一个基于 sentence-transformers 的模型，它在聚类数据集上对 Snowflake/snowflake-arctic-embed-l-v2.0 进行了微调。该模型可以将句子和段落映射到一个1024维的密集向量空间，可用于语义文本相似度计算和语义搜索。

Snowflake/snowflake-arctic-embed-l-v2.0 模型进一步使用韩语数据进行了训练，以提高其在 韩语检索任务 中的性能。它是一个强大的模型，在多个检索基准测试中达到了 最先进（SOTA）的性能。

🚀 快速开始

本模型可以将句子和段落映射到1024维的密集向量空间，可用于语义文本相似度计算和语义搜索。

✨ 主要特性

• 基于 Snowflake/snowflake-arctic-embed-l-v2.0 微调，在聚类数据集上训练。
• 支持韩语和英语，在韩语检索任务中表现出色。
• 可将句子和段落映射到1024维的密集向量空间，用于语义文本相似度计算和语义搜索。
• 在多个检索基准测试中达到了最先进（SOTA）的性能。

📦 安装指南

首先，你需要安装 sentence-transformers 库和 xformers 库：

pip install -U sentence-transformers
pip install xformers

💻 使用示例

基础用法

使用 sentence-transformers 库加载模型并进行推理：

from sentence_transformers import SentenceTransformer

# 加载模型
# 半精度推理时请使用bf16
model_name = 'dragonkue/snowflake-arctic-embed-l-v2.0-ko'
model = SentenceTransformer(model_name)

# 定义查询和文档
queries = ['대한민국의 수도는 어디인가?', '한글을 만든 사람은 누구인가?']
documents = ['대한민국의 수도는 서울이다.', '한글은 세종대왕이 창제하였다.']

# 计算嵌入：使用 `prompt_name="query"` 对查询进行编码！
query_embeddings = model.encode(queries, prompt_name="query") 
document_embeddings = model.encode(documents)

# 计算余弦相似度分数
scores = model.similarity(query_embeddings, document_embeddings)

# 输出结果
for query, query_scores in zip(queries, scores):
    doc_score_pairs = list(zip(documents, query_scores))
    doc_score_pairs = sorted(doc_score_pairs, key=lambda x: x[1], reverse=True)
    print("Query:", query)
    for document, score in doc_score_pairs:
        print(score, document)

高级用法

使用 transformers 包加载模型并进行推理：

import torch
from transformers import AutoModel, AutoTokenizer

# 加载模型
# 半精度推理时请使用bf16
model_name = 'dragonkue/snowflake-arctic-embed-l-v2.0-ko'
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModel.from_pretrained(model_name, add_pooling_layer=False)
model.eval()

# 定义查询和文档
query_prefix = 'query: '
queries  = ['대한민국의 수도는 어디인가?', '한글을 만든 사람은 누구인가?']
queries_with_prefix = ["{}{}".format(query_prefix, i) for i in queries]
query_tokens = tokenizer(queries_with_prefix, padding=True, truncation=True, return_tensors='pt', max_length=8192)

documents = ['대한민국의 수도는 서울이다.', '한글은 세종대왕이 창제하였다.']
document_tokens = tokenizer(documents, padding=True, truncation=True, return_tensors='pt', max_length=8192)

# 计算令牌嵌入
with torch.no_grad():
    query_embeddings = model(**query_tokens)[0][:, 0]
    document_embeddings = model(**document_tokens)[0][:, 0]

# 归一化嵌入
query_embeddings = torch.nn.functional.normalize(query_embeddings, p=2, dim=1)
document_embeddings = torch.nn.functional.normalize(document_embeddings, p=2, dim=1)

scores = torch.mm(query_embeddings, document_embeddings.transpose(0, 1))

for query, query_scores in zip(queries, scores):
    doc_score_pairs = list(zip(documents, query_scores))
    doc_score_pairs = sorted(doc_score_pairs, key=lambda x: x[1], reverse=True)
    # 输出段落和分数
    print("Query:", query)
    for document, score in doc_score_pairs:
        print(score, document)

📚 详细文档

模型详情

模型来源

• 文档：Sentence Transformers Documentation
• 仓库：Sentence Transformers on GitHub
• Hugging Face：Sentence Transformers on Hugging Face

完整模型架构

SentenceTransformer(
  (0): Transformer({'max_seq_length': 8192, 'do_lower_case': False}) with Transformer model: XLMRobertaModel 
  (1): Pooling({'word_embedding_dimension': 1024, 'pooling_mode_cls_token': True, 'pooling_mode_mean_tokens': False, 'pooling_mode_max_tokens': False, 'pooling_mode_mean_sqrt_len_tokens': False, 'pooling_mode_weightedmean_tokens': False, 'pooling_mode_lasttoken': False, 'include_prompt': True})
  (2): Normalize()
)

评估

• 评估参考了KURE GitHub仓库（https://github.com/nlpai-lab/KURE）。
• 在 MTEB 中注册的所有 韩语检索基准测试 上进行了评估。

韩语检索基准测试

• Ko-StrategyQA：一个韩语 开放域问答多跳检索数据集，从StrategyQA翻译而来。
• AutoRAGRetrieval：一个 韩语文档检索数据集，通过解析五个领域（金融、公共、医疗、法律和商业）的PDF构建。
• MIRACLRetrieval：一个基于维基百科的 韩语文档检索数据集。
• PublicHealthQA：一个专注于 医疗和公共卫生领域 的韩语 检索数据集。
• BelebeleRetrieval：一个基于FLORES-200的 韩语文档检索数据集。
• MrTidyRetrieval：一个基于维基百科的 韩语文档检索数据集。
• MultiLongDocRetrieval：一个涵盖多个领域的韩语 长文档检索数据集。
• XPQARetrieval：一个 跨领域韩语文档检索数据集。

指标

• 标准指标：NDCG@10

信息检索

基准测试之外的能力

• 支持短语查询：除了完整句子查询外，还支持基于短语的查询。例如："What products does Samsung sell?" 或 "Samsung's products"。
• 处理多样化查询格式：经过训练，能够处理各种查询格式，无论措辞如何变化。例如："Tell me about Samsung."、"I'm curious about Samsung."、"What is Samsung?"。
• 优化Markdown表格搜索：针对Markdown表格搜索进行了优化，当文档中存在表格时，可以检索嵌入在表格中的答案。
• 高效聚类：无需硬负样本即可进行高效聚类。在同一批次内对样本进行聚类，通过将 Snowflake/snowflake-arctic-embed-l-v2.0 模型的嵌入截断为256维来进行高效的聚类嵌入形成。聚类方法受到以下论文的启发：
  • Embedding And Clustering Your Data Can Improve Contrastive Pretraining
  • CONTEXTUAL DOCUMENT EMBEDDINGS
• 跨领域表现出色：在不同领域表现出色。《Arctic-Embed 2.0: Multilingual Retrieval Without Compromise》论文指出：“虽然像mE5、mGTE和BGE-M3这样的模型在MIRACL上表现出色，但它们在CLEF上的性能明显弱于我们的模型和闭源产品，这表明它们可能过度拟合了MIRACL或其基于维基百科的领域。” 根据经验，Snowflake/snowflake-arctic-embed-l-v2.0 在不同领域始终优于 BGE-M3。

偏差、风险和限制

为了防止过高的GPU使用成本，模型在训练时的最大序列长度为 1300 个令牌。因此，在像 MultiLongDocRetrieval (MLDR) 这样的基准测试中，其性能可能会下降。

之前的模型 BGE-m3-ko 在训练时的令牌长度为 1024，这限制了它在MLDR基准测试中的性能。

对于 snowflake-arctic-embed-l-v2.0-ko，如果文档长度超过 1300 个令牌或约 2500 个字符，建议考虑以下模型：

训练详情

• 损失函数：CachedGISTEmbedLoss，使用以下参数：

训练超参数

非默认超参数

• eval_strategy: steps
• per_device_train_batch_size: 20000
• per_device_eval_batch_size: 4096
• learning_rate: 2e-05
• num_train_epochs: 2
• lr_scheduler_type: warmup_stable_decay
• lr_scheduler_kwargs: {'num_decay_steps': 160}
• warmup_ratio: 0.05
• bf16: True
• batch_sampler: no_duplicates

所有超参数

框架版本

• Python: 3.10.12
• Sentence Transformers: 3.4.1
• Transformers: 4.49.0
• PyTorch: 2.6.0+cu124
• Accelerate: 1.4.0
• Datasets: 3.3.2
• Tokenizers: 0.21.0

🔧 技术细节

• 模型基于 Snowflake/snowflake-arctic-embed-l-v2.0 微调，在聚类数据集上训练。
• 使用 CachedGISTEmbedLoss 损失函数进行训练。
• 训练时使用了特定的超参数，如 eval_strategy、per_device_train_batch_size 等。
• 模型架构包括 Transformer、Pooling 和 Normalize 层。
• 支持韩语和英语，在韩语检索任务中表现出色。
• 在多个检索基准测试中达到了最先进（SOTA）的性能。

📄 许可证

Arctic采用 Apache 2.0 许可证。发布的模型可以免费用于商业目的。

📚 引用

BibTeX

Sentence Transformers

@inproceedings{reimers-2019-sentence-bert,
    title = "Sentence-BERT: Sentence Embeddings using Siamese BERT-Networks",
    author = "Reimers, Nils and Gurevych, Iryna",
    booktitle = "Proceedings of the 2019 Conference on Empirical Methods in Natural Language Processing",
    month = "11",
    year = "2019",
    publisher = "Association for Computational Linguistics",
    url = "https://arxiv.org/abs/1908.10084"
}

KURE

@misc{KURE,
  publisher = {Youngjoon Jang, Junyoung Son, Taemin Lee},
  year = {2024},
  url = {https://github.com/nlpai-lab/KURE}
}

Arctic-Embed 2.0

@article{yu2024arcticembed,
  title = "Arctic-Embed 2.0: Multilingual Retrieval Without Compromise",
  author = "Puxuan Yu, Luke Merrick, Gaurav Nuti, Daniel Campos",
  journal = "arXiv preprint arXiv:2412.04506",
  year = "2024",
  url = "https://arxiv.org/abs/2412.04506"
}

Embedding And Clustering Your Data Can Improve Contrastive Pretraining

@article{merrick2024embedding,
  title = "Embedding And Clustering Your Data Can Improve Contrastive Pretraining",
  author = "Luke Merrick",
  journal = "arXiv preprint arXiv:2407.18887",
  year = "2024",
  url = "https://arxiv.org/abs/2407.18887"
}

Contextual Document Embeddings

@article{morris2024contextual,
  title = "Contextual Document Embeddings",
  author = "John X. Morris, Alexander M. Rush",
  journal = "arXiv preprint arXiv:2410.02525",
  year = "2024",
  url = "https://arxiv.org/abs/2410.02525"
}