基于預訓練模型的代碼分類研究

梁 瑤, 洪慶成, 王 霞, 謝春麗

(江蘇師范大學計算機科學與技術學院, 江蘇 徐州 221116)

0 引言(Introduction)

隨著開源代碼倉庫的出現(xiàn),網(wǎng)絡上存在大量的源代碼文件,對代碼進行功能分類可幫助開發(fā)者快速找到要復用的代碼或組件,有利于開發(fā)過程中代碼的重用、理解、查找和維護,因此代碼分類是各種軟件開發(fā)任務的基礎工作。HINDLE等[1]的研究表明程序語言和自然語言都包含豐富的統(tǒng)計特性,利用自然語言的文本分類技術能夠有效地實現(xiàn)源代碼分類?；谏疃葘W習的文本分類從淺層表示轉(zhuǎn)變到深層學習表示,分類的準確性與穩(wěn)定性得到了顯著提升。但是,目前的方法多屬于有監(jiān)督學習,需要大量帶標注的數(shù)據(jù)集作為訓練樣本,而人工標注數(shù)據(jù)集不僅代價高,而且耗時長。針對這一現(xiàn)狀,提出了基于CodeBERT的代碼預訓練模型,首先利用無監(jiān)督學習技術針對大量無標簽的源代碼學習其語法、語義、上下文語境等特征信息,其次構建分類器,在少量有監(jiān)督數(shù)據(jù)集上進行微調(diào),實現(xiàn)代碼分類任務。

1 相關工作(Related work)

代碼分類是通過一定的標準、規(guī)則將實現(xiàn)不同功能、語義不相似的代碼進行區(qū)分,最終將語義相似、功能相同的源代碼歸為同一類別的過程。代碼分類首先要從代碼不同的抽象層次上提取代碼文本中蘊含的特征信息,其次根據(jù)分類器映射到不同的類別中,不同的表征方式會影響特征提取的有效性。對于實現(xiàn)不同功能的源代碼進行相似度檢測的前提,是將源代碼按照一定的粒度級別切分為基本代碼單元,即檢測粒度。目前,代碼分類研究根據(jù)粒度劃分級別主要采用固定切分和自由切分兩種方式。固定切分主要包括按行切分、按塊切分、按函數(shù)切分、按文件切分和按項目切分等,對于固定切分,如果劃分的粒度級別太大,則源代碼切分后的丟失檢查概率會很高,如果劃分的粒度級別太小,則檢測工作會耗費更多的時間與精力。因此,在具體的應用中,可以采用多種粒度對源代碼進行切分操作,例如按照Token切分、按字符串切分、按特征向量切分、按樹結(jié)點切分和按子圖結(jié)構切分等,不同源代碼的表征粒度直接影響相似性計算的精度。

早期的代碼分類是借鑒文本分類的方法,20世紀50年代至60年代,文本分類領域的資深學者通過積累的工作經(jīng)驗定義一些規(guī)則,從而實現(xiàn)分類,這種方法不僅耗時費力,還要求研究人員必須對某一領域有足夠深入的了解,才能總結(jié)分類規(guī)則并進行分類篩選,因此文本分類方法的局限性較大。針對該問題,基于統(tǒng)計學中字詞頻率的思想,提出了支持向量機(Support Vector Machine,SVM)技術,推動了文本分類的進一步發(fā)展。與此同時,深度學習技術在自然語言領域的成功應用給代碼分類帶來了新的突破,MOU等[2]通過一個單層神經(jīng)網(wǎng)絡學習抽象語法樹(Abstract Syntax Tree,AST)中上層結(jié)點的詞向量,提出一種基于樹的卷積神經(jīng)網(wǎng)絡模型(Tree-Based Convolutional Neural Network,TBCNN)進行代碼分類。BEN等[3]針對代碼本身具有的一些結(jié)構化特性,提出了連續(xù)空間的語句向量(Inst2vec)模型,該模型在代碼分類和性能預測等多種任務中獲得了很好的效果。ZHANG等[4]提出了語句級別的向量嵌入,有效捕獲了代碼的語法和語義信息,學習的代碼向量被用于代碼克隆檢測和代碼分類,取得了良好的分類效果。預訓練模型能從大量無標簽數(shù)據(jù)中學習代碼的特征信息,因此被廣泛應用于代碼表征任務[5]。例如,自監(jiān)督預訓練模型(InferCode)[6]和基于BERT的預訓練模型(CodeBERT)[7]。本文將CodeBERT預訓練模型引入代碼分類任務,一方面利用預訓練模型強大的特征學習能力提高代碼分類的性能,另一方面避免進行昂貴的人工標記,降低了實驗成本。

2 代碼分類模型(Code classification model)

代碼分類模型主要由源碼預處理、源代碼文本表征以及構建分類器三個部分組成,數(shù)據(jù)預處理主要是克服源代碼文本的篇幅長度、空白字符等消極因素對模型的影響。源代碼文本表征指利用CodeBERT模型將源代碼文本中的每一個詞映射為詞向量,所有詞向量經(jīng)過拼接處理、特征提取后,可以得到整篇文本的向量化表示,文本表征的輸出向量是分類器的輸入,將源代碼的表征向量送入Softmax層進行測試并評估結(jié)果,實現(xiàn)代碼分類。

2.1 源碼預處理

首先將獲取的源代碼數(shù)據(jù)集進行預處理,使數(shù)據(jù)集盡可能地符合預訓練模型的輸入要求,其次對處理后的源代碼數(shù)據(jù)集進行向量化操作,使得分類的源代碼文本向量可以被模型訓練學習。目前,常見的文本預處理方法有清除文本雜質(zhì)、去停用詞、類別匹配和文本過濾等[8]。清除中文文本雜質(zhì)的具體流程一般包括消除難以識別的特殊符號、刪除多余的空白字符、將繁體漢字轉(zhuǎn)換為簡體漢字三個過程。由于本文數(shù)據(jù)集是源代碼文本,因此本文中去除源代碼文本雜質(zhì)的流程包括清除計算機難以辨識的特殊符號和刪除多余的空白字符兩項操作,使源代碼文本的特征表示只關注源代碼中每一個token的語義和特征信息,從而提取源代碼文本中有價值的特征信息,以獲得更加有效的代碼表征。由于源代碼文本中通常存在輸入輸出、語句分隔符、括號等反復出現(xiàn)且無重要信息的詞,所以本文將這類詞加入停用詞詞庫進行篩選,避免無特殊含義的詞作為輸入詞。由于在分類任務中需要少量有監(jiān)督的樣本,因此需要預先知道源代碼實現(xiàn)的功能并對其進行標注處理,將實現(xiàn)相同功能的源代碼文本整理在同一個文件夾中。為防止源代碼文本篇幅長度超過本文模型的輸入范圍和樣本數(shù)據(jù)分布不一致問題,本文將對源代碼文本的篇幅長度進行篩選并對所有類別的文本進行整體過濾。

2.2 源代碼文本表征

傳統(tǒng)研究方法通過詞向量(Word2vec)模型[9]得到源代碼文本的詞向量,經(jīng)過此模型訓練得到的不同詞向量的大小固定,而一篇完整的文本可能由許多不同的字詞組成,真正有價值的字詞往往只占很小的一部分。此外,由于Word2vec模型訓練得到的詞向量相同,而相同詞向量對應的字詞在不同的語境中表達的含義不同,因此為符合現(xiàn)實的語義要求,使用CodeBERT模型替代Word2vec模型訓練文本向量。

在預訓練階段,CodeBERT模型的輸入通常設置為自然語言片段與程序語言片段組成的一個句子對,它們之間用一個特殊分隔符[SEP]進行分隔,具體的輸入形式為[CLS],w1,w2,…,wn,[SEP],c1,c2,…,cm,[EOS],其中第一個片段w1,w2,…,wn是自然語言文本,n為文本長度,第二個片段c1,c2,…,cm則是程序語言文本,m為程序長度。[CLS]、[SEP]和[EOS]是具有特殊作用的標志,[CLS]固定位于輸入語句對的開頭,[SEP]標志僅用于分隔輸入語句對中不同模態(tài)的語言,在后面的訓練過程中一般無特殊作用,[EOS]則固定位于輸入語句對的結(jié)尾,它本身不具有任何特殊的含義。

CodeBERT的輸出包括每個token的上下文向量表示和[CLS]的表示。[CLS]作為聚合序列表示蘊含整篇文本的豐富語義信息,在下游任務中一般使用[CLS]表征向量,對其參數(shù)進行微調(diào)可以實現(xiàn)不同的功能應用。由于實驗數(shù)據(jù)集使用的是C++源碼,其本質(zhì)不含有任何自然語言的單模態(tài)數(shù)據(jù),所以輸入的具體形式為[CLS],c1,c2,…,cm,[EOS]。通過模型訓練可以得到源代碼中每一個token的向量化表示和蘊含整篇文本信息的[CLS]向量,在后期的分類任務中,利用[CLS]向量初始化分類器。

2.3 構建分類器

2.3.1 Softmax回歸模型

本文使用Softmax回歸模型[10]實現(xiàn)代碼功能的多分類任務。Softmax回歸模型是邏輯回歸模型在多分類問題中的延伸與擴展,一般采用極大似然估計的方法進行參數(shù)估計,主要是對已知標簽類型的樣本進行學習訓練從而優(yōu)化模型,因此經(jīng)過預處理的數(shù)據(jù)非常適合有監(jiān)督學習任務。例如,{(x1,y1),(x2,y2),…,(xN,yN)},xi(i=1,2,…,N)表示第i篇源代碼文本對應的向量表征,N是訓練集中源代碼文本的數(shù)量,yi∈{1,2,…,K},K是分類數(shù)量,yi表示第i篇源代碼文本對應的標簽,即該篇源代碼實現(xiàn)的功能,由于本文研究的代碼分類屬于多分類問題,所以K>2。

在測試集中,輸入樣本xi,Softmax回歸模型會根據(jù)分布函數(shù)公式計算條件概率,即給定測試集中樣本xi屬于某一個類別的概率,其中出現(xiàn)概率最大的類別即當前樣本xi所屬的代碼功能類別。

因此,最終分布函數(shù)會輸出一個K維向量,每一維度的數(shù)值表示當前樣本屬于已知類別中某一類別的概率,并且模型將K維向量進行求和運算,然后做歸一化處理,使得所有向量元素和為1,Softmax回歸模型的判別函數(shù)hθ(xi)如公式(1)所示:

(1)

判別函數(shù)中的概率公式p(yi=j|xi;θ)(j=1,2,…,K)是當前樣本屬于目標類別j的概率,根據(jù)網(wǎng)絡的前向運算獲得K個輸出概率值,選擇具有最大輸出值的類別yk作為本文研究的預測值,θ為網(wǎng)絡權值參數(shù)。

2.3.2 代碼分類模型構建

本文提出基于CodeBERT的代碼分類模型,其輸入為訓練集U的源代碼文本U={(x1,y1),(x2,y2),…,(xN,yN)},(xi,yi)∈U(i=1,2,…,N),xi為第i篇源代碼文本,yi為第i篇源代碼文本對應的標簽,即該篇源代碼實現(xiàn)的功能,輸出是源代碼分類模型,具體步驟如下。

(1)訓練集U經(jīng)過預處理后得到訓練集U′,源代碼文本預處理首先刪除文本空白字符,降低輸入樣本的特征維度,其次去除停用詞,消除沒有價值的字詞對模型訓練的消極影響,最后進行長度切分處理,使輸入源代碼文本的篇幅長度符合模型的輸入要求。

(2)使用CodeBERT預訓練模型在訓練集U上進行微調(diào),通過CodeBERT模型處理之后輸出得到訓練集U′對應的特征表示V={v1,v2,…,vN},其中vi∈V(i=1,2,…,N)是每條源代碼文本xi對應句子級別的特征向量。在具體的實驗過程中,輸入的源代碼文本經(jīng)過CodeBERT模型訓練得到的[CLS]向量包含該篇源代碼的綜合語義信息,提取了每一語句中重要字詞的特征信息。

(3)將每一篇源代碼對應的特征表示V輸入Softmax回歸模型進行分類處理,Softmax回歸模型會根據(jù)分布函數(shù)的條件概率公式計算出該篇源代碼最有可能對應的功能類別。

(4)源代碼文本分類模型輸出每一篇源代碼的標簽。

3 實驗與分析(Experiment and analysis)

本文主要進行了兩個實驗,即基于CodeBERT和基于Text-CNN的代碼分類,本節(jié)主要介紹實驗中使用的實驗數(shù)據(jù)、實驗過程及實驗結(jié)果。基于CodeBERT的代碼分類如圖1所示。

圖1 基于CodeBERT的代碼分類Fig.1 Code classification based on CodeBERT

3.1 數(shù)據(jù)集

本文收集了江蘇師范大學教學科研輔助平臺中學生提交的C++源代碼課程作業(yè)作為數(shù)據(jù)集,其中包括35個種類,共905篇源代碼,每個文件夾種類都代表一種功能的實現(xiàn),包括但不限于進制轉(zhuǎn)換、歐幾里得算法、求逆序數(shù)、迪杰斯特拉算法及判斷閏年。雖然每個種類的文件夾內(nèi)又包含了數(shù)量不等的源代碼文本,但是不同種類的源代碼文本的數(shù)量符合相同的數(shù)學分布,即具有大致相同的均值與方差,所以將這些數(shù)據(jù)進行訓練學習可以避免實驗誤差。本文設置的兩組實驗均采用相同的數(shù)據(jù)集,模型評價指標主要為精確率(Precision)、準確率(Accuracy)、召回率(Recall)和F1值。

3.2 實驗

3.2.1 基于CodeBERT的實驗

本文實驗主要使用CodeBERT預訓練模型對預處理后帶標簽的源代碼數(shù)據(jù)進行訓練與驗證,最終得到每一篇源代碼的特征向量,該特征向量不僅含有源代碼的語法與語義信息,還包含豐富的上下文語境信息,最終將每一篇源代碼的特征向量送入全連接層進行拼接,連接Softmax層實現(xiàn)分類。本文實驗將預處理后生成的帶有標簽的數(shù)據(jù)集劃分為訓練集、驗證集,便于后續(xù)CodeBERT模型讀取和處理,數(shù)據(jù)文件的每一行都表示為源代碼標簽。訓練集和驗證集均經(jīng)過充分的打亂,服從同樣的數(shù)據(jù)分布,便于模型訓練的調(diào)參和結(jié)果評估。

由于本文實驗設備存在內(nèi)存限制與運行效率低的局限性,因此對本文實驗的數(shù)據(jù)進行劃分,按批次輸送到模型中。本文模型在驗證集上取得的精確率最高為0.983,F1值最高為0.96,具體的實驗結(jié)果如表1所示。

表1 CodeBERT模型實驗結(jié)果Tab.1 Experimental results based on CodeBERT

3.2.2 基于Text-CNN的實驗

本文實驗對數(shù)據(jù)集進行預處理,匹配并標記源碼,按8∶1∶1的比例將其劃分為訓練集、測試集和驗證集。由于Text-CNN模型數(shù)據(jù)最終的輸入形式為源碼的路徑信息,因此需要將樣本數(shù)據(jù)集轉(zhuǎn)化為文本文件,該文本文件的具體內(nèi)容為組合后源代碼的不同路徑信息及它們的標簽。本文實驗同樣劃分數(shù)據(jù),將其按批次輸送到模型中計算余弦相似度衡量代碼語義相似度。

本文模型經(jīng)過反復訓練,相似度閾值范圍取[0,1],F1值的最好結(jié)果為0.846(如圖2所示),其對應的精確率為0.853,召回率為0.840。

圖2 基于Text-CNN實驗結(jié)果Fig.2 Experimental results based on Text-CNN model

3.3 實驗結(jié)果與分析

本文的兩個實驗使用了相同的C++源碼數(shù)據(jù)集,并采用了兩種不同的方式從本質(zhì)上實現(xiàn)了代碼分類。因此,本文實驗對比分析具有可信度且從對比結(jié)果來看,本文實驗為代碼分類研究提供了新方法、新工具?；贑odeBERT的模型訓練和基于Text-CNN的模型訓練對比結(jié)果如表2所示。

表2 實驗對比結(jié)果Tab.2 Comparison of experimental results

實驗結(jié)果表明,基于CodeBERT預訓練模型的方法可以更全面、準確地反映代碼之間的結(jié)構關系、語義關系,具有更好的性能指標。但是,該實驗存在模型整體結(jié)構單一、在識別粒度更小的結(jié)構上存在不足等問題。對于相似度計算實驗,后期將進一步研究關鍵詞、語義、語法、結(jié)構等方面對代碼分類結(jié)果的影響,設計更全面的檢測方法。

4 結(jié)論(Conclusion)

在數(shù)據(jù)量日益激增的信息化時代,如何高效地處理與提取文本數(shù)據(jù)的有效信息具有非常重要的研究價值。本文在解決源代碼的分類問題時,使用CodeBERT預訓練模型替代傳統(tǒng)的Word2vec模型作詞嵌入操作,利用學習得到的文本向量實現(xiàn)分類。實驗通過更新參數(shù)為模型挑選適合的優(yōu)化器,其性能比一般基于文本的詞向量方法優(yōu)越。

通過實驗對比發(fā)現(xiàn),本文實驗中源代碼的表征方法可以充分利用文本的結(jié)構與上下文信息,但對代碼文本的局部特征信息進行學習時,其效果不如基于CNN模型表征學習得到的局部信息豐富。因此,在充分抽取文本局部特征信息的同時,如何利用源代碼文本的結(jié)構和上下文信息提高分類效率,是未來需要進一步研究的內(nèi)容。