CodeGemma: Open Code Models Based on Gemma

最近几天，对大模型生成代码的能力比较感兴趣，特此学习了一下大模型怎么在IDE上进行，自动补全你想要的代码。
大模型源码理解-以Gemma为例子
最近也在使用阿里的通义灵码，感觉还挺好用的。通义灵码

摘要：

CodeGemma基于谷歌开源的文本大模型Gemma，结构上跟LLama差不太多。CodeGemma（或者使用在IDE中的大模型）常规的使用GPT或者其它大模型，写prompt去生成代码还是略有区别的，下面就是常规的使用大模型写代码：

但是在IDE中，不能写prompt让大模型，写代码，这么太方便了，应该类似于自动补全，猜你想写什么代码，增加效率，下面是通义灵码的效果：

发布版本

公开了三个版本：

• CodeGemma 7B 预训练版本
• CodeGemma 7B 指令微调版本
• CodeGemma 2B 预训练版本
  
  7B和2B训练数据略有差距，7B进行4：1代码数据和自然语言数据进行混合。

如何实现IDE自动生成代码

IDE自动生成代码，和常规的大预言模型训练方式基本上使一致的，或者说和GPT、LLama的训练任务是一致的，但是为了实现你写几个字或者代码，大模型可以自动生成你需要代码。但是在数据构造的时候和常规的方式不一致，则是将一个代码文本信息拆分成三部分，Prefix，Middle，Suffix即三段式，开头，中间，结尾，这三个部分。数据在使用的时候有两种改造方式：

• PSM Prefix-Suffix-Middle
• SPM Suffix-Prefix-Middle
  在训练的时候会增加一下特殊的token进行分隔开，如下图所示：
  
  PSM具体样例：
  

实践Prompt:

在具体使用时候，采用以下形式prompt。

简单代码实现：

BEFORE_CURSOR = "<|fim_prefix|>"
AFTER_CURSOR = "<|fim_suffix|>"
AT_CURSOR = "<|fim_middle|>"
FILE_SEPARATOR = "<|file_separator|>"

stop_token
END_TOKEN = gemma_lm.preprocessor.tokenizer.end_token

stop_tokens = (BEFORE_CURSOR, AFTER_CURSOR, AT_CURSOR, FILE_SEPARATOR, END_TOKEN)

stop_token_ids = tuple(gemma_lm.preprocessor.tokenizer.token_to_id(x) for x in stop_tokens)
def format_completion_prompt(before, after):
    return f"{BEFORE_CURSOR}{before}{AFTER_CURSOR}{after}{AT_CURSOR}"

before = "import "
after = """if __name__ == "__main__":\n    sys.exit(0)"""
prompt = format_completion_prompt(before, after)
print(prompt)

实验结果

简单体验：

在kaggle平台上简单体验了，网址，测试的时候最大的感触，对于代码生成，或者基于大模型代码自动补全，prompt得严格按照，推荐的prompt，要不效果很差，或者可以说没效果。
一开始prompt采用以下方式，就随便用了一个别人测试方式，进行自己得需求魔改：

def format_completion_prompt(before, after):
    return f"{BEFORE_CURSOR}{before}{AFTER_CURSOR}{after}{AT_CURSOR}"

before = "import "
after = """def sort if __name__ == "__main__":\n    sys.exit(0)"""
prompt = format_completion_prompt(before, after)
print(prompt)

无法实现我们需要的排序方式， 只输出了sys

严格按照推荐的prompt方式

def format_completion_prompt(test):
    return f"""<|fim_prefix|>import <|fim_suffix|>if __name__ == "__main__":
    sys.exit(0)<|fim_middle|>{test}"""

test = "def quick_sort"
prompt = format_completion_prompt(test)
print(prompt)