CodeBERT-base-mlm

CodeBERT: 面向编程与自然语言的预训练模型的预训练权重

训练数据

模型基于CodeSearchNet代码语料库训练

训练目标

该模型以Roberta-base初始化，采用MLM（掩码语言建模）目标进行训练

使用示例

from transformers import RobertaTokenizer, RobertaForMaskedLM, pipeline  

model = RobertaForMaskedLM.from_pretrained('microsoft/codebert-base-mlm')  
tokenizer = RobertaTokenizer.from_pretrained('microsoft/codebert-base-mlm')  

code_example = "if (x is not None) <mask> (x>1)"  
fill_mask = pipeline('fill-mask', model=model, tokenizer=tokenizer)  

outputs = fill_mask(code_example)  
print(outputs)  

预期输出：

{'sequence': '<s> if (x is not None) and (x>1)</s>', 'score': 0.6049249172210693, 'token': 8}  
{'sequence': '<s> if (x is not None) or (x>1)</s>', 'score': 0.30680200457572937, 'token': 50}  
{'sequence': '<s> if (x is not None) if (x>1)</s>', 'score': 0.02133703976869583, 'token': 114}  
{'sequence': '<s> if (x is not None) then (x>1)</s>', 'score': 0.018607674166560173, 'token': 172}  
{'sequence': '<s> if (x is not None) AND (x>1)</s>', 'score': 0.007619690150022507, 'token': 4248}  

相关资源

1. 采用MLM+RTD目标训练的双模态CodeBERT（适用于代码搜索与文档生成）
2. 🤗 Hugging Face的CodeBERTa轻量版（6层精简架构）

引用文献

@misc{feng2020codebert,  
    title={CodeBERT: 面向编程与自然语言的预训练模型},  
    author={张寅峰 and 郭大禹 and 唐杜宇 and 段楠 and 冯晓成 and 龚铭 and 寿林钧 and 覃兵 and 刘挺 and 蒋达新 and 周明},  
    year={2020},  
    eprint={2002.08155},  
    archivePrefix={arXiv},  
    primaryClass={cs.CL}  
}