File size: 3,635 Bytes
3a464db
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
import evaluate as hf_evaluate
import re
from datasets import Dataset
import ast
# Hugging Face code_eval
try:
    pass_at_k = hf_evaluate.load("code_eval")
except Exception as e:
    raise e

def process_docs(dataset: Dataset) -> Dataset:
    def _process_doc(doc):
        doc["text"] = doc["prompt"]
        doc["test_list_str"] = "\n".join(doc["test"])
        return doc
    return dataset.map(_process_doc)


def pass_at_1(references, predictions):
    processed_predictions = []
    for pred in predictions:
        pred_after_step1 = humaneval_postprocess(pred)
        pred_after_step2 = _process_answer(pred_after_step1)
        processed_predictions.append(pred_after_step2)

    results = pass_at_k.compute(
        references=list(references),
        predictions=[processed_predictions],
        k=[1],
    )
    return results[0]["pass@1"]


def humaneval_postprocess(text: str) -> str:
    blocks = re.findall(r'```(?:python)?\n(.*?)\n```', text, re.DOTALL)
    if blocks:
        return blocks[0].strip()
    return text.strip()

# def _process_answer(text: str) -> str:
#     patterns = [
#         r"'(.*)'\s*$$DONE$$",
#         r"$$BEGIN$$\s*'(.*)'\s*$$DONE$$",
#         r"BEGIN\s*'(.*)'\s*$$DONE$$",
#         r"$$BEGIN$$\s*'(.*)'\s*DONE",
#         r"BEGIN\s*'(.*)'\s*DONE",
#         r"$$BEGIN$$\s*'(.*)\s*$$DONE$$",
#         r"BEGIN\s*'(.*)\s*$$DONE$$",
#         r"$$BEGIN$$\s*'(.*)\s*DONE",
#         r"BEGIN\s*'(.*)\s*DONE",
#         r'$$BEGIN$$\s*(.*)\s*$$DONE$$',
#         r'BEGIN\s*(.*)\s*$$DONE$$',
#         r'$$BEGIN$$\s*(.*)\s*DONE',
#         r'BEGIN\s*(.*)\s*DONE',
#         r'```python\s*(.*)\s*```',
#         r'```\s*(.*)\s*```',
#         r'```python\s*(.*)\s*$',
#         r'```\s*(.*)\s*$',
#         r'(.*)\s*```.*',
#         r"\[BEGIN\]\s*'(.*)",
#         r'\[BEGIN\](.*)',
#     ]

#     for p in patterns:
#         match = re.search(p, text, re.DOTALL)
#         if match:
#             text = match.group(1)
#             break
            
#     text = text.split('```')[0]
#     text = re.split(r"'?\s*\$\$?DONE\$\$?", text)[0]
#     text = text.replace('\\_', '_')
#     text = text.strip()
#     return text


def _process_answer(text: str) -> str:
    patterns = [
        r"'(.*)'\s*$$DONE$$",
        r"$$BEGIN$$\s*'(.*)'\s*$$DONE$$",
        r"BEGIN\s*'(.*)'\s*$$DONE$$",
        r"$$BEGIN$$\s*'(.*)'\s*DONE",
        r"BEGIN\s*'(.*)'\s*DONE",
        r"$$BEGIN$$\s*'(.*)\s*$$DONE$$",
        r"BEGIN\s*'(.*)\s*$$DONE$$",
        r"$$BEGIN$$\s*'(.*)\s*DONE",
        r"BEGIN\s*'(.*)\s*DONE",
        r'$$BEGIN$$\s*(.*)\s*$$DONE$$',
        r'BEGIN\s*(.*)\s*$$DONE$$',
        r'$$BEGIN$$\s*(.*)\s*DONE',
        r'BEGIN\s*(.*)\s*DONE',
        r'```python\s*(.*)\s*```',
        r'```\s*(.*)\s*```',
        r'```python\s*(.*)\s*$',
        r'```\s*(.*)\s*$',
        r'(.*)\s*```.*',
        r"\[BEGIN\]\s*'(.*)",
        r'\[BEGIN\](.*)',
    ]

    for p in patterns:
        match = re.search(p, text, re.DOTALL)
        if match:
            text = match.group(1)
            break
            
    text = text.split('```')[0]
    text = re.split(r"'?\s*\$\$?DONE\$\$?", text)[0]
    text = text.replace('\\_', '_')
    text = text.strip()

    try:
        tree = ast.parse(text)
        filtered_nodes = [
            node for node in tree.body 
            if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef, ast.ClassDef, ast.Import, ast.ImportFrom))
        ]
        if filtered_nodes:
            import astor 
            text = "\n\n".join([ast.unparse(node) for node in filtered_nodes])
    except Exception:
        pass

    return text.strip()