Log each output key

This commit is contained in:
Lance Martin
2024-02-20 19:56:22 -08:00
parent a33bfdac2b
commit a902b3c19b
@@ -335,6 +335,7 @@
" question = state_dict[\"question\"]\n",
" docs = state_dict[\"docs\"]\n",
" code_solution = state_dict[\"generation\"]\n",
" prefix = code_solution[0].prefix\n",
" imports = code_solution[0].imports\n",
" code = code_solution[0].code\n",
" code_block = imports +\"\\n\"+ code\n",
@@ -354,8 +355,13 @@
" # No errors occurred\n",
" error = \"None\"\n",
"\n",
" return {\"keys\": {\"generation\": code_solution, \"question\": question, \"error\": error, \"docs\": docs}}\n",
"\n",
" return {\"keys\": {\"generation\": code_solution, \n",
" \"question\": question, \n",
" \"error\": error, \n",
" \"docs\": docs,\n",
" \"prefix\":prefix,\n",
" \"imports\":imports,\n",
" \"code\":code}}\n",
"\n",
"### Edges\n",
"\n",
@@ -364,7 +370,7 @@
" Determines whether to test code execution, or re-try answer generation.\n",
"\n",
" Args:\n",
" state (dict): The current state of the agent, including all keys.\n",
" state (dict): The current graph state\n",
"\n",
" Returns:\n",
" str: Next node to call\n",
@@ -391,7 +397,7 @@
" Determines whether to finish.\n",
"\n",
" Args:\n",
" state (dict): The current state of the agent, including all keys.\n",
" state (dict): The current graph state\n",
"\n",
" Returns:\n",
" str: Next node to call\n",
@@ -480,30 +486,29 @@
"@run_evaluator\n",
"def check_import(run: Run, example: Union[Example, None] = None):\n",
" model_outputs = run.outputs[\"keys\"]\n",
" code_block = model_outputs['generation'][0]['repr']\n",
" print(\"CODE BLOCK!\")\n",
" print(code_block)\n",
" imports = model_outputs['imports']\n",
" try:\n",
" exec(code_block.imports)\n",
" exec(imports)\n",
" score = 1\n",
" print(\"Score: 1!\")\n",
" except:\n",
" score = 0\n",
" print(\"Score: 0!\")\n",
" return EvaluationResult(key=\"check_import\", score=score)\n",
"\n",
"@run_evaluator\n",
"def check_execution(run: Run, example: Union[Example, None] = None):\n",
" model_outputs = run.outputs[\"keys\"]\n",
" code_block = model_outputs['generation'][0]['repr']\n",
" print(\"CODE BLOCK!\")\n",
" print(code_block)\n",
" imports = code_block.imports\n",
" code = code_block.code\n",
" imports = model_outputs['imports']\n",
" code = model_outputs['code']\n",
" code_to_execute = imports +\"\\n\"+ code\n",
" try:\n",
" exec(code_to_execute)\n",
" score = 1\n",
" print(\"Score: 1!\")\n",
" except:\n",
" score = 0\n",
" print(\"Score: 0!\")\n",
" return EvaluationResult(key=\"check_execution\", score=score)"
]
},