| 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] |
