Views
No views yet
1pip install transformers
2pip install peft
3pip install bitsandbytes
4pip install acceleratebash merge.sh1from transformers import AutoTokenizer, AutoModelForCausalLM
2from peft import LoraConfig, get_peft_model, prepare_model_for_int8_training
3import torch
4
5
6# load model and tokenizer
7
8tokenizer = AutoTokenizer.from_pretrained('bigcode/starcoderbase', use_auth_token=True)
9
10model = AutoModelForCausalLM.from_pretrained(
11 "StarCoder-15B_for_NMT/Epoch_1/-merged",
12 use_auth_token=True,
13 use_cache=True,
14 load_in_8bit=True,
15 device_map="auto"
16 )
17
18model = prepare_model_for_int8_training(model)
19
20lora_config = LoraConfig(
21 r=16,
22 lora_alpha=32,
23 lora_dropout=0.05,
24 bias="none",
25 task_type="CAUSAL_LM",
26 target_modules = ["c_proj", "c_attn", "q_attn"]
27)
28
29model = get_peft_model(model, lora_config)
30
31
32# a bug-fix pairs
33
34buggy_code = """
35 public MultiplePiePlot(CategoryDataset dataset){
36 super();
37// bug_start
38 this.dataset=dataset;
39// bug_end
40 PiePlot piePlot=new PiePlot(null);
41 this.pieChart=new JFreeChart(piePlot);
42 this.pieChart.removeLegend();
43 this.dataExtractOrder=TableOrder.BY_COLUMN;
44 this.pieChart.setBackgroundPaint(null);
45 TextTitle seriesTitle=new TextTitle("Series Title",new Font("SansSerif",Font.BOLD,12));
46 seriesTitle.setPosition(RectangleEdge.BOTTOM);
47 this.pieChart.setTitle(seriesTitle);
48 this.aggregatedItemsKey="Other";
49 this.aggregatedItemsPaint=Color.lightGray;
50 this.sectionPaints=new HashMap();
51 }
52"""
53
54fixed_code = """
55// fix_start
56 setDataset(dataset);
57// fix_end
58"""
59
60# model inference
61
62input_text = '<commit_before>\n' + buggy_code + '\n<commit_after>\n'
63input_ids = tokenizer(input_text, return_tensors="pt").input_ids.to(0)
64
65eos_id = tokenizer.convert_tokens_to_ids(tokenizer.eos_token)
66generated_ids = model.generate(
67 input_ids=input_ids,
68 max_new_tokens=256,
69 num_beams=10,
70 num_return_sequences=10,
71 early_stopping=True,
72 pad_token_id=eos_id,
73 eos_token_id=eos_id
74)
75
76for generated_id in generated_ids:
77 generated_text = tokenizer.decode(generated_id, skip_special_tokens=False)
78 patch = generated_text.split('\n<commit_after>\n')[1]
79 patch = patch.replace('<|endoftext|>','')
80 print(patch)
81
82