SFR Embedding Code 2B R

🚀 Salesforce/SFR-Embedding-Code-2B_R

Salesforce ResearchによるSFR-Embeddingです。

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

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

また、400Mサイズのモデル Salesforce/SFR-Embedding-Code-400_R も提供しています。

✨ 主な機能

• 多言語およびマルチタスクのコードとテキスト検索に対応
• 複数のコード検索タスクで、オープンソースのコード埋め込みモデルよりも優れた性能を発揮

📚 ドキュメント

倫理的な考慮事項

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

ライセンス声明

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

このリリースされたモデルはGemmaの微調整バージョンであり、Gemmaはai.google.dev/gemma/termsに記載されているGemma利用規約の下で提供され、その規約の対象となります。さらに、このモデルの使用は、ai.google.dev/gemma/prohibited_use_policyに記載されているGemma禁止使用ポリシー（「禁止使用ポリシー」）に定められた通りに制限されます。この禁止使用ポリシーは、この契約に参照により組み込まれます。

CoIRベンチマークでの性能

開発チーム

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

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

📦 インストール

このREADMEには具体的なインストール手順が記載されていないため、このセクションは省略されます。

💻 使用例

基本的な使用法

Transformers

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

# Each query needs to be accompanied by an corresponding instruction describing the task.
query_instruction_example = "Given Code or Text, retrieval relevant content"
queries = [
    "how to implement quick sort in Python?"
]

# No instruction needed for retrieval passages
passages = [
    "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"
]

# load model with tokenizer
model = AutoModel.from_pretrained('Salesforce/SFR-Embedding-Code-2B_R', trust_remote_code=True)

# get the embeddings
max_length = 32768
query_embeddings = model.encode_queries(queries, instruction=query_instruction_example, max_length=max_length)
passage_embeddings = model.encode_corpus(passages, max_length=max_length)

# normalize embeddings
query_embeddings = F.normalize(query_embeddings, p=2, dim=1)
passage_embeddings = F.normalize(passage_embeddings, p=2, dim=1)

scores = (query_embeddings @ passage_embeddings.T) * 100
print(scores.tolist())
# [[69.26929473876953, 58.41606903076172]]

Sentence Transformers

from sentence_transformers import SentenceTransformer

# Each query needs to be accompanied by an corresponding instruction describing the task.
query_instruction_example = "Instruct: Given Code or Text, retrieval relevant content\nQuery: "
queries = ["how to implement quick sort in Python?"]

# No instruction needed for retrieval passages
passages = [
    "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"
]

# Load the Sentence Transformer model, including pooling
model = SentenceTransformer('Salesforce/SFR-Embedding-Code-2B_R', trust_remote_code=True)

# Compute the embeddings for both queries and passages. Use 'prompt' for queries only
query_embeddings = model.encode(queries, prompt=query_instruction_example)
passage_embeddings = model.encode(passages)

# Compute the similarities between the queries and passages
similarities = model.similarity(query_embeddings, passage_embeddings)
print(similarities)
# tensor([[0.6927, 0.5842]])

引用

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

📄 ライセンス

このモデルはCC BY-NC 4.0ライセンスの下で提供されています。