license: mit

ERNIE-Code

ERNIE-Code：超越以英语为中心的编程语言跨语言预训练

ERNIE-Code是一个统一的大型语言模型（LLM），连接了116种自然语言和6种编程语言。我们采用两种预训练方法进行通用跨语言预训练：从单语种自然语言或编程语言中学习模式的片段掩码语言建模；以及依赖于多种自然语言与编程语言并行数据的基于枢轴的翻译语言建模。大量实验结果表明，ERNIE-Code在代码智能的多种终端任务上（包括多语言代码到文本、文本到代码、代码到代码以及文本到文本生成）均优于以往面向编程语言或自然语言的多语言LLM。我们进一步展示了其在多语言代码摘要和文本到文本翻译任务上的零样本提示优势。

ACL 2023（发现论文） | arXiv

使用方式

import torch
from transformers import (
    AutoModelForSeq2SeqLM,
    AutoModelForCausalLM,
    AutoTokenizer
)

model_name = "baidu/ernie-code-560m"

model = AutoModelForSeq2SeqLM.from_pretrained(model_name)
tokenizer = AutoTokenizer.from_pretrained(model_name)
# 注意：可使用前述`clean_up_code_spaces`函数预处理代码


def format_code_with_spm_compatablity(line: str):
    format_dict = {
        " " : "<|space|>"
    }
    tokens = list(line)
    i = 0
    while i < len(tokens):
        if line[i] == "\n":
            while i+1 < len(tokens) and tokens[i+1] == " ":
                tokens[i+1] = format_dict.get(" ")
                i += 1
        i += 1
    formatted_line = ''.join(tokens)
    return formatted_line

"""
TYPE="code" # 定义输入类型为("code", "text")
input="arr.sort()"
prompt="将Python翻译为Java: \n%s" % (input)  # 此处填写你的提示词
"""

TYPE="text" # 定义输入类型为("code", "text")
input="快速排序"
prompt="将英语翻译为日语: \n%s" % (input)  # 此处填写你的提示词

assert TYPE in ("code", "text")

# 代码输入预处理
if TYPE=="code":
    prompt = format_code_with_spm_compatablity(prompt)

model_inputs = tokenizer(prompt, max_length=512, padding=False, truncation=True, return_tensors="pt")

model = model.cuda() # 默认设置
input_ids = model_inputs.input_ids.cuda() # 默认设置
attention_mask = model_inputs.attention_mask.cuda() # 默认设置

output = model.generate(input_ids=input_ids, attention_mask=attention_mask, 
        num_beams=5, max_length=20) # 按需调整参数

# 需根据具体需求自定义`clean_up_code_spaces`输出的后处理逻辑
output = tokenizer.decode(output.flatten(), skip_special_tokens=True)


# 代码生成后处理
def clean_up_code_spaces(s: str):
    # 后处理步骤
    # ===========================
    new_tokens = ["<pad>", "</s>", "<unk>", "\n", "\t", "<|space|>"*4, "<|space|>"*2, "<|space|>"]
    for tok in new_tokens:
        s = s.replace(f"{tok} ", tok)

    cleaned_tokens = ["<pad>", "</s>", "<unk>"]
    for tok in cleaned_tokens:
        s = s.replace(tok, "")
    s = s.replace("<|space|>", " ")
    return s
output = [clean_up_code_spaces(pred) for pred in output]

可参考序列到序列翻译代码进行微调。

也可查看PaddleNLP官方推理代码。

零样本示例

• 多语言代码到文本生成（零样本）

• 多语言文本到文本翻译（零样本）

引用文献

@inproceedings{chai-etal-2023-ernie,
    title = "{ERNIE}-Code：超越以英语为中心的编程语言跨语言预训练",
    author = "柴叶坤 和 王硕欢 和 庞超 和 孙宇 和 田浩 和 吴华",
    booktitle = "ACL 2023发现论文集",
    month = "7月",
    year = "2023",
    address = "加拿大多伦多",
    publisher = "计算语言学协会",
    url = "https://aclanthology.org/2023.findings-acl.676",
    pages = "10628--10650",
    abstract = "使用相同编程语言（PL）的软件工程师可能说不同的自然语言（NL），反之亦然，这为沟通和工作效率设置了巨大障碍。近期研究证明了生成式预训练在计算机程序中的有效性，但这些模型始终以英语为中心。本研究中，我们致力于为大型语言模型（LLM） bridging 多语言NL与多语言PL之间的鸿沟。我们发布了ERNIE-Code——一个支持116种NL和6种PL的统一预训练语言模型。我们采用两种通用跨语言预训练方法：从单语种NL或PL学习模式的片段掩码语言建模；以及依赖多NL与PL并行数据的基于枢轴的翻译语言建模。大量实验表明，ERNIE-Code在代码智能的广泛终端任务（包括多语言代码到文本、文本到代码、代码到代码及文本到文本生成）上优于以往面向PL或NL的多语言LLM。我们进一步展示了其在多语言代码摘要和文本到文本翻译任务上的零样本提示优势。我们公开了代码和预训练模型。",
}