init project

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]