Views
No views yet
1from transformers import AutoModelForCausalLM, AutoTokenizer
2
3faulty ="""
4def add (x, y):
5 \"\"\"return sum of x and y\"\"\"
6 return x * y
7"""
8
9
10PROGRAM_REPAIR_TEMPLATE = f"""
11You are an expert in the field of software testing.
12You are given a buggy Python program, you are supposed to first generate testcases that can expose the bug,
13and then generate the corresponding fixed code. The two tasks are detailed as follows.
14
151. **Generate a comprehensive set of test cases to expose the bug**:
16 - Each test case should include an input and the expected output.
17 - Output the test cases as a JSON list, where each entry is a dictionary with keys `"test_input"` and `"test_output"`.
18 - Write in ```json ``` block.
19
202. **Provide a fixed version**:
21 - Write a correct Python program to fix the bug.
22 - Write in ```python ``` block.
23 - The code should read from standard input and write to standard output, matching the input/output format specified in the problem.
24
25Here is an example.
26The faulty Python program is:
27\`\`\`python
28\"\"\"Please write a Python program to sum two integer inputs\"\"\"
29def add (x, y):
30 return x - y
31x = int(input())
32y = int(input())
33print(add(x,y))
34\`\`\`
35
36Testcases that can expose the bug:
37\`\`\`json
38[
39 {{
40 \"test_input\":\"1\n2\",
41 \"test_output\":\"3\"
42 }},
43 {{
44 \"test_input\":\"-1\n1\",
45 \"test_output\":\"0\"
46 }},
47 {{
48 \"test_input\":\"-1\n2\",
49 \"test_output\":\"1\"
50 }}
51]
52\`\`\`
53
54Fixed code:
55\`\`\`python
56def add (x, y):
57 return x + y
58x = int(input())
59y = int(input())
60print(add(x,y))
61\`\`\`
62
63Now, you are given a faulty Python function, please return:
641. **Testcases** that helps expose the bug.
652. **Fixed code** that can pass all testcases.
66
67The faulty function is:
68\`\`\`python
69{faulty}
70\`\`\`
71<|assistant|>
72"""
73
74model = AutoModelForCausalLM.from_pretrained(
75 "Neo111x/Falcon3-3B-Instruct-RL-CODE-RL",
76 trust_remote_code=True
77)
78tokenizer = AutoTokenizer.from_pretrained(
79 "Neo111x/Falcon3-3B-Instruct-RL-CODE-RL",
80 trust_remote_code=True
81)
82
83messages = [
84 {"role": "user", "content": PROGRAM_REPAIR_TEMPLATE}
85]
86text = tokenizer.apply_chat_template(
87 messages,
88 tokenize=False,
89)
90model_inputs = tokenizer([text], return_tensors="pt").to(model.device)
91
92generated_ids = model.generate(
93 **model_inputs,
94 max_new_tokens=512
95)
96generated_ids = [
97 output_ids[len(input_ids):] for input_ids, output_ids in zip(model_inputs.input_ids, generated_ids)
98]
99
100response = tokenizer.batch_decode(generated_ids, skip_special_tokens=True)[0]
101print(response)