167 lines
5.6 KiB
Python
167 lines
5.6 KiB
Python
import argparse
|
|
import json
|
|
import os
|
|
import random
|
|
|
|
from datasets import load_dataset
|
|
|
|
from eval.mbpp_eval.utils import compute_code_eval
|
|
|
|
|
|
def get_fewshot():
|
|
return """
|
|
You are an expert Python programmer, and here is your task: Write a function to find the shared elements from the given two lists. Your code should pass these tests:
|
|
|
|
assert set(similar_elements((3, 4, 5, 6),(5, 7, 4, 10))) == set((4, 5))
|
|
assert set(similar_elements((1, 2, 3, 4),(5, 4, 3, 7))) == set((3, 4))
|
|
assert set(similar_elements((11, 12, 14, 13),(17, 15, 14, 13))) == set((13, 14))
|
|
|
|
[BEGIN]
|
|
def similar_elements(test_tup1, test_tup2):
|
|
res = tuple(set(test_tup1) & set(test_tup2))
|
|
return (res)
|
|
[DONE]
|
|
|
|
You are an expert Python programmer, and here is your task: Write a python function to identify non-prime numbers. Your code should pass these tests:
|
|
|
|
assert is_not_prime(2) == False
|
|
assert is_not_prime(10) == True
|
|
assert is_not_prime(35) == True
|
|
assert is_not_prime(37) == False
|
|
|
|
[BEGIN]
|
|
import math
|
|
def is_not_prime(n):
|
|
result = False
|
|
for i in range(2,int(math.sqrt(n)) + 1):
|
|
if n % i == 0:
|
|
result = True
|
|
return result
|
|
[DONE]
|
|
|
|
You are an expert Python programmer, and here is your task: Write a function to find the n largest integers from a given list of numbers, returned in descending order. Your code should pass these tests:
|
|
|
|
assert heap_queue_largest( [25, 35, 22, 85, 14, 65, 75, 22, 58],3)==[85, 75, 65]
|
|
assert heap_queue_largest( [25, 35, 22, 85, 14, 65, 75, 22, 58],2)==[85, 75]
|
|
assert heap_queue_largest( [25, 35, 22, 85, 14, 65, 75, 22, 58],5)==[85, 75, 65, 58, 35]
|
|
|
|
[BEGIN]
|
|
import heapq as hq
|
|
def heap_queue_largest(nums,n):
|
|
largest_nums = hq.nlargest(n, nums)
|
|
return largest_nums
|
|
[DONE]
|
|
|
|
"""
|
|
|
|
|
|
EXAMPLE_TEMPLATE = '''You are an expert Python programmer, and here is your task: {text}\nYour code should pass these tests:\n\n{tests}\n\n'''
|
|
EXAMPLE_TEMPLATE_493 = '''You are an expert Python programmer, and here is your task: {text}\n\ncalculate_polygons(startx, starty, endx, endy, radius)\n\n'''
|
|
|
|
|
|
def remove_extra_symbols(code):
|
|
lines = code.split("\n")
|
|
|
|
# se_lines = [line.startswith("```") for line in lines]
|
|
outputs = []
|
|
if "```" in lines[0]:
|
|
lines = lines[1:]
|
|
|
|
for line in lines:
|
|
if not line.startswith("```"):
|
|
outputs.append(line)
|
|
else:
|
|
break
|
|
|
|
return "\n".join(outputs)
|
|
|
|
|
|
def extract_code(raw_completions):
|
|
missing = 0
|
|
for item in raw_completions:
|
|
if "[BEGIN]" not in item["completion"] or "[END]" not in item["completion"]:
|
|
if "```python" in item["completion"] or "```" in item["completion"]:
|
|
s1 = item["completion"].find("```python")
|
|
s2 = item["completion"].find("```")
|
|
if s1 == -1:
|
|
s = s2 + 3
|
|
else:
|
|
s = s1 + len("```python")
|
|
e = item["completion"].find("```", s)
|
|
if e == -1:
|
|
missing += 1
|
|
print(f"Warning: {item['completion']}")
|
|
continue
|
|
code = item["completion"][s:e].strip()
|
|
else:
|
|
missing += 1
|
|
continue
|
|
else:
|
|
s = item["completion"].index("[BEGIN]") + len("[BEGIN]")
|
|
e = item["completion"].index("[END]")
|
|
code = item["completion"][s:e].strip()
|
|
code = remove_extra_symbols(code)
|
|
|
|
item["completion"] = code
|
|
print(f"Missing {missing} segments of code.")
|
|
return raw_completions
|
|
|
|
|
|
def main():
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument("--prediction_file", type=str)
|
|
parser.add_argument("--sanitized", default=False, action="store_true")
|
|
parser.add_argument("--save_dir", type=str)
|
|
args = parser.parse_args()
|
|
|
|
outputs = [json.loads(line) for line in open(args.prediction_file).readlines()]
|
|
|
|
random.seed(42)
|
|
|
|
if not os.path.exists(args.save_dir):
|
|
os.makedirs(args.save_dir, exist_ok=True)
|
|
|
|
if args.sanitized:
|
|
test_data = load_dataset("mbpp", "sanitized", split="test").to_list()
|
|
else:
|
|
test_data = load_dataset("mbpp", split="test").to_list()
|
|
|
|
print("Number of examples:", len(test_data))
|
|
|
|
assert len(test_data) == len(outputs)
|
|
# predictions = [{"task_id": example["task_id"], "prompt": example[prompt_key], "completion": output} for
|
|
# example, output in zip(duplicate_test_data, outputs)]
|
|
predictions = extract_code(outputs)
|
|
|
|
predictions_code_only = [[] for _ in range(len(test_data))]
|
|
for i in range(len(predictions)):
|
|
# predictions_code_only[i // args.unbiased_sampling_size_n].append(predictions[i]["completion"])
|
|
predictions_code_only[i].append(predictions[i]["completion"])
|
|
reference_test_list = ["\n".join(example["test_list"]) for example in test_data]
|
|
assert len(predictions_code_only) == len(reference_test_list)
|
|
|
|
os.environ["HF_ALLOW_CODE_EVAL"] = "1"
|
|
os.environ["TOKENIZERS_PARALLELISM"] = "false"
|
|
pass_at_k_results, eval_results = compute_code_eval(
|
|
references=reference_test_list,
|
|
predictions=predictions_code_only,
|
|
num_workers=1,
|
|
)
|
|
|
|
for item, result in zip(predictions, eval_results.values()):
|
|
result.sort()
|
|
item["passed"] = result[0][1]["passed"]
|
|
|
|
prediction_save_path = os.path.join(args.save_dir, "mbpp_eval_predictions.json")
|
|
with open(prediction_save_path, "w") as fout:
|
|
json.dump(predictions, fout)
|
|
|
|
print(pass_at_k_results)
|
|
|
|
with open(os.path.join(args.save_dir, "metrics.json"), "w") as fout:
|
|
json.dump(pass_at_k_results, fout)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|