Views
No views yet
| Model | Java pass@1 | Java Average number of test scenarios | Python pass@1 | Python Average number of test scenarios | Javascript pass@1 | Javascript Average number of test scenarios |
|---|---|---|---|---|---|---|
| TestGPT-7B | 48.6% | 4.37 | 35.67% | 3.56 | 36% | 2.76 |
| CodeLlama-13B-Instruct | 40.54% | 1.08 | 30.57% | 1.65 | 31.7% | 3.13 |
| Qwen-14B-Chat | 10.81% | 2.78 | 15.9% | 1.32 | 9.15% | 4.22 |
| Baichuan2-13B-Chat | 13.5% | 2.24 | 12.7% | 2.12 | 6.1% | 3.31 |
| Model | pass@1 | Percentage of strong validation |
|---|---|---|
| TestGPT-7B | 71.1% | 100% |
1from modelscope import AutoModelForCausalLM, AutoTokenizer, snapshot_download, AutoConfig
2import torch
3
4HUMAN_ROLE_START_TAG = "<s>human\n"
5BOT_ROLE_START_TAG = "<s>bot\n"
6
7if __name__ == '__main__':
8 # 模型地址, 可以替换为本地模型地址
9 model_dir = snapshot_download('codefuse-ai/TestGPT-7B', revision = 'v1.0.0')
10
11 # 加载tokenizer
12 tokenizer = AutoTokenizer.from_pretrained(model_dir, trust_remote_code=True, use_fast=False, legacy=False)
13
14 eos_token = '</s>'
15 pad_token = '<unk>'
16
17 try:
18 tokenizer.eos_token = eos_token
19 tokenizer.eos_token_id = tokenizer.convert_tokens_to_ids(eos_token)
20 except:
21 print(tokenizer.eos_token, tokenizer.eos_token_id)
22
23 try:
24 tokenizer.pad_token = pad_token
25 tokenizer.pad_token_id = tokenizer.convert_tokens_to_ids(pad_token)
26 except:
27 print(tokenizer.pad_token, tokenizer.pad_token_id)
28
29 tokenizer.padding_side = "left"
30 print(f"tokenizer's eos_token: {tokenizer.eos_token}, pad_token: {tokenizer.pad_token}")
31 print(f"tokenizer's eos_token_id: {tokenizer.eos_token_id}, pad_token_id: {tokenizer.pad_token_id}")
32
33 # 配置
34 config, unused_kwargs = AutoConfig.from_pretrained(
35 model_dir,
36 use_flash_attn=True,
37 use_xformers=True,
38 trust_remote_code=True,
39 return_unused_kwargs=True)
40
41 # 加载模型
42 model = AutoModelForCausalLM.from_pretrained(
43 model_dir,
44 config=config,
45 device_map="auto",
46 torch_dtype=torch.bfloat16,
47 trust_remote_code=True,
48 use_safetensors=False,
49 ).eval()
50
51 # 推理生成测试用例
52 # 被测代码prompt，分为用例生成和assert补全
53 # 用例生成格式
54 prompt = '为以下Python代码生成单元测试\n' \
55 '```Python\ndef add(lst):\n return sum([lst[i] for i in range(1, len(lst), 2) if lst[i]%2 == 0])\n```\n'
56
57 # assert补全格式，目前仅支持java语言
58 # prompt = '下面是被测代码\n' \
59 # '```java\n' \
60 # 'public class BooleanUtils {\n ' \
61 # 'public static boolean and(final boolean... array) {\n ' \
62 # 'ObjectUtils.requireNonEmpty(array, "array");\n ' \
63 # 'for (final boolean element : array) {\n ' \
64 # 'if (!element) {\n return false;\n }\n }\n ' \
65 # 'return true;\n }\n}\n```\n' \
66 # '下面代码是针对上面被测代码生成的用例，请补全用例，生成assert校验\n' \
67 # '```java\n' \
68 # '@Test\npublic void testAnd_withAllTrueInputs() {\n ' \
69 # 'boolean[] input = new boolean[] {true, true, true};\n ' \
70 # 'boolean result = BooleanUtils.and(input);\n}\n\n@Test\npublic void testAnd_withOneFalseInput() {\n ' \
71 # 'boolean[] input = new boolean[] {true, false, true};\n ' \
72 # 'boolean result = BooleanUtils.and(input);\n}\n' \
73 # '```\n'
74
75 # 输入格式化处理
76 prompt = f"{HUMAN_ROLE_START_TAG}{prompt}{BOT_ROLE_START_TAG}"
77 inputs = tokenizer(prompt, return_tensors='pt', padding=True, add_special_tokens=False).to("cuda")
78
79 # 推理
80 outputs = model.generate(
81 inputs=inputs["input_ids"],
82 max_new_tokens=2048,
83 top_p=0.95,
84 temperature=0.2,
85 do_sample=True,
86 eos_token_id=tokenizer.eos_token_id,
87 pad_token_id=tokenizer.pad_token_id,
88 num_return_sequences=1,
89 )
90
91 # 结果处理
92 outputs_len = len(outputs)
93 print(f"output len is: {outputs_len}")
94 for index in range(0, outputs_len):
95 print(f"generate index: {index}")
96 gen_text = tokenizer.decode(outputs[index], skip_special_tokens=True)
97 print(gen_text)
98 print("===================")
99