LLM大模型: 源代码语义相似度检测
1、源代码相似度检测的用途很多，比如：
- 代码抄袭克隆
- 高危/漏洞代码检测
- 软件成分的分析
以往的字符串/文本方法简单粗暴：直接用字符串正则匹配或其变种(比如字符串的指纹、字符串的Levenshtein 编辑距离等)的方式检测，这种方式的缺点也很明显：抄袭者不会傻到直接ctrl+c、ctrl+v，或多或少都会对代码做各种更改，这种情况直接匹配字符串的方式就会失效；根本原因就是没能提取代码深层次的语义特征，只使用了表层的代码符号；代码符号一旦改变，即使底层的语义没变也检测不出来啦！
2、提取普通文本的语义特征简单粗暴：直接上transformer的encoder就完了！我首次尝试了bert-base-uncased、google--mt5-small、gpt2、Langboat--bloom-1b4-zh、longformer-base-4096等LLM后，效果出人意料地非常差，细想了一下，发现这些模型的训练语料都是普通的语料，针对普通文本处理的效果还不错，但源代码和普通文本的语义相差较大，比如：
- 普通的文本不论中文、英语，词的数量都有限，所以理论上讲vocab是能穷举所有词的，但源码就不一样了：源码的类名、函数名、变量名是用户自定义的，vocab大概率是没有的，这部分的token是没法通过encoder转成向量的，语义自然也无从提取！
- 普通文本token的顺序一旦发生变化，语义很有可能就变了（当然也有倒装，但毕竟是少数），比如“我打你”和“你打我"，这两句话所有的token都是一样的，但顺序不同，语义截然相反，所以transformer的embedding层会使用position编码来表征token的位置信息；但代码就不同了：比如变量定义，int a=1；float b=2；这两行代码的顺序完全可以改变，完全不影响最终的执行结果，所以语义也没任何改变！
针对上述的问题，我偶先想到的就是抽象语法树AST：做代码分析这个是绕不开的，所以也尝试先用AST把源码转成树形结构，然后再把AST树通过encoder转成向量后再匹配代码的相似度。demo如下：
这棵树以文本的形式打印了出来，既然是文本，肯定能够直接输入encoder编码，但是这样做效果咋样了？我尝试过了，效果还是不理想，可能的原因如下：
- 不论源码是怎样的，树节点的名称不变，这部分token是么有任何区分度的
- 控制流一旦改变，树形的结构就会变：比如return a+b和 c=a+b return c；这两个代码片段的结果逻辑完全一样，但是ast树形结构是不一样的！还有把for循环改成while循环，树形结构也会变，直接把ast输入encoder得到向量后计算向量距离的效果也不咋的！下面是这两个代码片段的ast树对比：层次结构明显不同，但代码的功能/语义确是一样的！
import ast
code1 = """ def add(): return a + b """
tree1 = ast.parse(code1)
ast.dump(tree1, indent=4)
code1的抽象语法树如下：
Module(
body=[
FunctionDef(
name='add',
args=arguments(
posonlyargs=[],
args=[],
vararg=None,
kwonlyargs=[],
kw_defaults=[],
kwarg=None,
defaults=[]
),
body=[
Return(
value=BinOp(
left=Name(id='a', ctx=Load()),
op=Add(),
right=Name(id='b', ctx=Load())
)
)
],
decorator_list=[],
returns=None,
type_comment=None
)
],
type_ignores=[]
)
code2 = """ def add(): c = a + b return c """ tree2 = ast.parse(code2) ast.dump(tree2, indent=4) code2的抽象语法树如下：
Module( body=[ FunctionDef( name='add', args=arguments( posonlyargs=[], args=[], vararg=None, kwonlyargs=[], kw_defaults=[], kwarg=None, defaults=[] ), body=[ Assign( targets=[ Name(id='c', ctx=Store()) ], value=BinOp( left=Name(id='a', ctx=Load()), op=Add(), right=Name(id='b', ctx=Load()) ) ), Return( value=Name(id='c', ctx=Load()) ) ], decorator_list=[], returns=None, type_comment=None ) ], type_ignores=[] )
3、（1）代码的核心功能就是处理各种数据，所以这类技术的原始名称叫information technology，后续演变出来了很多职业：后台开发是crud boy，数仓开发是sql boy，本质就是处理各种数据；如果不论是什么代码，只要处理数据的输出结果是一样，是不是代表的这些代码的功能/语义是一样的了？这就是程序分析最核心的功能之一：数据流分析！核心分析思路就是：数据从哪来？经过了什么处理？输出了什么？关注的是数据的流转，而不是代码本身的结构！
相比通过ast树分析代码，数据流分析的优势：
- 只考虑数据的处理和流转，梳理清楚了“数值从哪里来”的关系！而ast会深入分析代码的结构层次，这里其实是没必要的
所以接下来会采用数据流的方式分析代码片段的相似性！这里采用微软的 microsoft/graphcodebert-base 模型。GraphCodeBERT 不仅考虑了代码的序列信息，还考虑了代码的结构信息。它使用数据流Data Flow Graph（一种代码的语义级结构，编码了变量之间的“值从何处来”的关系）而不是抽象语法树（AST）来获取代码的结构信息。这种语义级结构简洁明了，不会带来 AST 的不必要深层次结构，使得模型更加高效。此外，GraphCodeBERT 还引入了两个结构感知的预训练任务，一个是预测代码结构边，另一个是对齐源代码和代码结构的表示。该模型的详细信息详见文章末尾的参考链接1；
和普通的LLM比，GraphCodeBERT的差异之处在于：
- 把注释、代码、DFG分别作为输入进入encoder，这里明显多出了DFG； 更多细节参考：https://blog.csdn.net/qq_44370676/article/details/114384343
DFG的信息在attention层的处理方式： https://blog.csdn.net/qq_44370676/article/details/114384343
（2）demo代码如下：先加载这个LLM
import torch from transformers import RobertaTokenizer, RobertaModel from sklearn.metrics.pairwise import cosine_similarity # 使用GraphCodeBERT模型和Tokenizer model_name = "/root/huggingface/graphcodebert-base" tokenizer = RobertaTokenizer.from_pretrained(model_name) model = RobertaModel.from_pretrained(model_name)
再调用大模型的encoder计算代码片段的向量，每个token都采用最后一个hidden layer的输出，然后取平均：
def encode_code(code, model, tokenizer): inputs = tokenizer(code, return_tensors='pt', max_length=512, truncation=True, padding='max_length') with torch.no_grad(): outputs = model(**inputs) last_hidden_state = outputs.last_hidden_state vector = last_hidden_state.mean(dim=1).squeeze() return vector.numpy()
最后直接用cosin计算向量的距离：
1 code1 = """ 2 void CWE369_Divide_by_Zero__float_fscanf_04_bad() 3 { 4 float data; 5 /* Initialize data */ 6 data = 0.0F; 7 if(STATIC_CONST_TRUE) 8 { 9 /* POTENTIAL FLAW: Use a value input from the console using fscanf() */ 10 fscanf (stdin, \"%f\", &data); 11 } 12 if(STATIC_CONST_TRUE) 13 { 14 { 15 /* POTENTIAL FLAW: Possibly divide by zero */ 16 int result = (int)(100.0 / data); 17 printIntLine(result); 18 } 19 } 20 } 21 """ 22 23 code2 = """ 24 void good() 25 { 26 int divid = 0.0F; 27 const int a = 1; 28 while(a) 29 { 30 fscanf (stdin, \"%i\", &divid); 31 } 32 while(a) 33 { 34 { 35 int output = (int)(100.0 / data); 36 printIntLine(output); 37 } 38 } 39 } 40 """ 41 42 # 编码代码片段 43 vector1 = encode_code(code1, model, tokenizer) 44 vector2 = encode_code(code2, model, tokenizer) 45 46 # 计算余弦相似度 47 similarity = cosine_similarity([vector1], [vector2])[0][0] 48 print(f"Semantic similarity: {similarity}")
两个代码片段：更改函数名、变量类型、变量名、控制方式，去掉注释，cosin相识度高达0.93！
注意：
1、GraphCodeBERT 只能处理最多512个token，超过这个数量，要么分片多次调用encoder，要么换其他LLM；
2、实际使用时可配合FAISS这类向量数据库，效果更佳
3、汇编代码：二进制逆向的时候需要解析汇编代码的功能，所以也需要能识别汇编代码的功能，方式之一就是求代码的相似度。比如把可疑的汇编代码片段反编译成c伪代码后，和常见的加密算法源码计算相似度，就能大概猜测这段汇编代码所使用的sign算法！
4、https://github.com/microsoft/CodeBERT/blob/master/GraphCodeBERT/clonedetection/run.py 实际inference的时候，需要先提取代码的DFG，官方的方式如下：
#remove comments, tokenize code and extract dataflow def extract_dataflow(code, parser,lang): #remove comments try: code=remove_comments_and_docstrings(code,lang) except: pass #obtain dataflow if lang=="php": code="<?php"+code+"?>" try: tree = parser[0].parse(bytes(code,'utf8')) root_node = tree.root_node tokens_index=tree_to_token_index(root_node) code=code.split('\n') code_tokens=[index_to_code_token(x,code) for x in tokens_index] index_to_code={} for idx,(index,code) in enumerate(zip(tokens_index,code_tokens)): index_to_code[index]=(idx,code) try: DFG,_=parser[1](root_node,index_to_code,{}) except: DFG=[] DFG=sorted(DFG,key=lambda x:x[1]) indexs=set() for d in DFG: if len(d[-1])!=0: indexs.add(d[1]) for x in d[-1]: indexs.add(x) new_DFG=[] for d in DFG: if d[1] in indexs: new_DFG.append(d) dfg=new_DFG except: dfg=[] return code_tokens,dfg
5、选择大模型的时候，可以看看vocab.json文件有没有自己任务涉及到的token；如果没有，说明这个大模型就没有提取这些token的语义，用这个大模型开发特定应用效果肯定很差！
总结：
大模型的transformer block由attention和FFN构成，FFN十几年前基本定型，暂时没发现有突破的创新；但是attention（及embedding层）就不同了，完全可以根据不同的任务类型（比如某些垂直领域，如code相关的任务）有不同类型的输入，和设计不同的attention/embedding机制；比如structcode的attention机制如下：https://github.com/reddy-lab-code-research/structcoder
代码相关大模型举例：
参考：
1、https://arxiv.org/abs/2009.08366 GraphCodeBERT: Pre-training Code Representations with Data Flow
2、https://huggingface.co/microsoft/graphcodebert-base GraphCodeBERT is a graph-based pre-trained model based on the Transformer architecture for programming language
3、https://blog.csdn.net/nbyvy/article/details/127371318 https://blog.csdn.net/qq_44370676/article/details/114384343 GraphCodeBERT 介绍
4、https://github.com/reddy-lab-code-research/structcoder https://blog.csdn.net/nbyvy/article/details/127282241 structcode设计
5、https://github.com/microsoft/CodeBERT/tree/master/GraphCodeBERT https://www.jianshu.com/p/ba75337b5575 GraphCodeBERT训练
