Spaces:
Sleeping
Sleeping
| from langchain import OpenAI | |
| from langchain.prompts import PromptTemplate, FewShotPromptTemplate | |
| from langchain.chains import LLMChain | |
| EXAMPLES_PROMPT_TEMPLATE = PromptTemplate( | |
| input_variables=["input", "output"], | |
| template="Input: {input}\nOutput: {output}" | |
| ) | |
| PLAN_MOVE_PROMPT_EXAMPLES = [ | |
| {"input": "The piles contain 3, 5, 7 sticks", "output": "I'll take one stick from pile A"}, | |
| {"input": "The piles contain 2, 5, 7 sticks", "output": "I'll take one stick from pile B"}, | |
| {"input": "The piles contain 2, 5, 7 sticks", "output": "I'll take five stick from pile B"}, | |
| {"input": "The piles contain 1, 2, 3 sticks", "output": "I'll take two sticks from pile C"}, | |
| {"input": "The piles contain 0, 2, 3 sticks", "output": "I'll take one stick from pile C"}, | |
| {"input": "The piles contain 0, 2, 0 sticks", "output": "I'll take two sticks from pile B"}, | |
| ] | |
| PLAN_MOVE_PROMPT_FROM_STRING_EXAMPLES = FewShotPromptTemplate( | |
| examples=PLAN_MOVE_PROMPT_EXAMPLES, | |
| example_prompt=EXAMPLES_PROMPT_TEMPLATE, | |
| prefix="Nim is a two-player game of strategy in which players take turns removing objects from separate piles. " | |
| "The goal of the game is to remove the last sticks from a pile when the other piles contain 0 sticks. Each " | |
| "of these inputs represent a game state. For each of these game states please express a logical move that " | |
| "consists of taking some number of sticks from a pile. " | |
| "You may not take any sticks from a pile that contains 0 sticks. " | |
| "You may not take more sticks from a pile than it contains. " | |
| "You may only take sticks from one pile. ", | |
| suffix="Input: {text_game_state}\nOutput:", | |
| input_variables=["text_game_state"], | |
| example_separator="\n\n" | |
| ) | |
| EXEC_MOVE_PROMPT_EXAMPLES = [ | |
| {"input": "I'll take two sticks from pile A", "output": "0,2"}, | |
| {"input": "I'll take 3 sticks from the first pile", "output": "0,3"}, | |
| {"input": "I'll take two sticks from pile C", "output": "2,2"}, | |
| {"input": "I'll take one stick from the third pile", "output": "2,1"}, | |
| {"input": "From pile B remove 2 sticks", "output": "1,2"}, | |
| {"input": "I'll take the last stick from pile C", "output": "2,1"}, | |
| ] | |
| EXEC_MOVE_PROMPT_FROM_STRING_EXAMPLES = FewShotPromptTemplate( | |
| examples=EXEC_MOVE_PROMPT_EXAMPLES, | |
| example_prompt=EXAMPLES_PROMPT_TEMPLATE, | |
| prefix="Express every input as two numbers separated by a comma, where the first number is the zero index pile " | |
| "number and the second number is the number of sticks to remove.", | |
| suffix="Input: {move_to_express}\nOutput:", | |
| input_variables=["move_to_express"], | |
| example_separator="\n\n" | |
| ) | |
| def plan_move(text_game_state, temperature, api_key): | |
| llm = OpenAI(model_name='text-davinci-003', temperature=temperature, max_tokens=100, | |
| openai_api_key=api_key) | |
| llm_chain = LLMChain(llm=llm, prompt=PLAN_MOVE_PROMPT_FROM_STRING_EXAMPLES, verbose=False) | |
| planned_move = llm_chain.run({'text_game_state': text_game_state}).strip() | |
| return planned_move | |
| def execute_move(move_to_express, nim_game_env, api_key): | |
| llm = OpenAI(model_name='text-davinci-003', temperature=0.0, max_tokens=10, | |
| openai_api_key=api_key) | |
| llm_chain = LLMChain(llm=llm, prompt=EXEC_MOVE_PROMPT_FROM_STRING_EXAMPLES, verbose=False) | |
| step_tuple_str = llm_chain.run({'move_to_express': move_to_express}) | |
| step_tuple = tuple(int(x) for x in step_tuple_str.split(',')) | |
| try: | |
| step_result = nim_game_env.step(step_tuple) | |
| except ValueError: | |
| return "Invalid move!", [0, 0, 0], 0, True, None | |
| text_observation = "The piles contain " + ", ".join(str(x) for x in step_result[0]) + " sticks." | |
| return text_observation, step_result[0], step_result[1], step_result[2], step_result[3] | |