codesage-small - 多语言支持的开源码理解与嵌入模型

codesage-small 项目介绍

项目简介

CodeSage 是一组新型的开放代码嵌入模型，采用编码器架构，支持广泛的源码理解任务。这个项目在一篇名为《Code Representation Learning At Scale》的论文中被首次介绍。该论文的主要作者包括Dejiao Zhang、Wasi Uddin Ahmad、Ming Tan、Hantian Ding、Ramesh Nallapati、Dan Roth、Xiaofei Ma和Bing Xiang。每位作者对论文的贡献是均等的。

预训练数据

CodeSage-Small 使用了来自 The Stack 的数据集进行训练，该数据集可以在 Hugging Face 上找到。支持的编程语言共有九种，分别是：C、C#、Go、Java、JavaScript、TypeScript、PHP、Python 和 Ruby。

训练过程

该模型的训练分为两个阶段。首先，通过掩码语言模型（MLM）技术对代码数据进行训练。然后，在文本-代码对的数据上进行进一步训练。具体的技术细节可以参考前述的论文进行了解。

如何使用

CodeSage-Small 提供了一个编码器（130M参数规模），可以用于提取1024维的代码嵌入。这一模型可以通过 transformers 库的 AutoModel 功能方便地加载，并使用 Starcoder tokenizer 进行分词。

以下是简单的代码示例：

from transformers import AutoModel, AutoTokenizer

checkpoint = "codesage/codesage-small"
device = "cuda"  # 使用GPU时用"cuda"，使用CPU时用"cpu"

# 注意：CodeSage 在每个分词结果序列的末尾需要加上 eos token，以确保良好的性能
tokenizer = AutoTokenizer.from_pretrained(checkpoint, trust_remote_code=True, add_eos_token=True)

model = AutoModel.from_pretrained(checkpoint, trust_remote_code=True).to(device)

inputs = tokenizer.encode("def print_hello_world():\tprint('Hello World!')", return_tensors="pt").to(device)
embedding = model(inputs)[0]
print(f'嵌入的维度: {embedding[0].size()}')
# 嵌入的维度: torch.Size([14, 1024])

引用信息

如果在研究中使用了 CodeSage-Small，请引用以下参考文献：

@inproceedings{
    zhang2024code,
    title={{CODE} {REPRESENTATION} {LEARNING} {AT} {SCALE}},
    author={Dejiao Zhang and Wasi Uddin Ahmad and Ming Tan and Hantian Ding and Ramesh Nallapati and Dan Roth and Xiaofei Ma and Bing Xiang},
    booktitle={The Twelfth International Conference on Learning Representations},
    year={2024},
    url={https://openreview.net/forum?id=vfzRRjumpX}
}

通过这样一个架构，CodeSage-Small 为代码理解任务提供了一种高效且灵活的解决方案。它利用深度学习的技术从大规模代码数据中学习丰富的表示，对于相关研究人员和开发者都是一种强大的工具。