Views
No views yet
1import torch
2from transformers import AutoTokenizer, AutoModelForCausalLM
3from peft import PeftModel
4import re
5
6model_id_pretrained = 'WizardLMTeam/WizardCoder-Python-13B-V1.0'
7model_id_finetuned = 'szalontaib/MegaBugInject'
8
9tokenizer = AutoTokenizer.from_pretrained(model_id_pretrained, add_eos_token=False)
10model = AutoModelForCausalLM.from_pretrained(model_id_pretrained, device_map='auto', dtype=torch.float16, trust_remote_code=True)
11model = PeftModel.from_pretrained(model, model_id_finetuned)
12
13
14def extract_diff(model_output):
15 pattern = re.compile(r'\s*\[DIFF\](.*?)\[/DIFF\]\s*', re.DOTALL)
16 matches = pattern.findall(model_output)
17 if matches:
18 return matches[0].strip('\n')
19 return None
20
21def diff2code(diff : str) -> str:
22 return '\n'.join(
23 line[2:] for line in diff.splitlines()
24 if not line.startswith('-')
25 ).strip()
26
27def corrupt(program, model, tokenizer, **generation_kwargs):
28 prompt = f'[PYTHON]\n{program.strip()}\n[/PYTHON]\n[DIFF]\n'
29 model_inputs = tokenizer([prompt], return_tensors="pt").to(model.device)
30 generated_ids = model.generate(**model_inputs, **generation_kwargs)
31 outputs = tokenizer.batch_decode(generated_ids, skip_special_tokens=True)
32 diffs = [extract_diff(output) for output in outputs]
33 corrupted_programs = [diff2code(diff) for diff in diffs if diff is not None]
34 return corrupted_programs
35
36
37test_code = '''
38def bitcount(n):
39 count = 0
40 while n:
41 n &= n - 1
42 count += 1
43 return count
44'''.strip()
45
46corrupted_programs = corrupt(
47 test_code, model, tokenizer,
48 do_sample=True,
49 temperature=0.5,
50 max_new_tokens=4096,
51 num_return_sequences=5,
52)
53
54for corrupted_program in corrupted_programs:
55 print('-'*30)
56 print(corrupted_program)@misc{szalontai2026diffbasedcodecorruptionusing,
title={Diff-Based Code Corruption using LLMs for Large-Scale Bugfix Benchmarking},
author={Balázs Szalontai and Ábel Szauter and Balázs Márton and Péter Verebics and Balázs Pintér and Tibor Gregorics},
year={2026},
eprint={2606.29088},
archivePrefix={arXiv},
primaryClass={cs.SE},
url={https://arxiv.org/abs/2606.29088},
}