🚀 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)
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]
embeddings = F.normalize(embeddings, p=2, dim=1)
scores = (embeddings[:1] @ embeddings[1:].T) * 100
print("Similarity Scores:", scores.tolist())
Sentence Transformers 框架
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)
✨ 主要特性
- 通用性:是一個通用的嵌入模型家族,適用於多語言和多任務的代碼與文本檢索。
- 高性能:在多個代碼檢索任務中,與各種開源代碼嵌入模型相比,展現出了卓越的性能。
📚 詳細文檔
道德考量
本版本僅用於支持學術論文的研究目的。我們的模型、數據集和代碼並非專門為所有下游用途而設計或評估。我們強烈建議用戶在部署此模型之前,評估並解決與準確性、安全性和公平性相關的潛在問題。我們鼓勵用戶考慮人工智能的常見侷限性,遵守適用法律,並在選擇用例時採用最佳實踐,特別是在錯誤或濫用可能對人們的生活、權利或安全產生重大影響的高風險場景中。有關用例的更多指導,請參閱我們的 AUP 和 AI AUP。
許可證聲明
用戶需要自行評估與原始數據集和數據的相應許可證或條款和條件相關的任何義務或責任。本版本僅用於支持學術論文的研究目的。
CoIR 基準測試性能
模型 |
模型大小 |
CoIR AVG (NDCG@10) |
SFR-Embedding-Code |
2B |
67.4 |
CodeSage-Large-v2 |
1.3B |
64.2 |
CodeSage-Large |
1.3B |
61.0 |
SFR-Embedding-Code |
400M |
61.9 |
CodeRankEmbed |
137M |
60.1 |
CodeSage-Base |
356M |
57.5 |
Voyage-Code-002 |
- |
56.3 |
CodeSage-Small |
130M |
54.4 |
📄 許可證
本項目採用 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 †