SFR Embedding Code 400M R

🚀 Salesforce/SFR-Embedding-Code-400M_R

Salesforce ResearchによるSFR-Embeddingです。

Salesforce/SFR-Embedding-Codeは、多言語およびマルチタスクのコードとテキスト検索用の汎用的な埋め込みモデルファミリーです。複数のコード検索タスクにおいて、様々なオープンソースのコード埋め込みモデルと比較して卓越した性能を示しています。

詳細については、論文をご覧ください！

🚀 クイックスタート

倫理的な考慮事項

このリリースは、学術論文の支援のための研究目的のみです。当社のモデル、データセット、およびコードは、すべての下流の目的に特に設計または評価されていません。ユーザーは、このモデルをデプロイする前に、精度、安全性、および公正性に関連する潜在的な懸念事項を評価し、対処することを強く推奨します。ユーザーは、AIの一般的な制限を考慮し、適用される法律を遵守し、特にエラーや誤用が人々の生活、権利、または安全に重大な影響を与える可能性のある高リスクシナリオにおいて、ユースケースを選択する際にベストプラクティスを活用することをお勧めします。ユースケースに関するさらなるガイダンスについては、AUPおよびAI AUPを参照してください。

ライセンス声明

ユーザーは、元のデータセットおよびデータに関する対応するライセンスまたは利用規約に基づく義務や責任について、独自の評価を行う必要があります。このリリースは、学術論文の支援のための研究目的のみです。

CoIRベンチマークでの性能

SFR-Embeddingチーム (†は共同リーダーを示します)

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

💻 使用例

基本的な使用法

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]])

📄 ライセンス

@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}
}