SFR Embedding Code 400M R

🚀 Salesforce/SFR-Embedding-Code-400M_R

SFR-Embedding 由 Salesforce Research 研发。

Salesforce/SFR-Embedding-Code 是一个通用的嵌入模型家族，适用于多语言和多任务的代码与文本检索。在多个代码检索任务中，与各种开源代码嵌入模型相比，它展现出了卓越的性能。

查看我们的 论文 以了解更多详情！

🚀 快速开始

本项目提供了在不同框架下使用模型的示例，以下将详细介绍如何运行代码。

Transformers 框架

import torch.nn.functional as F
from transformers import AutoModel, AutoTokenizer

input_texts = [
    "how to implement quick sort in Python?",
    "def quick_sort(arr):\n    if len(arr) <= 1:\n        return arr\n    pivot = arr[len(arr) // 2]\n    left = [x for x in arr if x < pivot]\n    middle = [x for x in arr if x == pivot]\n    right = [x for x in arr if x > pivot]\n    return quick_sort(left) + middle + quick_sort(right)",
    "def bubble_sort(arr):\n    n = len(arr)\n    for i in range(n):\n        for j in range(0, n-i-1):\n            if arr[j] > arr[j+1]:\n                arr[j], arr[j+1] = arr[j+1], arr[j]\n    return arr",
]

model_path = 'Salesforce/SFR-Embedding-Code-400M_R'
tokenizer = AutoTokenizer.from_pretrained(model_path)
model = AutoModel.from_pretrained(model_path, trust_remote_code=True)

# Tokenize the input texts
batch_dict = tokenizer(input_texts, max_length=8192, padding=True, truncation=True, return_tensors='pt')

outputs = model(**batch_dict)
embeddings = outputs.last_hidden_state[:, 0]

# normalize embeddings
embeddings = F.normalize(embeddings, p=2, dim=1)
scores = (embeddings[:1] @ embeddings[1:].T) * 100
print("Similarity Scores:", scores.tolist())
# Similarity Scores: [[74.84745025634766, 65.39266967773438]]

Sentence Transformers 框架

# Requires sentence_transformers>=2.7.0
from sentence_transformers import SentenceTransformer
from sentence_transformers.util import cos_sim

sentences = [
    "how to implement quick sort in Python?",
    "def quick_sort(arr):\n    if len(arr) <= 1:\n        return arr\n    pivot = arr[len(arr) // 2]\n    left = [x for x in arr if x < pivot]\n    middle = [x for x in arr if x == pivot]\n    right = [x for x in arr if x > pivot]\n    return quick_sort(left) + middle + quick_sort(right)",
    "def bubble_sort(arr):\n    n = len(arr)\n    for i in range(n):\n        for j in range(0, n-i-1):\n            if arr[j] > arr[j+1]:\n                arr[j], arr[j+1] = arr[j+1], arr[j]\n    return arr",
]

model = SentenceTransformer('Salesforce/SFR-Embedding-Code-400M_R', trust_remote_code=True)
embeddings = model.encode(sentences)
similarities = cos_sim(embeddings[0], embeddings[1:])
print(similarities)
# tensor([[0.7485, 0.6539]])

✨ 主要特性

• 通用性：是一个通用的嵌入模型家族，适用于多语言和多任务的代码与文本检索。
• 高性能：在多个代码检索任务中，与各种开源代码嵌入模型相比，展现出了卓越的性能。

📚 详细文档

道德考量

本版本仅用于支持学术论文的研究目的。我们的模型、数据集和代码并非专门为所有下游用途而设计或评估。我们强烈建议用户在部署此模型之前，评估并解决与准确性、安全性和公平性相关的潜在问题。我们鼓励用户考虑人工智能的常见局限性，遵守适用法律，并在选择用例时采用最佳实践，特别是在错误或滥用可能对人们的生活、权利或安全产生重大影响的高风险场景中。有关用例的更多指导，请参阅我们的 AUP 和 AI AUP。

许可证声明

用户需要自行评估与原始数据集和数据的相应许可证或条款和条件相关的任何义务或责任。本版本仅用于支持学术论文的研究目的。

CoIR 基准测试性能

📄 许可证

本项目采用 CC BY-NC 4.0 许可证。

📖 引用

@article{liu2024codexembed,
  title={CodeXEmbed: A Generalist Embedding Model Family for Multiligual and Multi-task Code Retrieval},
  author={Liu, Ye and Meng, Rui and Jot, Shafiq and Savarese, Silvio and Xiong, Caiming and Zhou, Yingbo and Yavuz, Semih},
  journal={arXiv preprint arXiv:2411.12644},
  year={2024}
}

👥 团队成员

SFR-Embedding 团队（† 表示共同负责人）

• Ye Liu
• Rui Meng
• Shafiq Rayhan Joty
• Silvio Savarese
• Caiming Xiong †
• Yingbo Zhou †
• Semih Yavuz †