Log each output key

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