(model_output: str, lmstyle: LMStyle, cot: bool = False)
| 46 | |
| 47 | |
| 48 | def extract_execution_code(model_output: str, lmstyle: LMStyle, cot: bool = False): |
| 49 | if cot: |
| 50 | if "[ANSWER]" in model_output: |
| 51 | model_output = model_output.split("[ANSWER]")[1].strip() |
| 52 | if "==" in model_output: |
| 53 | model_output = model_output.split("==")[1].strip() |
| 54 | if "[/ANSWER]" in model_output: |
| 55 | model_output = model_output.split("[/ANSWER]")[0].strip() |
| 56 | else: |
| 57 | model_output = model_output.split("\n")[0].strip() |
| 58 | return model_output.strip() |