🚀 🦉CodeModernBERT-Owl
CodeModernBERT-Owl 是一个专门为代码搜索和代码理解任务从头设计的预训练模型。与之前的版本如 CodeHawks-ModernBERT 和 CodeMorph-ModernBERT 相比,该模型现在支持 Rust 语言,并提高了在 Python、PHP、Java、JavaScript、Go 和 Ruby 等语言中的搜索准确性。
🚀 快速开始
这个模型可以使用 Hugging Face Transformers 库轻松加载。
⚠️ 重要提示
需要 transformers >= 4.48.0
🔗 Colab 演示(替换为 "CodeModernBERT-Owl")
模型的加载
from transformers import AutoModelForMaskedLM, AutoTokenizer
tokenizer = AutoTokenizer.from_pretrained("Shuu12121/CodeModernBERT-Owl")
model = AutoModelForMaskedLM.from_pretrained("Shuu12121/CodeModernBERT-Owl")
获取代码嵌入
import torch
def get_embedding(text, model, tokenizer, device="cuda"):
inputs = tokenizer(text, return_tensors="pt", truncation=True, max_length=256)
if "token_type_ids" in inputs:
inputs.pop("token_type_ids")
inputs = {k: v.to(device) for k, v in inputs.items()}
outputs = model.model(**inputs)
embedding = outputs.last_hidden_state[:, 0, :]
return embedding
embedding = get_embedding("def my_function(): pass", model, tokenizer)
print(embedding.shape)
✨ 主要特性
- ✅ 支持长达 2048 个标记的长序列(相比微软 512 个标记的模型)
- ✅ 针对代码搜索、代码理解和代码克隆检测进行了优化
- ✅ 在 GitHub 开源仓库(Java、Rust)上进行了微调
- ✅ 在 CodeHawks/CodeMorph 系列中实现了最高的准确性
- ✅ 多语言支持:Python、PHP、Java、JavaScript、Go、Ruby 和 Rust
📚 详细文档
模型参数
属性 |
详情 |
模型类型 |
预训练模型 |
训练数据 |
Shuu12121/rust-codesearch-dataset-open、Shuu12121/java-codesearch-dataset-open、code-search-net/code_search_net、google/code_x_glue_ct_code_to_text |
vocab_size |
50,004 |
hidden_size |
768 |
num_hidden_layers |
12 |
num_attention_heads |
12 |
intermediate_size |
3,072 |
max_position_embeddings |
2,048 |
type_vocab_size |
2 |
hidden_dropout_prob |
0.1 |
attention_probs_dropout_prob |
0.1 |
local_attention_window |
128 |
rope_theta |
160,000 |
local_attention_rope_theta |
10,000 |
评估结果
数据集
- 📌 在
code_x_glue_ct_code_to_text
上进行测试,候选池大小为 100。
- 📌 使用
Shuu12121/rust-codesearch-dataset-open
进行特定于 Rust 的评估。
主要评估指标的比较(同一种子值)
语言 |
CodeModernBERT-Owl |
CodeHawks-ModernBERT |
Salesforce CodeT5+ |
Microsoft CodeBERT |
GraphCodeBERT |
Python |
0.8793 |
0.8551 |
0.8266 |
0.5243 |
0.5493 |
Java |
0.8880 |
0.7971 |
0.8867 |
0.3134 |
0.5879 |
JavaScript |
0.8423 |
0.7634 |
0.7628 |
0.2694 |
0.5051 |
PHP |
0.9129 |
0.8578 |
0.9027 |
0.2642 |
0.6225 |
Ruby |
0.8038 |
0.7469 |
0.7568 |
0.3318 |
0.5876 |
Go |
0.9386 |
0.9043 |
0.8117 |
0.3262 |
0.4243 |
- ✅ 在所有目标语言中实现了最高的准确性。
- ✅ 使用额外微调的 GitHub 数据显著提高了 Java 的准确性。
- ✅ 优于之前的模型,特别是在 PHP 和 Go 语言中。
Rust (独自数据集)
指标 |
CodeModernBERT-Owl |
MRR |
0.7940 |
MAP |
0.7940 |
R - Precision |
0.7173 |
K 别评估指标
K |
Recall@K |
Precision@K |
NDCG@K |
F1@K |
Success Rate@K |
Query Coverage@K |
1 |
0.7173 |
0.7173 |
0.7173 |
0.7173 |
0.7173 |
0.7173 |
5 |
0.8913 |
0.7852 |
0.8118 |
0.8132 |
0.8913 |
0.8913 |
10 |
0.9333 |
0.7908 |
0.8254 |
0.8230 |
0.9333 |
0.9333 |
50 |
0.9887 |
0.7938 |
0.8383 |
0.8288 |
0.9887 |
0.9887 |
100 |
1.0000 |
0.7940 |
0.8401 |
0.8291 |
1.0000 |
1.0000 |
其他推荐模型
1. CodeSearch-ModernBERT-Owl🦉 (https://huggingface.co/Shuu12121/CodeSearch-ModernBERT-Owl)
如果您需要一个更专门用于代码搜索的模型,强烈推荐此模型。
2. CodeModernBERT-Snake🐍 (https://huggingface.co/Shuu12121/CodeModernBERT-Snake)
如果您需要一个支持更长序列或更小模型大小的预训练模型,这个模型是理想的选择。
- 最大序列长度:8192 个标记
- 更小的模型大小:约 75M 参数
3. CodeSearch-ModernBERT-Snake🐍 (https://huggingface.co/Shuu12121/CodeSearch-ModernBERT-Snake)
对于那些想要一个结合了长序列长度和代码搜索专业性的模型的人来说,这个模型是最佳选择。
📝 结论
- ✅ 在所有语言中表现出色
- ✅ 通过数据集扩充成功添加了对 Rust 的支持
- ✅ 使用更好的数据集可能会进一步提高性能
📄 许可证
📄 Apache - 2.0
📧 联系方式
📩 如有任何问题,请联系:
📧 shun0212114@outlook.jp