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_NTR/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
54repair_template = "OtherTemplate"
55
56fixed_code = "
57// fix_start
58 setDataset(dataset);
59// fix_end
60"
61
62# model inference
63
64input_text = '<commit_before>\n' + buggy_code + '\n<commit_msg>\n' + repair_template + '\n<commit_after>\n'
65input_ids = tokenizer(input_text, return_tensors="pt").input_ids.to(0)
66
67eos_id = tokenizer.convert_tokens_to_ids(tokenizer.eos_token)
68generated_ids = model.generate(
69 input_ids=input_ids,
70 max_new_tokens=256,
71 num_beams=10,
72 num_return_sequences=10,
73 early_stopping=True,
74 pad_token_id=eos_id,
75 eos_token_id=eos_id
76)
77
78for generated_id in generated_ids:
79 generated_text = tokenizer.decode(generated_id, skip_special_tokens=False)
80 patch = generated_text.split('\n<commit_after>\n')[1]
81 patch = patch.replace('<|endoftext|>','')
82 print(patch)
83
84