@@ -51,17 +51,23 @@ def main(
51
51
dataset : Iterator [Union [RepoInstance , SimpleInstance ]] = load_dataset (
52
52
dataset_name , split = dataset_split
53
53
) # type: ignore
54
+ dataset_name = dataset_name .lower ()
54
55
spec = None
55
56
example = None
56
57
repo_name = None
57
58
dataset_type = None
58
59
for example in dataset :
59
60
if repo_or_repo_dir .endswith ("/" ):
60
61
repo_or_repo_dir = repo_or_repo_dir [:- 1 ]
61
- if "swe" in dataset_name . lower () :
62
+ if "swe" in dataset_name :
62
63
repo_name = example ["instance_id" ]
63
64
dataset_type = "swebench"
64
- elif "humaneval" in dataset_name .lower ():
65
+ elif (
66
+ "humaneval" in dataset_name
67
+ or "mbpp" in dataset_name
68
+ or "bigcodebench" in dataset_name
69
+ or "codecontests" in dataset_name
70
+ ):
65
71
repo_name = example ["instance_id" ]
66
72
dataset_type = "simple"
67
73
else :
@@ -130,7 +136,7 @@ def main(
130
136
)
131
137
132
138
# make patch file
133
- if "swe" in dataset_name . lower () :
139
+ if "swe" in dataset_name :
134
140
if branch == "reference" :
135
141
patch = (
136
142
example ["test" ]["patch" ] + "\n \n " + example ["test" ]["test_patch" ]
@@ -164,7 +170,7 @@ def main(
164
170
+ example ["test" ]
165
171
)
166
172
else :
167
- solution = open ( test_ids ). read ()
173
+ solution = test_ids
168
174
prompt = example ["prompt" ] if "prompt" in example .keys () else ""
169
175
matches = extract_code_blocks (solution )
170
176
if len (matches ) > 0 :
0 commit comments