DeepSCC: Source Code Classification Based on Fine-Tuned RoBERTa

In this section, we show the framework of DeepSCC in Figure 1. In particular, we first preprocess the corpus, including data cleaning, filtering, and word segmentation. Then we fine-tune the pre-trained model RoBERTa to predict the type of programming language.

In this phase, the data cleaning and filtering are consistent with previous work [4]. However, we find that the previous code classification methods treat the code word as the basic unit. Its disadvantage is that it cannot effectively solve the out-of-vocabulary (OOV) problem. That means there exist some words, which are not in the training set but in the testing set. To solve the OOV problem, we use the Byte-Pair Encoding (BPE) proposed by Sennrich et al. [15]. It is a mixture between character-level and word-level representations. Using BPE can avoid a large number of ”[UNK]” symbols in the test set, as ”[UNK]” symbols may decrease the performance of the pre-trained model. For example, the original code snippet is “def split_lines(s): return s.split(‘$\backslash$n’)”, and the result after using BPE segmentation is “def”, “Ġsplit”, “_”, “lines”, “(”, “s”, “)”, “:”, “Ġreturn”, “Ġs”, “.”, “split”, “(”, “’”, “Ċ”, “’”, “)”. Here Ġ means that it is the first subword of a subword division.

In this phase, we continue to pre-train the RoBERTa model on the code corpus with the MLM (mask language model) method, and then use the pre-trained model to fine-tune the code classification task. RoBERTa [6] is similar to Bert (Bidirectional Encoder Representations from Transformers)[12], and DeepSCC uses Transformer as the method’s main framework because Transformer can more thoroughly capture the bidirectional relationship in the text. In particular, we treat the code as text and use the method MLM to constantly fine-tuning the roberta-base model²²2https://huggingface.co/roberta-base on the corpus to obtain our pre-trained language model. During the fine-tuning process, we do not tune the parameters of the model’s bias and LayerNorm.weight weights, and use the AdamW method to fine-tune the other parameters.

Consider that different layers of the neural network can capture different levels of syntactic and semantic information. In our study, we choose the last layer of Encoder as the feature representation of the whole code snippet, feed it into the linear layer, and obtain the model prediction label by Softmax, which can be used to calculate the cross-entropy loss with the real label. Then we use AdamW as the optimizer to perform gradient descent and back propagation to update the model parameters. Finally, we can obtain our fine-tuned model.

We choose the corpus shared by Alreshedy et al. [7] as our experimental subject. Alreshedy et al. gathered code snippets from 21 programming languages (i.e., Bash, C, C#, C++, CSS, Haskell, HTML, Java, JavaScript, Lua, Objective-C, Perl, PHP, Python, R, Ruby, Scala, SQL, Swift, Visual Basic, and Markdown). After manual analysis on their gathered corpus, we find: (1) The number of the code snippets related to Markdown is only 1,210, which is significantly lower than that of other languages. (2) In the code snippets related to HTML, we find most of these code snippets also include CSS and JavaScript code segments. Therefore, we remove the code snippets related to these two languages. Finally, we use 179,556 code snippets for model training and 44,889 code snippets for model testing via stratified sampling.

To compare the performance between our proposed method and the baselines, we choose the following four performance measures: Accuracy, Precision, Recall, and F1. Before introducing these measures, we first illustrate the following concepts:

• •

  True Positive (TP): The positive sample is successfully predicted as positive.

• •

  True Negative (TN): The negative sample is successfully predicted as negative.

• •

  False Positive (FP): The negative sample is wrongly predicted as positive.

• •

  False Negative (FN): The positive sample is wrongly predicted as negative.

Then the four performance measures can be computed as follows:

Accuracy is the most intuitive performance measure, and it is the ratio of the correctly predicted observations to the total observations. Precision is the ratio of the correct predicted positive observations to the total predicted positive observations. Recall indicates how many positive examples in the sample are predicted correctly. F1 is the average of Precision and Recall.

In the RQ, we first compare our proposed method DeepSCC with two state-of-the-art methods from source code classification (i.e., Random Forest and XGBoost methods used in SCC++ [4]). We also choose TextCNN [8] and Transformer [11] with/without pre-trained word vectors (i.e., FastText [9] or Word2Vec [10]) from the neural text classification field as baselines. Besides, we also select BERT [12] as a baseline for the pre-trained model.

In our study, we use Pytorch 1.6.0 to implement our proposed method. For baselines in the source code classification field, we run their shared code on our preprocessed corpus. For baselines in the neural text classification field, we re-implemented these baselines according to the corresponding description by Pytorch. For BERT and RoBERTa, we pre-train the model in the transformer library.

It needs to be noticed that pre-trained models (i.e., BERT and RoBERTa) use the method BPE for code segmentation by default. For other baselines, we choose the word_tokenize method provided by the NLTK library for code segmentation.

We run all the experiments on a computer with an Inter(R) Core(TM) i7-9750H 4210 CPU and a GeForce GTX3090 GPU with 24 GB memory. The running OS platform is Windows 10.

Table 1 shows the comparison results between DeepSCC and baselines. Table 2 shows the detailed results for each language type in terms of three performance measures (Precision, Recall, and F1) and support (i.e., the number of code snippets related to the given programming language in the test set).

According to the analysis of the experimental results, we can find: (1) From Table 1, we can find that our method can outperform baselines and achieves the best performance in source code classification. Specifically, it can achieve a maximum performance improvement of 17%, 16%, 19%, and 18% in terms of Accuracy, Precision, Recall, and F1 respectively. The results show that the two-way transformer encoder can learn the deep semantics of the code snippets more effectively, which is helpful to obtain a better classification performance. (2) Not all the baselines in the neural text classification field outperform the baselines in the code classification field. That means some traditional machine learning methods can outperform deep learning-based methods in this task. (3) For baselines in the field of neural text classification, Transformer is not as effective as TextCNN in code classification. This may be because Transformer learns too little code semantics. After adding pre-trained word vectors (such as Word2Vec and FastText), the performance of TextCNN can be slightly improved, but the performance of Transformer is decreased. This shows that pre-trained word vectors can better capture the feature representation of the code when the structure is CNN in this task. (4) From Table 2, we can find that DeepSCC can achieve high performance in most of the programming languages. Then we analyze the cause of the poor performance when the programming languages are C/C++ and CSS/JavaScript. Specifically, 8% of the code snippets with the actual category of C are predicted to be C++, and 10% of the code snippets with the actual category of C++ are predicted to be C. Since C++ is almost a superset of C, this indicates that some C++ code snippets and C code snippets are indistinguishable, which poses a challenge for the source code classification problem. 6% of the code fragments with the actual category of CSS are predicted to be JavaScript, and 7% of the code snippets with the actual category of JavaScript are predicted to be CSS. Because CSS as a style language often appears in the scripting language JavaScript, it is used to dynamically update page elements. This leads to the simultaneous appearance of JavaScript and CSS in the code snippets, which also poses another challenge for the source code classification problem.