融合多結構信息的代碼注釋生成模型*

余天賜,高 尚

(江蘇科技大學計算機學院,江蘇 鎮江 212100)

1 引言

代碼注釋用于描述程序中關鍵代碼的功能、注意事項和約束條件等,幫助開發人員更好地理解代碼的實現意圖和工作原理,在軟件開發和維護過程中起到重要作用。現如今,隨著軟件規模不斷增大,很多軟件項目存在著代碼注釋缺失、描述不清晰等情況,導致開發人員需要花費更多的時間和精力去閱讀和理解代碼,從而降低了開發和維護效率。代碼注釋自動生成通過機器學習的方法構建編程語言到自然語言注釋的映射模型,訓練好的模型可以分析代碼的語法結構、標識符名稱、函數參數和返回值等元素,自動生成高質量的代碼注釋,不僅可以提高開發人員的編碼效率,還提高了代碼的可讀性、可維護性和可重用性。

代碼注釋生成研究中存在2個主要問題,一是代碼的表示,二是注釋生成。不同的源代碼表示將直接影響抽取源代碼信息的質量,從而影響生成注釋的質量。現有的研究方法主要用抽象語法樹AST(Abstract Syntax Tree)一種結構信息表示代碼,沒有考慮代碼其他類型的結構信息,例如數據流和控制流,使得模型生成的注釋存在準確率低和可讀性差的問題。

為了提升模型生成注釋的質量,本文提出一種融合多結構信息的代碼注釋生成模型,模型先將AST處理為特定格式的序列,使用Transformer編碼器對AST序列進行編碼,捕獲全局信息。使用圖神經網絡GNN(Graph Neural Network)[1]對數據流圖進行特征提取,提供變量之間的計算依賴關系等信息,然后使用跨模態注意力機制(Crossmodal Attention)融合抽象語法樹和數據流2種特征,經過Transformer解碼器生成相應的注釋。通過使用Java和Python編程語言的代碼-注釋數據集驗證了本文模型性能優于6種主流對比模型,消融實驗也驗證了融合抽象語法樹和數據流圖2種特征的有效性。

2 相關工作

近年來,研究人員設計出不同的基于深度學習的代碼注釋生成模型,這些模型采用了不同的代碼表示方法和不同的神經網絡結構,但它們普遍使用序列到序列(Sequence-to-Sequence)的框架。在代碼注釋生成模型中有2個重要環節:代碼表示和注釋生成。代碼注釋生成過程通常可以分為3個階段:數據處理階段、模型訓練階段和模型測試階段。代碼表示模型是代碼注釋自動生成中的核心問題,根據代碼的不同表現形式,可將現有的代碼表示模型分為3類。

第1類是基于詞元(token)序列的代碼表示模型。在這些模型中,代碼的變量名、方法名和標識符被看作普通文本。Haiduc等[2]采用信息檢索技術對源代碼進行分析處理,從中找出合適的詞語,生成由關鍵詞組成的代碼注釋。Zheng等[3]提出的模型利用注意力機制配合編碼器獲取代碼語義信息,模型識別代碼中的循環、選擇等重要結構,并將其中的關鍵字嵌入表示為詞向量,作為編碼器的輸入。Hu等[4]利用應用程序編程接口API(Application Programming Interface)序列輔助模型生成代碼注釋。

第2類是基于抽象語法樹的代碼表示模型。Alon等[5]提出的模型將代碼表示為AST的多條路徑集,利用LSTM對每個路徑序列進行編碼,在注意力機制的配合下解碼生成代碼注釋。Zhang等[6]將AST拆分成一系列小型語句樹,并通過捕獲代碼語句的詞匯和語法知識將語句樹編碼為向量。該方法解決了AST結構復雜的問題,提高了計算性能。Hu等[7]利用2種編碼器分別編碼代碼序列和AST序列,在注意力機制的幫助下融合2種特征,最后使用解碼器生成代碼注釋。

第3類是基于圖的代碼表示模型。這類模型利用圖結構表征程序的AST、數據流和控制流等信息。戎珂瑤等[8]利用異質圖結構和圖神經網絡,將源代碼的抽象語法樹、控制流圖和數據流圖等融合并構建為具有多種節點和邊的異質表示圖,提出了基于多維度異質圖結構的代碼注釋自動生成模型。LeClair等[9]使用圖卷積神經網絡處理圖結構表征的AST,并將基于圖卷積神經網絡的編碼器輸出與基于代碼token序列編碼器的輸出相結合。模型可以學習到代碼結構與序列間的關系。Allamanis等[10]設計了一種基于門控圖神經網絡的代碼表示模型。模型使用圖表示代碼的AST和數據流信息。訓練好的模型可以為給定代碼段預測變量名。與前2類代碼表示模型相比,基于圖的代碼表示模型更適合表示源代碼的結構信息。

本文利用抽象語法樹和數據流圖可以表示代碼中不同類型的結構信息,提出一種融合多結構信息的代碼注釋生成模型。

3 本文模型

3.1 問題定義

代碼注釋生成任務可以形式化為有監督機器學習問題。本文使用C={c1,…,ci}表示代碼片段,S={s1,…,si}表示與代碼對應的注釋句子,構建由(C,S)對組成的訓練集,通過對模型訓練,得到注釋生成模型f,當輸入代碼片段ci時,模型f可以自動生成注釋si。

3.2 模型結構

本文提出的融合多結構信息的代碼注釋生成模型結構如圖1所示,模型主要包含以下3部分:抽象語法樹編碼器、數據流圖編碼器和解碼器。抽象語法樹編碼器將原始的AST序列轉換為一個固定維度的向量,以捕捉AST中的語義和結構信息。數據流圖編碼器使用圖神經網絡將代碼數據流圖轉換成固定維度的向量,以捕捉代碼數據流中的變量依賴關系等關鍵信息。解碼器接收跨模態注意力機制生成的中間向量表示,然后將其解碼為代碼注釋序列。其中,抽象語法樹編碼器由N個相同的編碼器層堆疊而成,解碼器由N個相同解碼器層堆疊而成。

Figure 1 Structure of the proposed model

3.3 編碼器

3.3.1 抽象語法樹編碼器

抽象語法樹AST是由節點和邊組成的樹狀結構,每個節點表示代碼中的一個語法結構。AST的非葉節點主要表示代碼中的各種語句,例如循環、判斷語句以及聲明語句。而葉子節點表示標識符等,例如變量名、參數名和函數名。本文使用Transformer[11]的編碼器來提取代碼抽象語法樹的上下文信息。為了使AST能夠作為編碼器的輸入,需要將AST轉換為特定格式的序列,本文參考文獻[12]的方法,處理過程如下:從根節點開始,首先使用一對括號來表示樹結構,并將根節點本身放在右括號后面。然后,遍歷根節點的子樹,并將子樹的所有根節點放入括號中。如此遞歸遍歷,直到遍歷完所有節點,得到最終序列。

圖2表示一段Java代碼,其對應的抽象語法樹如圖3所示。

Figure 2 Java code

Figure 3 Abstract syntax tree of Java code in figure 2

將如圖3抽象語法樹轉換成序列可以表示為:(MethodDeclaration(int)int(minNum)minNum(parameters)parameters(BlockStmt(…))BlockStmt)MethodDeclaration。使用這種格式的序列,可以使模型從AST中學習到代碼的語義和語法信息。

抽象語法樹編碼器使用X={x1,…,xl}表示某一代碼片段的AST序列,l表示序列長度,X∈Rl×dAST,dAST表示嵌入的維度。獲得AST序列嵌入表示后,利用位置編碼捕獲序列的位置信息。

PE(i,2m)=sin(i/100002m/dAST)

(1)

PE(i,2m+1)=cos(i/100002m/dAST)

(2)

其中,i表示序列X中token的位置,m是位置編碼的維度索引,2m表示偶數維度,2m+1表示奇數維度,當位置編碼維度為偶數時使用式(1)計算,當位置編碼維度為奇數時使用式(2)計算。

將序列嵌入向量和位置編碼相加后,可以得到AST序列的輸入特征Xinput,如式(3)所示:

Xinput=X+PE

(3)

(4)

A=[a1,a2,…,ah]

(5)

Z=LayerNorm(Xinput+AWO)

(6)

XAST=LayerNorm(Z+FFN(Z))

(7)

其中,dr表示矩陣Qj和Kj的維度,h表示注意力頭的個數。LayerNorm(·)表示歸一化操作,WO表示權重矩陣,FFN(·)表示前饋神經網絡。抽象語法樹編碼器由N個相同編碼器層堆疊而成。

3.3.2 數據流圖編碼器

數據流圖DFG(Data Flow Graph)以圖的形式記錄了代碼段中變量的賦值和傳播過程。其中,圖的節點表示變量,邊表示變量的來源。這種圖結構為代碼理解提供了關鍵的代碼語義信息。圖4表示圖2中Java代碼的數據流圖。

Figure 4 Data flow graph of Java code in figure 2

數據流圖編碼器使用圖神經網絡GNN[1]提取數據流圖特征,GNN通過聚合鄰居節點的信息來更新節點的特征表示,最終得到數據流圖的嵌入表示。數據流圖是一種有向圖,為了使圖神經網絡在處理有向圖時不丟失方向信息,需要同時從節點的傳入和傳出方向進行消息傳遞。數據流圖編碼器具體計算過程如下:使用G=(V,E)表示數據流圖,V表示變量節點的集合,E表示邊的集合。圖中節點的連接關系使用鄰接矩陣A∈RN×N表示,N表示節點個數。vi∈V表示節點,(vi,vj)∈E表示從節點vi指向vj的一條有向邊,1≤i,j≤N。用N├(v)={u∈V|(v,u)∈E}表示節點v所有前向鄰居的集合,即v指向的節點。用N┤(v)={u∈V|(u,v)∈E}表示節點v所有后向鄰居的集合,即指向v的節點。使用xv表示節點v經過嵌入后的文本特征向量,矩陣Xnode∈RN×dnode存儲所有節點的文本特征向量,其中dnode表示特征向量的維度。圖4中變量x1的一階前向鄰居為x5和x8,二階前向鄰居為z7,z7的一階后向鄰居為x8。圖神經網絡主要通過聚合節點K階鄰居的信息來更新節點特征。圖神經網絡主要計算過程如下:對所有節點,先初始化節點的前向表示和后向表示為其文本特征向量,如式(8)所示:

(8)

(9)

其中,AGGREGATE(·)表示聚合操作。

(10)

其中,σ(·)表示非線性激活函數,Wk表示權重矩陣。

同理,將相同的過程應用于后向表示。

(11)

(12)

本文使用均值聚合方法。

(13)

經過K次聚合過程,拼接節點最終的前向表示和后向表示得到向量Zv,作為節點v的特征向量。

(14)

最后將節點的特征向量傳入一個前饋神經網絡,并使用最大池化方法得到數據流圖的特征向量XDFG,如式(15)所示:

XDFG=max({σ(WZv+b),v∈V})

(15)

其中,W表示權重矩陣,b表示偏置向量。

3.4 特征融合

本文使用一個跨模態注意力模塊Crossmodal attention[13]來融合抽象語法樹特征和數據流圖特征,這樣可以使抽象語法樹編碼器接收數據流圖編碼器的信息。

本文使用XAST∈RTα×dα表示抽象語法樹特征,XDFG∈RTβ×dβ表示數據流圖特征,其中,T(·)表示特征序列長度,d(·)表示特征的維度。分別定義查詢矩陣Qα,索引矩陣Kβ和值矩陣Vβ。經過下列計算可得融合后的特征ZFusion。

Qα=XASTWQα

(16)

Kβ=XDFGWKβ

(17)

Vβ=XDFGWVβ

(18)

(19)

其中,WQα∈Rdα×ds,WKβ∈Rdβ×ds,WVβ∈Rdβ×dv為權重矩陣。其中,ds和dv分別表示權重矩陣對應的維度。

3.5 解碼器

本文使用Transformer的解碼器來生成代碼注釋。解碼器部分由N個解碼器層堆疊而成,每個解碼器層由3個子層連接組成,第1個子層包括帶掩碼的多頭自注意力機制、殘差連接以及規范化層。使用帶掩碼的多頭自注意力機制可以確保預測僅依賴于已生成的輸出詞元。第2個子層包括連接融合特征和解碼器的多頭注意力機制、殘差連接以及規范化層。其中的多頭注意力機制與之前抽象語法樹編碼器中多頭注意力機制不同,查詢來自前一個解碼器層的輸出,而鍵和值來自于融合特征,這樣能夠使解碼器捕獲到抽象語法樹序列和數據流圖特征信息。第3個子層包括前饋神經網絡、殘差連接以及規范化層。經過N個解碼器層的計算后,使用線性層得到指定維度的向量。最后經過softmax層得到概率向量,模型選擇具有最高概率的單詞作為輸出。

本文使用交叉熵損失函數來訓練模型,損失函數如式(20)所示:

(20)

其中,n表示整個輸出序列的長度,yt表示模型預測第t時間步的token。訓練過程中,通過最小化損失函數得到模型參數。

4 實驗與結果分析

4.1 數據集

為了有效評估本文模型,實驗數據集采用CodeXGLUE(General Language Understanding Evaluation benchmark for Code)中Java和Python語言代碼-注釋數據集[14]。Java語言數據集包含181 061對代碼-注釋對。其中,164 923對數據用于訓練模型,10 955對數據用于測試模型,5 183對數據作為驗證集。Python語言數據集包含280 652對代碼-注釋對。其中,251 820對數據用于訓練模型,14 918對數據用于測試模型,13 914對數據作為驗證集。

4.2 數據預處理

在模型訓練之前,需要對代碼和注釋數據進行預處理,將其轉換為可以作為模型輸入的數據結構。

本文使用開源工具tree-sitter為Java和Python代碼構建抽象語法樹。tree-sitter是一個生成解析器的工具。它可以為多種語言的代碼構建語法樹。生成抽象語法樹的主要步驟如下:首先,對代碼進行詞法分析,將其分割成多個詞法單元,例如關鍵字、標識符和操作符等。然后進入語法分析階段,語法分析器按照不同編程語言的語法規則,對詞法單元進行分析和組織,最后構建出一棵抽象語法樹。為了方便模型訓練,在構建抽象語法樹的過程中,需要刪除節點數超過1 000以及因語法錯誤導致構建抽象語法樹失敗的數據集,最后將處理好的數據存儲在JSONL格式文件中。

數據流圖的構建主要參考文獻[15]中使用的方法,具體構建步驟如下:根據抽象語法樹中的葉子節點來識別變量序列,這些變量序列中的每一個元素都會成為數據流圖中的一個節點。然后根據抽象語法樹中變量之間的依賴關系,生成數據流圖的有向邊。以集合形式存儲變量節點,使用鄰接矩陣表示節點間的連接關系。

4.3 參數設置

實驗包含基本參數設置。抽象語法樹編碼器中編碼器層數N為6,多頭注意力機制頭的個數h為8,嵌入維度dAST為512。數據流圖編碼器中節點嵌入維度dnode為512,聚合節點鄰居階數K為2。解碼器中解碼器層數N為6,帶掩碼的多頭注意力機制頭的個數為8。使用Adam優化器訓練模型,初始學習率設置為0.001。

4.4 評估指標

實驗采用BLEU(BiLingual Evaluation Understudy)[16]、METEOR(Metric for Evaluation of Translation with Explicit ORdering)[17]和ROUGE(Recall-Oriented Understudy for Gisting Evaluation)[18]3種評估指標來評估模型生成代碼注釋的質量。

BLEU是一種基于n-gram精確度的評估指標,可以自動分析模型生成注釋與參考注釋之間的相似度。BLEU分數的取值范圍是0到1,越接近1表示模型生成注釋的質量越好。

METEOR在評估中加入召回率來反映模型生成注釋與參考注釋的相似程度。分數越高表示模型生成注釋越接近參考注釋。與BLEU不同的是,METEOR將模型生成注釋與參考注釋進行對齊,以便更好地評估模型生成注釋的質量。

ROUGE是一種基于召回率的文本摘要質量評估指標。它含有4種指標:ROUGE-L、ROUGE-S、ROUGE-N以及ROUGE-W。在代碼注釋質量評估中常使用ROUGE-L,它是基于最長公共子序列的召回率和精確度來計算的。

4.5 對比模型

為了驗證模型有效性,本文選用以下幾種相關代碼注釋生成模型進行對比分析。

(1)Tree2Seq[19]:使用基于樹的LSTM作為編碼器捕獲AST的結構信息。

(2)code2seq[20]:使用AST的多條隨機路徑集表示代碼片段,通過LSTM編碼AST路徑,在注意力機制的配合下解碼生成代碼注釋。

(3)Code-NN[21]:將代碼表示為token序列,然后利用LSTM和注意力機制生成代碼注釋。

(4)DeepCom[11]:對代碼AST進行結構遍歷得到可以作為模型輸入的特定格式序列,這種序列可以更好地表示代碼的結構信息。

(5)CodeGNN[8]:包含2個編碼器,一個使用GRU編碼源代碼序列,另一個使用圖神經網絡編碼代碼的AST。

(6)Transformer[22]:使用多頭注意力機制編碼輸入和解碼輸出,并利用位置編碼學習代碼序列中token的位置關系。

4.6 實驗結果

4.6.1 對比實驗

表1展示了本文模型和對比模型在CodeXGLUE中Java和Python語言代碼-注釋數據集上的實驗結果。通過表1可以看出,本文提出的融合多結構信息的代碼注釋生成模型在2個數據集上的BLEU-4、METEOR和ROUGE-L 指標均優于對比模型。

Table 1 Results of different models on Java and Python datasets

具體地,基于代碼token序列的Code-NN模型,僅考慮了代碼的語義信息,從代碼中抽取關鍵字,沒有充分利用代碼結構特征,因此模型生成的注釋質量較低。基于抽象語法樹的代碼注釋生成模型,例如Tree2Seq、code2seq和DeepCom,這類模型通常用AST的節點序列、路徑集合來表示源代碼。因為抽象語法樹可以表示代碼的抽象句法結構,所以這類模型可以學習到代碼的結構信息,生成的注釋質量也比基于代碼token序列的注釋生成模型要高。使用圖神經網絡的注釋生成模型CodeGNN利用圖結構表征代碼的AST。由于圖更適合表示代碼的結構信息,所以CodeGNN與基于AST序列的模型相比,BLEU、METEOR和ROUGE-L的得分更高。Transformer能捕獲長距離依賴關系,因此使用Transformer的注釋生成模型3種指標分數要高于Code-NN模型。本文模型利用Transformer和圖神經網絡的優點,融合抽象語法樹和數據流圖2種特征來表示代碼,能更全面地覆蓋代碼所包含的語義和結構等信息。模型生成的注釋的BLEU、METEOR和ROUGE-L分數均高于對比模型。

4.6.2 消融實驗

為了進一步分析抽象語法樹編碼器和數據流圖編碼器的作用,以及圖神經網絡中聚合節點鄰居階數和特征融合方式對模型性能的影響,本文設計了消融實驗。在Java和Python數據集上的消融實驗結果如表2所示,加粗表示最優的結果。AST only表示只有抽象語法樹編碼器,DFG only表示只有數據流圖編碼器。其中,K表示聚合節點鄰居階數,concat表示在特征融合時使用拼接操作,CM attention表示在特征融合時使用跨模態注意力機制。

Table 2 Results of ablation experiments on Java and Python datasets

由表2中結果可知,圖神經網絡聚合數據流圖節點鄰居階數會影響模型性能。在只有數據流圖編碼器的實驗中,當K=2時,模型在Java和Python數據集上BLEU-4、METEOR和ROUGE-L 3種評估指標分數最高,增加聚合節點鄰居階數會出現過平滑問題,即經過多次鄰居節點特征的聚合,所有節點的特征會變得相似,模型不易學習到節點的特征,從而導致模型性能下降。在特征融合時,不同的融合方式也會影響模型性能,從表2中數據可知,使用跨模態注意力機制的模型在Java和Python數據集上BLEU-4、METEOR和ROUGE-L 3種評估指標分數都要高于使用拼接操作融合2種特征的模型,驗證了跨模態注意力機制的有效性。另外,只使用抽象語法樹編碼器或者數據流圖編碼器會降低注釋生成模型的性能。抽象語法樹和數據流圖都是代碼的結構化表示形式,將2種特征融合可以幫助模型更全面地理解代碼的結構和行為。具體來說,AST中含有代碼的語法結構信息,而DFG可以描述代碼中的數據流和計算依賴關系,其中包含變量的定義、使用以及傳遞過程。使用跨模態注意力機制融合AST和DFG 2種特征的注釋生成模型在Java數據集上BLEU-4、METEOR和ROUGE-L 3種指標得分分別為20.36%,15.79%和36.73%,在Python數據集上3種指標得分分別為22.49%,14.96%和32.58%。

消融實驗結果表明,融合代碼的抽象語法樹和數據流圖2種特征可以提升模型生成注釋的質量。

4.6.3 案例展示

本節主要展示模型生成注釋的案例。從數據集中隨機抽取2段Java和Python代碼,其參考注釋和模型生成的注釋如表3所示。從表3中可以看出,模型生成的注釋可以準確描述代碼功能,且具有良好的可讀性。

Table 3 Cases of code summarization

5 結束語

本文提出了一種融合多結構信息的代碼注釋生成模型,模型在代碼抽象語法樹的基礎上,增加了數據流圖結構信息來表示代碼。模型使用Transformer編碼器對AST序列進行編碼,捕獲代碼全局信息。使用圖神經網絡對數據流圖進行特征提取,提供變量之間的計算依賴關系等信息,然后使用跨模態注意力機制融合抽象語法樹和數據流2種特征,最后經過Transformer解碼器生成相應的注釋。對比實驗結果表明,本文模型的性能要優于6種主流注釋生成模型。消融實驗也驗證了融合抽象語法樹和數據流2種特征的有效性。在未來工作中,一是需要設計更適合抽取代碼結構信息的模型,提高模型學習能力。二是需要探究多種方法融合的注釋生成算法。三是提升模型的泛化能力,可以將注釋生成模型應用到不同的代碼語言中。