Views
No views yet
1from transformers import BartForConditionalGeneration, BartTokenizer
2model = BartForConditionalGeneration.from_pretrained("NTUYG/ComFormer")
3tokenizer = BartTokenizer.from_pretrained("NTUYG/ComFormer")
4code = '''
5public static void copyFile( File in, File out )
6 throws IOException
7 {
8 FileChannel inChannel = new FileInputStream( in ).getChannel();
9 FileChannel outChannel = new FileOutputStream( out ).getChannel();
10 try
11 {
12// inChannel.transferTo(0, inChannel.size(), outChannel); // original -- apparently has trouble copying large files on Windows
13
14 // magic number for Windows, 64Mb - 32Kb)
15 int maxCount = (64 * 1024 * 1024) - (32 * 1024);
16 long size = inChannel.size();
17 long position = 0;
18 while ( position < size )
19 {
20 position += inChannel.transferTo( position, maxCount, outChannel );
21 }
22 }
23 finally
24 {
25 if ( inChannel != null )
26 {
27 inChannel.close();
28 }
29 if ( outChannel != null )
30 {
31 outChannel.close();
32 }
33 }
34 }
35 '''
36code_seq, sbt = utils.transformer(code) #can find in https://github.com/NTDXYG/ComFormer
37input_text = code_seq + sbt
38input_ids = tokenizer.encode(input_text, return_tensors="pt", max_length=256, truncation=True)
39summary_text_ids = model.generate(
40 input_ids=input_ids,
41 bos_token_id=model.config.bos_token_id,
42 eos_token_id=model.config.eos_token_id,
43 length_penalty=2.0,
44 max_length=30,
45 min_length=2,
46 num_beams=5,
47)
48comment = tokenizer.decode(summary_text_ids[0], skip_special_tokens=True)
49print(comment)@misc{yang2021comformer,
title={ComFormer: Code Comment Generation via Transformer and Fusion Method-based Hybrid Code Representation},
author={Guang Yang and Xiang Chen and Jinxin Cao and Shuyuan Xu and Zhanqi Cui and Chi Yu and Ke Liu},
year={2021},
eprint={2107.03644},
archivePrefix={arXiv},
primaryClass={cs.SE}
}