Cleans up the generated code.
(text: str, language_type: str, reference)
| 170 | |
| 171 | |
| 172 | def _clean_up_code(text: str, language_type: str, reference) -> str: |
| 173 | """Cleans up the generated code.""" |
| 174 | try: |
| 175 | # for chatGLM related text |
| 176 | eval_text = eval(text) |
| 177 | except Exception: |
| 178 | pass |
| 179 | else: |
| 180 | if isinstance(eval_text, str): |
| 181 | text = eval_text |
| 182 | # extract code from code block |
| 183 | text = text.lstrip('\n') |
| 184 | if '```' in text: |
| 185 | blocks = re.findall(r'```(.*?)```', text, re.DOTALL) |
| 186 | if len(blocks) == 0: |
| 187 | text = text.split('```')[1] # fall back to default strategy |
| 188 | else: |
| 189 | text = blocks[0] # fetch the first code block |
| 190 | if not text.startswith('\n'): # in case starting with ```xxx |
| 191 | text = text[max(text.find('\n') + 1, 0):] |
| 192 | if language_type.lower() == 'python': |
| 193 | text = humaneval_postprocess_v2(text) |
| 194 | # we need to take care of the first line |
| 195 | # append extra space for first line for correct indentation |
| 196 | text = ' ' + text.lstrip() |
| 197 | |
| 198 | text_splits = text.split('\n') |
| 199 | is_empty_line = False |
| 200 | ind_empty_line = None |
| 201 | for i, line in enumerate(text_splits): |
| 202 | if len(line.strip()) > 0 and line[0] != ' ' and line[0] != '\t': |
| 203 | is_empty_line = True |
| 204 | ind_empty_line = i |
| 205 | break |
| 206 | if is_empty_line: |
| 207 | text = '\n'.join(text_splits[:ind_empty_line]) |
| 208 | else: |
| 209 | end_words = [ |
| 210 | '\ndef', '\nclass', '\n#', '\nassert', '\n"""', '\nprint', |
| 211 | '\nif', '\n\n\n' |
| 212 | ] |
| 213 | for w in end_words: |
| 214 | if w in text: |
| 215 | text = text[:text.rfind(w)] |
| 216 | # strip function head for all other language |
| 217 | func_name = reference.strip().split('\n')[-1] |
| 218 | if func_name: |
| 219 | func_name = func_name.strip().strip('{') |
| 220 | if func_name in text: |
| 221 | text = '\n'.join(text[text.find(func_name):].split('\n')[1:]) |
| 222 | if language_type.lower() == 'java': |
| 223 | main_pos = text.find('public static void main') |
| 224 | if main_pos != -1: |
| 225 | text = text[:main_pos] + '}' |
| 226 | if '}' in text: |
| 227 | text = text[:text.rfind('}')] + '}' |
| 228 | if text.count('{') + 1 == text.count('}'): |
| 229 | text += '\n}' |
no test coverage detected