File size: 1,350 Bytes
fc35a48 |
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 |
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] |