Class Extractor

src/evaluate.py:22–100

class Extractor(object):
    def __init__(self, dataset, output_file) -> None:
        self.dataset = Dataset(dataset, data_path = os.path.join("test_datasets/", dataset.replace("/", "_")))
        self.output_file = output_file
        self.outputs = json.load(open(self.output_file, "r"))

        self.solutions = {}


    def get_entrypoint(self, instance):
        entry_point = ""
        if self.dataset.name == "openai_humaneval":
            entry_point += instance["entry_point"]

        return entry_point

    def process_solution(self, solution, instance):
        if self.dataset.name == "openai_humaneval" and not solution.startswith("def"):
            lines = solution.splitlines()
            if len(lines) > 0:
                lines[0] = "    " + lines[0]
            solution = instance["prompt"] + "\n".join(lines)
        elif self.dataset.name == "openai_humaneval":
            try:
                ast.parse(solution)
            except:
                lines = solution.splitlines()
                if len(lines) > 0:
                    lines[0] = "    " + lines[0]
                solution = instance["prompt"] + "\n".join(lines)

        return solution


    def get_solutions(self, codegen = False, chat = False):
        for index, prompt in enumerate(self.outputs):
            if prompt not in self.dataset.prompt2instance:
                print('Cannot find the prompt of instance #{} in dataset, skipped.'.format(index))
                continue
            instance = self.dataset.prompt2instance[prompt]
            solutions = []
            if not self.outputs[prompt][1]:
                continue
            if not isinstance(self.outputs[prompt][0], list):
                print(self.outputs[prompt][0])
                continue
            for code in self.outputs[prompt][0]:
                solution = sanitize(code, self.get_entrypoint(instance), codegen = codegen, global_code = True if self.dataset.name in ["NTU-NLP-sg/xCodeEval", "codeparrot/apps"] else False, chat = chat)
                solution = self.process_solution(solution, instance)
                if solution not in solutions:
                    solutions.append(solution)
            
            self.solutions[prompt] = solutions


    def save_solutions(self):
        filename = self.output_file.replace(".json", "_SOLUTIONS.json")
        with open(filename, "w", encoding = "utf-8") as f: