Код программ отличается от естественного языка из-за его формализма и строгости, однако ничто не мешает воспринимать его как последовательность токенов и работать с ним, как с обычным языком. Существуют исследования, которые показали, что модель BERT, обученная на большом наборе данных, неплохо справляется с некоторыми задачами, связанными с обработкой программного кода. В этом посте я буду решать задачу автогенерации комментариев к нему. Вы узнаете, как подготовить данные для обучения, настроить нейросеть и получить результат.

Данные

Данные представлены в виде набора пар [функция — комментарий] для различных языков программирования (awesome Code Search Net Challenge dataset). Кстати говоря, этот набор изначально был создан не для этой задачи, однако его можно легко перепрофилировать под свои нужды. 

Я не стану очищать данные, это описано здесь. Я же буду использовать уже предварительно обработанные данные в объеме 1 % от общего количества образцов в наборе, так как обучение модели занимает довольно много времени. Но, как можно будет убедиться в будущем, генерация комментариев даже на 1 % данных выглядит неплохо. Если у вас есть время и ресурсы, можете обучить модель на всём наборе и получить результаты получше.

CodeBERT

Предварительно обученная модель, которую я буду использовать, взята из статьи исследовательского подразделения Microsoft. В этой модели также использовался набор данных CodeSearchNet, но вместо генерирования комментариев он использовался для обучения модели на основе RoBERTa удобному для восприятия представлению кода и естественного языка. Использование больших языковых моделей для представления текста удобным способом в настоящее время является обычной практикой, поскольку они показали свою эффективность для решения других задач. 

Загрузка, установка и импортирование библиотек

!pip install transformers
!git clone -q https://github.com/microsoft/CodeXGLUE.git


import json
from dataclasses import dataclass
import numpy as np
import pandas as pd
from transformers import AutoTokenizer

Здесь я прописываю пути до файлов с данными и оборачиваю их в структуру для более удобного дальнейшего использования:

PATH_TO_TRAIN_DATA = '/content/train.csv'
PATH_TO_TEST_DATA = '/content/test.csv'
PATH_TO_VALIDATION_DATA = '/content/validation.csv'
#validation, test and train
data_struct = {
	'train' : pd.read_csv(PATH_TO_TRAIN_DATA),
	'test' : pd.read_csv(PATH_TO_TEST_DATA),
	'valid' : pd.read_csv(PATH_TO_VALIDATION_DATA)
}

Инициализирую две вспомогательные функции: токенизации текста и записи DataFrame в JSON-файл, так как именно в таком формате требуется подавать данные для модели.

}
def write_into_json_file(json_file_name: str, data: pd.DataFrame):
	'''
	json_file_name - name output json file
	data - pandas data frame
	write your pandas data to json file
	'''
with open(json_file_name, 'w') as current_file:
	for index, current_row in data.iterrows():
		current_file.write(json.dumps(current_row.to_dict()) + '\n')


def split_data(split_column: str, new_column: str, data: pd.DataFrame)-> pd.DataFrame:
	'''
	split items in column
	data - your pandas data frame
	split_column - column in your pd.df
	'''
	data[new_column] = data[split_column].apply(lambda current_item: current_item.split())
	return data

Реализую небольшую предобработку данных с помощью функций, описанных выше:

#preproc data
for type_data, value in data_struct.items():
	#split target colums
	code_tokens_step = split_data('code', 'code_tokens', value)
	docs_tokens_step = split_data('comment', 'docstring_tokens', code_tokens_step)
	data_struct[type_data] = docs_tokens_step
	#create json file
	write_into_json_file(f'/content/{type_data}.jsonl', data_struct[type_data])

Далее создаю конфигурационный класс для модели и на его основе прописываю всю конфигурацию:

@dataclass
class ConfigurationModel:
	learning_rate : float 
	batch_size : int 
	beam : int
	test_file : str
	source_size : int
	target_size : int
	path_to_data_directory : str
	path_to_output_data_directory : str
	train_file : str
	dev_file : str
	count_epochs : int
	pretrained_model : str


configuration_codetext_model = ConfigurationModel(
	learning_rate = 5e-5,
	batch_size = 8,
	beam = 10,
	source_size = 256,
	target_size = 512,
	path_to_data_directory = '.',
	path_to_output_data_directory = 'model_for_java',
	train_file = '/content/train.jsonl',
	dev_file = '/content/valid.jsonl',
	test_file = '/content/test.jsonl',
	count_epochs = 10,
	pretrained_model = 'microsoft/codebert-base',
)
configuration_codetext_model

Обучение

Теперь, когда данные обработаны и представлены в удобном формате, можно приступать к обучению. Сделаю этоОбучу модель на обучающей выборке. В качестве метрики использую BLEU-4 (четвёрка означает, что количество словесных n-gram = 4), которая распределена от 0 до 1, но в нашем примере будет использоваться BLEU-4 * 100%. Эта метрика используется в задачах машинного перевода, но и для генерации текста она также хорошо подходит. Если брать задачи машинного перевода, то даже для человека bleu = [0.6:0.7] — отличный результат, потому что каждый человек может перевести текст по-разному. Точности в единицу достигнуть почти невозможно. 

Если посмотреть на исходную задачу, то, во-первых, модель должна сгенерировать текст, а во-вторых, это не просто текст, а осмысленный комментарий к коду. Поэтому ожидать больших значений метрики bleu не стоит. 

#run train model
!python /content/CodeXGLUE/Code-Text/code-to-text/code/run.py \
	--do_train \
	--do_eval \
	--do_lower_case \
	--model_type roberta \
	--model_name_or_path {configuration_codetext_model.pretrained_model} \
	--train_filename {configuration_codetext_model.train_file} \
	--dev_filename {configuration_codetext_model.dev_file} \
	--output_dir {configuration_codetext_model.path_to_output_data_directory} \
	--max_source_length {configuration_codetext_model.source_size} \
	--max_target_length {configuration_codetext_model.target_size} \
	--beam_size {configuration_codetext_model.beam} \
	--train_batch_size {configuration_codetext_model.batch_size} \
	--eval_batch_size {configuration_codetext_model.batch_size} \
	--learning_rate {configuration_codetext_model.learning_rate} \
	--num_train_epochs {configuration_codetext_model.count_epochs}

Обучение

После обучения модели её можно проверить на отдельной выборке:

binary_model_file = '/content/model_for_java/checkpoint-best-bleu/pytorch_model.bin'
!python  /content/CodeXGLUE/Code-Text/code-to-text/code/run.py \
    --do_test \
    --model_type roberta \
    --model_name_or_path microsoft/codebert-base \
    --load_model_path {binary_model_file} \
    --dev_filename {configuration_codetext_model.dev_file} \
    --test_filename {configuration_codetext_model.test_file} \
    --output_dir {configuration_codetext_model.path_to_output_data_directory} \
    --max_source_length {configuration_codetext_model.source_size} \
    --max_target_length {configuration_codetext_model.target_size} \
    --beam_size {configuration_codetext_model.beam} \
    --eval_batch_size {configuration_codetext_model.batch_size}

Как можно увидеть, bleu-4 = 11, и это неплохо для такой задачи, даже с учётом того, что bleu в нашем случае распределена от 0 до 100.

Далее считаю получившиеся результаты:

path_to_gold = '/content/model_for_java/test_1.gold'
path_to_output = '/content/model_for_java/test_1.output'

Инициализирую функцию считывания из txt-файла:

def read_result_txt_file(txt_file: str)-> list:
	with open(txt_file) as file:		return [' '.join(line.rstrip().replace('\t', ' ').split(' ')[1:]) for line in file]

И для удобства считаю всё в DataFrame:

def read_result_txt_file(txt_file: str)-> list:
#true comments and predicted
true_sent = read_result_txt_file(path_to_gold)
pred_sent = read_result_txt_file(path_to_output)
result_data_frame = pd.DataFrame(
	{
		'code' : data_struct['test']['code'],
		'true' : true_sent,
		'pred' : pred_sent
	}
)


result_data_frame.head(10)

Вывод 10 примеров кода, оригинальных комментариев и комментариев, сгенерированных моделью.

Теперь попробую субъективно сравнить оригинальный комментарий со сгенерированным по шкале от 1 до 5. Code — исходный код, true — исходный комментарий, pred — сгенерированный.

Пример 1:

Code: public t includeas(final class template) { blacklist = false; string[] properties = getallbeanpropertynames(template, false); include(properties); return _this(); }

True: defines included property names as public properties of given template class. sets to black list mode.

Pred: create a new resource

Оценка: 1 — абсолютно непонятно, о чём идёт речь.

Пример 2:

Code: int setdirect(int v0, int v1, int v2, int v3, int v4) { return offset + v0*stride0 + v1*stride1 + v2*stride2 + v3*stride3 + v4*stride4; }

True: experimental : should be package private

Pred: sets the value for the specified point.

Оценка: 4 — исходный комментарий абсолютно никак не отражает функциональность, в отличие от сгенерированного. 

Пример 3:

Code: static private servicetype checkifdap4(string location) throws ioexception { // strip off any trailing dap4 prefix if (location.endswith(".dap")) location = location.substring(0, location.length() - ".dap".length()); else if (location.endswith(".dmr")) location = location.substring(0, location.length() - ".dmr".length()); else if (location.endswith(".dmr.xml")) location = location.substring(0, location.length() - ".dmr.xml".length()); else if (location.endswith(".dsr")) location = location.substring(0, location.length() - ".dsr".length()); try (httpmethod method = httpfactory.get(location + ".dmr.xml")) { int status = method.execute(); if (status == 200) { header h = method.getresponseheader("content-type"); if ((h != null) && (h.getvalue() != null)) { string v = h.getvalue(); if (v.startswith("application/vnd.opendap.org")) return servicetype.dap4; } } if (status == httpstatus.sc_unauthorized || status == httpstatus.sc_forbidden) throw new ioexception("unauthorized to open dataset " + location); // not dods return null; } }

True: check for dmr

Pred: returns true if the given name is valid.

Оценка: 5 — попадание в точку.

Пример 4:

Code: public void setsize(dimension newsize) { if ( newsize != null ) { size.setsize( newsize ); firepropertychange( size_prop, null, size ); } }

True: set the size of this vertex. will not update the size if newsize is null.

Pred: sets the initializes the size of the shape.

Оценка: 3 — в целом, комментарии схожи, но вместо vertex используется shape, и в сгенерированном комментарии не отражено условие, которое прописано в оригинальном.

Пример 5:

Code: protected bufferedimage createbufferedimage(int w, int h, int imgtype) { bufferedimage bi = null; if (imgtype == 0) { bi = (bufferedimage) createimage(w, h); } else if ((imgtype > 0) && (imgtype < 14)) { bi = new bufferedimage(w, h, imgtype); } else if (imgtype == 14) { bi = createbinaryimage(w, h, 2); } else if (imgtype == 15) { bi = createbinaryimage(w, h, 4); } else if (imgtype == 16) { bi = createsgisurface(w, h, 32); } else if (imgtype == 17) { bi = createsgisurface(w, h, 16); } // store the buffered image size biw = w; bih = h; return bi; }

True: generates a fresh buffered image of the appropriate type.

Pred: creates a new image.

Оценка: 2 — в исходном комментарии сказано, что генерируется новое буферное изображение определённого типа, в сгенерированном такие уточнения отсутствуют.

Пример 6:

Code: public orientgraph gettx() { final orientgraph g; if (pool == null) { g = (orientgraph) gettxgraphimplfactory().getgraph(getdatabase(), user, password, settings); } else { // use the pool g = (orientgraph) gettxgraphimplfactory().getgraph(pool, settings); } initgraph(g); return g; }

True: gets transactional graph with the database from pool if pool is configured. otherwise creates a graph with new db instance. the graph instance inherits the factory's configuration.

Pred: get the graph for the graph.

Оценка: 1 — очень краткое и в то же время неверное описание.

Пример 7:

Code: public boundingbox getboundingbox(long geopackageid, string tablename) { boundingbox boundingbox = null; cursor result = db.rawquery("select min(" + geometrymetadata.column_min_x + "), min(" + geometrymetadata.column_min_y + "), max(" + geometrymetadata.column_max_x + "), max(" + geometrymetadata.column_max_y + ") from " + geometrymetadata.table_name + " where " + geometrymetadata.column_geopackage_id + " = ? and " + geometrymetadata.column_table_name + " = ?", new string[]{string.valueof(geopackageid), tablename}); try { if (result.movetonext()) { boundingbox = new boundingbox(result.getdouble(0), result.getdouble(1), result.getdouble(2), result.getdouble(3)); } } finally { result.close(); } return boundingbox; }

True: query for the bounds of the feature table index

Pred: get the bounding box.

Оценка: 3 — в целом, суть похожа.

Пример 8:

Code: public static list<element> getelements(stage stage, iterable<? extends module> modules) { recordingbinder binder = new recordingbinder(stage); for (module module : modules) { binder.install(module); } binder.scanforannotatedmethods(); for (recordingbinder child : binder.privatebinders) { child.scanforannotatedmethods(); } // free the memory consumed by the stack trace elements cache stacktraceelements.clearcache(); return collections.unmodifiablelist(binder.elements); }

True: records the elements executed by

Pred: returns the list of the given

Оценка: 3 — в целом, суть похожа.

Пример 9: 

Code: static proofnode<owlaxiom> canconvertstep(proofstep<owlaxiom> step) { if (step.getname() != elkclassinclusionexistentialcomposition.name) { return null; } list<? extends proofnode<owlaxiom>> premises = step.getpremises(); proofnode<owlaxiom> lastpremise = premises.get(premises.size() - 1); collection<? extends proofstep<owlaxiom>> lastpremisesteps = lastpremise .getinferences(); if (lastpremisesteps.size() != 1) { return null; } // else for (proofstep<owlaxiom> lastpremisestep : lastpremisesteps) { if (lastpremisestep .getname() == elkpropertyinclusionoftransitiveobjectproperty.name) { return lastpremisestep.getpremises().get(0); } } // else return null; }

True: checks if is derived by inference where the last premise is derived from

Pred: determine if the expression has been cleared.

Оценка: 2 — не очень похоже на правду, но сгенерированный комментарий вполне осмысленный.

Пример 10:

Code: public string asjsonstring(object o) { if ( getcoderspecific() instanceof jsonfactory == false ) { return "can be called on jsonconfiguration only"; } else { return new string(asbytearray(o), standardcharsets.utf_8); } }

True: utility/debug method. use "asbytearray" for programmatic use as the byte array will already by utf-8 and ready to be sent on network.

Pred: convert a json string to a json string.

========================================

Code: private void notifylisteners(string str) { writerlistener[] writerlisteners = null; synchronized (listeners) { writerlisteners = new writerlistener[listeners.size()]; listeners.toarray(writerlisteners); } for (int i = 0; i < writerlisteners.length; i++) { writerlisteners[i].write(str); } }

True: notify that a new string has been written.

Pred: notifies all listeners.

Оценка: 1 — модель не смогла уловить суть.

Средняя субъективная оценка: (1+4+5+3+2+1+3+3+2+1)/10 = 2,5 — вполне неплохой результат для модели, которая училась на 1 % от общего объёма обучающих данных. В целом, суть сгенерированных комментариев понятна, но если у вас есть ресурсы, чтобы обучить модель более чем на 1 % набора, то вы можете улучшить результат.

Заключение

Я показал, что после обучения модели даже на 1 % данных она выполняет свою цель и может вполне адекватно генерировать комментарии к коду. Также продемонстрирована предварительная обработка текста для языка Java. Если модель будет использоваться в исследовании целой кодовой базы, то лучше её всё же обучить на всех данных. 

Также следует сказать, что если обучить модель на большем объёме, то её можно встроить в IDE (VisualStudio, PyCharm и т.д.) Подробнее об этом можно посмотреть здесь.