File size: 1,392 Bytes
5374a2d |
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 extract_test_output_code(model_output: str):
outputlines = model_output.split("\n")
# find the last line startwith assert...
indexlines = [i for i, line in enumerate(outputlines) if line.startswith("assert")]
if indexlines:
return outputlines[indexlines[-1]]
# first try to extract ```python if not then try ```
indexlines = [
i
for i, line in enumerate(outputlines)
if "```python" in line or "```Python" in line
]
if indexlines:
start_index = indexlines[0]
else:
start_index = None
indexlines = [i for i, line in enumerate(outputlines) if "```" in line]
if start_index is not None:
indexlines = [i for i in indexlines if i > start_index]
indexlines = [start_index] + indexlines
if len(indexlines) < 2:
return ""
return "\n".join(outputlines[indexlines[0] + 1 : indexlines[1]])
def extract_execution_code(model_output: str, cot: bool = False):
if cot:
if "[ANSWER]" in model_output:
model_output = model_output.split("[ANSWER]")[1].strip()
if "==" in model_output:
model_output = model_output.split("==")[1].strip()
if "[/ANSWER]" in model_output:
model_output = model_output.split("[/ANSWER]")[0].strip()
else:
model_output = model_output.split("\n")[0].strip()
return model_output.strip() |