| 2 | |
| 3 | |
| 4 | def extract_code(model_output: str, lmstyle: LMStyle): |
| 5 | outputlines = model_output.split("\n") |
| 6 | if lmstyle == LMStyle.CodeLLaMaInstruct: |
| 7 | indexlines = [i for i, line in enumerate(outputlines) if "PYTHON]" in line] |
| 8 | if len(indexlines) < 2: |
| 9 | indexlines = [i for i, line in enumerate(outputlines) if "```" in line] |
| 10 | elif lmstyle == LMStyle.GenericBase: |
| 11 | return model_output.strip() |
| 12 | else: |
| 13 | indexlines = [i for i, line in enumerate(outputlines) if "```" in line] |
| 14 | if len(indexlines) < 2: |
| 15 | return "" |
| 16 | return "\n".join(outputlines[indexlines[0] + 1 : indexlines[1]]) |
| 17 | |
| 18 | |
| 19 | def extract_test_output_code(model_output: str, lmstyle: LMStyle = None): |