init project

Evaluation/HumanEval/utils/dataset.py

@@ -0,0 +1,61 @@
+import os
+import numpy as np
+import json
+
+class HumanEvalDataset:
+
+    def __init__(self, root, sample_num=1, language="python", issft=False):
+        """
+        root: the path to the HumanEval dataset
+        sample_num: the number of samples for each prompt
+        language: the language of the HumanEval dataset
+        issft: whether to use the SFT setting
+        """
+        self.root = root
+        self.data = open(os.path.join(self.root, f"humaneval-{language}.jsonl")).readlines()
+
+        tmp = self.get_qa_only_data(self.data, issft)
+        self.clean_data = []
+        for i in range(len(tmp)):
+            for j in range(sample_num):
+                self.clean_data.append(tmp[i])
+        self.stopwords = self.clean_data[0]["stopwords"]
+        np.random.seed(1234)
+        print(f"Read HumanEval from {root}, number of samples {len(self.clean_data)}")
+
+    def get_qa_only_data(self, data_json, sft=False):
+        """
+        data_json: the jsonl file of HumanEval
+        sft: whether to use the SFT setting
+        return: a list of dict, each dict contains the prompt, task_id and stopwords
+        """
+        ans = []
+        for line in data_json:
+            line = json.loads(line)
+            prompt = line["prompt"].strip()
+            if "prefix" in line:
+                origin_prompt = line["prefix"]
+            else:
+                origin_prompt = line["prompt"]
+
+            if sft:
+                prompt = f"""Below is an instruction that describes a task, paired with an input that provides further context.\nWrite a response that appropriately completes the request.\n\n### Instruction:\nWrite a program to perform the given task.\n\nInput:\n{prompt}\n\n### Response:\n"""
+            if "stop_tokens" in line:
+                s = line["stop_tokens"]
+            else:
+                s = []
+            ans.append({"prompt":prompt, "task_id":line["task_id"], "original_prompt": origin_prompt, "stopwords":s})
+        return ans
+
+    def __len__(self):
+        """
+        return the number of samples in the dataset
+        """
+        return len(self.clean_data)
+
+    def __getitem__(self, index):
+        """
+        return the sample at index
+        """
+        sample = self.clean_data[index]
+        return sample

Evaluation/HumanEval/utils/utils.py

@@ -0,0 +1,40 @@
+def cleanup_code(
+    code: str,
+    language_type: str = None,
+    dataset: str = None,
+    issft: bool = False,
+    stop_words = []
+):
+    """
+    Cleans up the generated code.
+    """
+
+    if language_type.lower() == "python":
+        if issft:
+            code = _clean_python_code_for_sft(code)
+        stop_words = ["\ndef", "\nclass", "\nif", "\n#", "\nprint"]
+        code = _truncate_code_at_stopwords(code, stop_words)
+    elif language_type.lower() == "ts":
+        code = _truncate_code_at_stopwords(code, stop_words + ["\nexport", "\nimport", "\nexport default", "\nimport default", "\nconsole.log"])
+    else:
+        code = _truncate_code_at_stopwords(code, stop_words)
+
+    return code
+
+def _clean_python_code_for_sft(code):
+    code = code.replace("\r", "")
+    if "```python" in code:
+        code_start_idx = code.index("```python")
+        code = code[code_start_idx:].replace("```python", "").strip()
+        end_idx = code.find("```") if "```" in code else len(code)
+        code = code[:end_idx].strip()
+
+    return code
+
+def _truncate_code_at_stopwords(code, stop_words):
+    min_stop_idx = len(code)
+    for stop_word in stop_words:
+        stop_index = code.find(stop_word)
+        if 0 <= stop_index < min_stop_idx:
+            min_stop_idx = stop_index
+    return code[:min_stop_idx]