1 line
32 KiB
JSON
1 line
32 KiB
JSON
{"repo": "211820252/temp-dobevi", "files": {"DoBeVi/src/search/deepseek_algo.py": {"bf32ca96d380d73392bbf88835c2b1f7863e87e4": {"t": 1785747502.115636, "n": 15829, "h": "af730d470958c2e2", "c": "import asyncio\nimport logging\nimport time\nfrom typing import List, Optional, Tuple\nimport torch\nimport math\nfrom openai import OpenAI\nimport time\nimport logging\nimport re\n\nfrom dojo import (\n TracedTheorem,\n Dojo,\n DojoCrashError,\n DojoInitError,\n LeanError,\n TacticState,\n ProofFinished,\n ProofGivenUp,\n)\n\nfrom search.search_algo import (\n Prover,\n SearchResult,\n client\n)\n\nfrom search.search_tree import (\n Status,\n SolvedNode,\n UnsolvedNode,\n InvalidNode,\n Edge,\n collect_success_edges,\n)\n\nfrom search.tactic_generator import (\n ModelEmptyOutputError,\n)\n\nfrom search.visual import (\n visualize_proof_tree,\n)\n\ndeepseek = OpenAI(api_key=\"sk-75c0610bc70f43df9a0aadf63e0f1a00\", base_url=\"https://api.deepseek.com\")\n\nasync def evaluate_func(current_state: str, tactic: str, new_state: str) -> float:\n if len(new_state) > 4096:\n logging.warning(\"New state exceeds 4096 characters, returning 0.00\")\n return 0.00\n attempts = 5 # 设置重试次数\n for attempt in range(attempts):\n try:\n response = deepseek.chat.completions.create(\n model=\"deepseek-chat\",\n # messages=[\n # {\n # \"role\": \"system\",\n # \"content\": \"\\\n # You are a professional math assistant specialized in Lean 4 theorem proving. \\\n # The user will provide the current proof state, the tactic applied, and the resulting new state after applying the tactic. \\\n # Your task is to evaluate the effectiveness of the applied tactic in advancing the proof. \\\n # Please provide a score between 0 and 1 with exactly two decimal places. \\\n # The score should follow these rules: \\\n # - 0.00: The tactic is ineffective or invalid. \\\n # - 1.00: The tactic is highly effective and advances the proof significantly. \\\n # - Intermediate values: Indicate varying degrees of effectiveness, rounded to two decimal places. \\\n # The score should be in the format of a number like 0.00, 0.25, 0.50, etc. \\\n # Make sure to output only the score as a number, with no additional explanation or text.\"\n # },\n # {\n # \"role\": \"user\",\n # \"content\": f\"Current State: {current_state}\\nTactic Applied: {tactic}\\nNew State: {new_state}\"\n # }\n # ],\n messages=[\n {\n \"role\": \"system\",\n \"content\": \"\\\n You are a professional math assistant specialized in Lean 4 theorem proving. \\\n The user will provide the current proof state, the tactic applied, and the resulting new state after applying the tactic. \\\n Your task is to evaluate the effectiveness of the applied tactic in advancing the proof. \\\n Please provide a score between 0 and 1 with exactly one decimal places. \\\n The score should follow these rules: \\\n - 0.0: The tactic is ineffective or invalid. \\\n - 1.0: The tactic is highly effective and advances the proof significantly. \\\n - Intermediate values: Indicate varying degrees of effectiveness, rounded to one decimal places. \\\n The score should be in the format of a number like 0.1, 0.5, 0.7, etc. \\\n Make sure to output only the score as a number, with no additional explanation or text.\"\n },\n {\n \"role\": \"user\",\n \"content\": f\"Current State: {current_state}\\nTactic Applied: {tactic}\\nNew State: {new_state}\"\n }\n ],\n stream=False\n )\n\n # print(response.choices[0].message.content)\n\n # score = re.search(r\"effectiveness of this tactic as \\*\\*([0-9.]+)\\*\\*\", response.choices[0].message.content)\n \n # if score:\n # score = score.group(1)\n # print(\"Extracted score:\", score)\n # else:\n # print(\"Score not found.\")\n\n score = float(response.choices[0].message.content)\n\n print(f\"score: {score}\")\n\n return score\n\n except Exception as e:\n # 捕获 OpenAI 的错误和空响应错误\n logging.error(f\"Error occurred: {e}. Attempt {attempt + 1} of {attempts} failed.\")\n if attempt < attempts - 1:\n time.sleep(3) # 等待 3 秒再重新执行\n else:\n logging.error(\"Max attempts reached, exiting.\")\n return 0.1 # 达到最大重试次数后退出\n\ndef log(x):\n if x <= 0:\n return float('-inf')\n return math.log(x)\n\nalpha = 1.00\n\nclass DeepSeekProver(Prover):\n def __init__(\n self,\n actor_id: int,\n num_gpus: int,\n search_timeout: int,\n max_expansions: Optional[int],\n num_sampled_tactics: int,\n result_save_path: str,\n ):\n super().__init__(actor_id, num_gpus, search_timeout, max_expansions, num_sampled_tactics, result_save_path)\n\n async def search(self, thm: TracedTheorem) -> None:\n \"\"\"\n Search for a proof of a theorem using best-first search.\n \"\"\"\n self._prepare()\n await client._test_connection()\n\n self.thm = thm\n \n self.min_beam_size = 4 \n self.max_beam_size = self.num_sampled_tactics\n\n try:\n # initialize the current root node\n self.dojo = Dojo(\n thm.root_dir, \n thm, \n self.leandojo_tactic_timeout,\n self.leandojo_num_threads,\n self.leandojo_memory_limit\n ).__enter__()\n\n current_state = self.dojo.init_state\n\n self.root = UnsolvedNode(\n leandojo_state=current_state,\n is_terminal=False,\n priority=0.0, \n depth=0,\n )\n\n self.nodes = {current_state.id: self.root}\n self.back_edges = []\n self.success_edges = []\n\n try:\n await self._best_first_search()\n except torch.OutOfMemoryError:\n logging.error(f\"🚨OOM when sampling theorem: {self.thm.name}\")\n torch.cuda.empty_cache()\n except DojoCrashError as e:\n logging.error(f\"🚨{e}\")\n \n except (asyncio.TimeoutError, asyncio.CancelledError):\n raise\n\n except DojoInitError as e:\n logging.error(f\"🚨Failed to initialize Dojo: {e}\")\n \n except Exception as e:\n logging.error(f\"🚨{type(e).__name__}: {e}\")\n \n finally: \n if hasattr(self, \"dojo\"):\n self.dojo.__exit__(None, None, None)\n \n async def _best_first_search(self) -> None:\n start_time = time.time()\n priority_queue = asyncio.PriorityQueue() # lowest priority first\n priority_queue.put_nowait((-self.root.priority / ((self.root.depth + 1) ** alpha), self.root))\n\n while True:\n if priority_queue.empty():\n logging.info(\"Search queue is empty.\")\n break\n try:\n await self._step(priority_queue)\n except (asyncio.TimeoutError, asyncio.CancelledError, DojoCrashError):\n raise\n except Exception as e:\n logging.error(f\"🚨{type(e).__name__}: {e}\")\n\n self.elapsed_time = time.time() - start_time\n if (self.elapsed_time > self.search_timeout) or (self.max_expansions and self.num_expansions >= self.max_expansions):\n if self.root.status == Status.SOLVED:\n logging.info(\"Search complete: proof found.\")\n else:\n logging.info(\"Exceeded time limit or max expansions.\")\n break\n\n if self.root.status == Status.INVALID:\n logging.info(\"Invalid tactic generated.\")\n break\n\n if self.root.status == Status.SOLVED:\n logging.info(\"Search complete: proof found.\")\n break\n\n async def _step(self, priority_queue: asyncio.PriorityQueue[Tuple[float, UnsolvedNode]]) -> None:\n _, search_node = priority_queue.get_nowait()\n \n logging.info(f\"Expanding node: {search_node}\")\n\n if isinstance(search_node.leandojo_state, TacticState):\n tactic_state_str = self._build_prompt(search_node)\n else:\n raise ValueError(f\"Invalid leandojo_state for search_node: {type(search_node.leandojo_state)}\")\n \n # generate tactics by LLM\n suggestions = await self._generate_tactics(tactic_state_str, search_node.depth)\n \n # try all the tactics\n results = []\n\n for tactic in suggestions:\n edge, proof_finished = await self._run_tactic(\n search_node, tactic, priority_queue\n )\n \n if edge is not None:\n results.append(edge)\n\n if proof_finished:\n break\n \n # expand the search node\n search_node.out_edges = results\n self.num_expansions += 1\n priority_queue.task_done()\n \n def _count_current_beam_size(\n self,\n depth: int,\n ) -> int:\n \n progress = self.num_expansions / self.max_expansions if self.max_expansions else 0\n factor = max(0, 1 - 15 * progress)\n beam = self.min_beam_size + (self.max_beam_size - self.min_beam_size) * factor\n \n return int(beam)\n\n @torch.no_grad()\n async def _generate_tactics(self, tactic_state_str: str, depth: int) -> List[Tuple[str, float]]:\n start_time = time.time()\n suggestions = await client.async_generate_tactic_sampling(\n self.select_tac_gen(),\n state=tactic_state_str,\n num_samples=self._count_current_beam_size(depth),\n # num_samples=self.num_sampled_tactics,\n )\n\n suggestions = [item['tactic'] for item in suggestions['suggestions']]\n \n if len(suggestions) == 0:\n raise ModelEmptyOutputError(\"No tactic generated.\")\n\n self.model_elapsed_time += time.time() - start_time\n return suggestions\n \n async def _run_tactic(\n self, \n search_node: UnsolvedNode, \n tactic: str, \n priority_queue: asyncio.PriorityQueue\n ) -> Tuple[Optional[Edge], bool]:\n # threshold = log(0.10)\n\n # run tactic in Lean server\n start_time = time.time()\n\n leandojo_new_state = await asyncio.to_thread(\n self.dojo.run_tac, search_node.leandojo_state, tactic\n )\n\n self.dojo_elapsed_time += time.time() - start_time # Tactics triggering timeout will not be counted\n\n # create a new child node\n depth = search_node.depth + 1\n if leandojo_new_state.id not in self.nodes:\n # proof finished\n if isinstance(leandojo_new_state, ProofFinished):\n score = log(1.00)\n child_node = SolvedNode(leandojo_state=leandojo_new_state, depth=depth)\n # invalid tactic\n elif type(leandojo_new_state) in (LeanError, ProofGivenUp):\n score = log(0.00)\n child_node = InvalidNode(leandojo_state=leandojo_new_state, depth=depth)\n # unsolved proof\n else:\n assert isinstance(leandojo_new_state, TacticState), f\"Expected TacticState, got{type(leandojo_new_state)}\"\n score = log(await evaluate_func(\n current_state=search_node.leandojo_state.pp,\n tactic=tactic,\n new_state=leandojo_new_state.pp\n ))\n # if score <= threshold: \n # return None, False # Skip tactics with low scores\n child_node = UnsolvedNode(\n leandojo_state=leandojo_new_state,\n is_terminal=False,\n priority=search_node.priority + score,\n depth=depth,\n )\n priority_queue.put_nowait((-child_node.priority / ((child_node.depth + 1) ** alpha), child_node))\n self.nodes[leandojo_new_state.id] = child_node\n edge = Edge(src=search_node, dst=child_node, tactic=tactic, score=score)\n else:\n assert isinstance(leandojo_new_state, TacticState), f\"Expected TacticState, got{type(leandojo_new_state)}\"\n child_node = self.nodes[leandojo_new_state.id]\n assert isinstance(child_node, UnsolvedNode), f\"Expected UnsolvedNode, got {type(child_node)}\"\n child_node.depth = min(child_node.depth, depth)\n if await asyncio.to_thread(child_node.is_descendant, search_node):\n edge = None\n else:\n score = log(await evaluate_func(\n current_state=search_node.leandojo_state.pp,\n tactic=tactic,\n new_state=leandojo_new_state.pp\n ))\n edge = Edge(src=search_node, dst=child_node, tactic=tactic, score=score)\n self.back_edges.append(edge)\n\n if isinstance(child_node, UnsolvedNode) and edge is not None:\n child_node.in_edges.append(edge)\n elif isinstance(child_node, SolvedNode) or isinstance(child_node, InvalidNode):\n child_node.in_edge = edge\n\n return edge, isinstance(leandojo_new_state, ProofFinished)\n \n def _build_prompt(\n self,\n search_node: UnsolvedNode,\n ) -> str:\n input_template = \"[GOAL]\\n{state}\\n[PROOFSTEP]\\n\"\n return input_template.format(state=search_node.leandojo_state.pp)\n \n def _normalize_scores(\n self,\n suggestions: List[Tuple[str, float]]\n ) -> List[Tuple[str, float, float]]:\n\n exps = [math.exp(score) for _, score in suggestions]\n total = sum(exps)\n \n return [\n (tactic, score, exp_score / total)\n for (tactic, score), exp_score in zip(suggestions, exps)\n ]\n \n async def get_result(self, visualize: bool = True) -> Optional[SearchResult]:\n if not self.thm or not self.root:\n return None\n \n # check if the search was successful\n if self.root.status == Status.SOLVED:\n proof = [e.tactic for e in await asyncio.to_thread(self.root.extract_proof)]\n if visualize:\n self.success_edges = await asyncio.to_thread(collect_success_edges, self.root)\n else:\n proof = None\n \n if visualize:\n await asyncio.to_thread(\n visualize_proof_tree,\n list(self.nodes.values()),\n self.success_edges, \n self.back_edges,\n self.result_save_path + \"/visual\",\n self.thm.name,\n ['simple','detail']\n )\n\n # box the result\n result = SearchResult(\n theorem=self.thm,\n status=self.root.status,\n proof=proof,\n num_total_nodes=len(self.nodes),\n num_expansions=self.num_expansions,\n elapsed_time=self.elapsed_time,\n dojo_elapsed_time=self.dojo_elapsed_time,\n model_elapsed_time=self.model_elapsed_time,\n )\n\n return result"}, "d008afcb4dcfac880788dbbfe98354ee3557d29e": {"t": 1785773629.4785745, "n": 15829, "h": "af730d470958c2e2", "c": "import asyncio\nimport logging\nimport time\nfrom typing import List, Optional, Tuple\nimport torch\nimport math\nfrom openai import OpenAI\nimport time\nimport logging\nimport re\n\nfrom dojo import (\n TracedTheorem,\n Dojo,\n DojoCrashError,\n DojoInitError,\n LeanError,\n TacticState,\n ProofFinished,\n ProofGivenUp,\n)\n\nfrom search.search_algo import (\n Prover,\n SearchResult,\n client\n)\n\nfrom search.search_tree import (\n Status,\n SolvedNode,\n UnsolvedNode,\n InvalidNode,\n Edge,\n collect_success_edges,\n)\n\nfrom search.tactic_generator import (\n ModelEmptyOutputError,\n)\n\nfrom search.visual import (\n visualize_proof_tree,\n)\n\ndeepseek = OpenAI(api_key=\"sk-75c0610bc70f43df9a0aadf63e0f1a00\", base_url=\"https://api.deepseek.com\")\n\nasync def evaluate_func(current_state: str, tactic: str, new_state: str) -> float:\n if len(new_state) > 4096:\n logging.warning(\"New state exceeds 4096 characters, returning 0.00\")\n return 0.00\n attempts = 5 # 设置重试次数\n for attempt in range(attempts):\n try:\n response = deepseek.chat.completions.create(\n model=\"deepseek-chat\",\n # messages=[\n # {\n # \"role\": \"system\",\n # \"content\": \"\\\n # You are a professional math assistant specialized in Lean 4 theorem proving. \\\n # The user will provide the current proof state, the tactic applied, and the resulting new state after applying the tactic. \\\n # Your task is to evaluate the effectiveness of the applied tactic in advancing the proof. \\\n # Please provide a score between 0 and 1 with exactly two decimal places. \\\n # The score should follow these rules: \\\n # - 0.00: The tactic is ineffective or invalid. \\\n # - 1.00: The tactic is highly effective and advances the proof significantly. \\\n # - Intermediate values: Indicate varying degrees of effectiveness, rounded to two decimal places. \\\n # The score should be in the format of a number like 0.00, 0.25, 0.50, etc. \\\n # Make sure to output only the score as a number, with no additional explanation or text.\"\n # },\n # {\n # \"role\": \"user\",\n # \"content\": f\"Current State: {current_state}\\nTactic Applied: {tactic}\\nNew State: {new_state}\"\n # }\n # ],\n messages=[\n {\n \"role\": \"system\",\n \"content\": \"\\\n You are a professional math assistant specialized in Lean 4 theorem proving. \\\n The user will provide the current proof state, the tactic applied, and the resulting new state after applying the tactic. \\\n Your task is to evaluate the effectiveness of the applied tactic in advancing the proof. \\\n Please provide a score between 0 and 1 with exactly one decimal places. \\\n The score should follow these rules: \\\n - 0.0: The tactic is ineffective or invalid. \\\n - 1.0: The tactic is highly effective and advances the proof significantly. \\\n - Intermediate values: Indicate varying degrees of effectiveness, rounded to one decimal places. \\\n The score should be in the format of a number like 0.1, 0.5, 0.7, etc. \\\n Make sure to output only the score as a number, with no additional explanation or text.\"\n },\n {\n \"role\": \"user\",\n \"content\": f\"Current State: {current_state}\\nTactic Applied: {tactic}\\nNew State: {new_state}\"\n }\n ],\n stream=False\n )\n\n # print(response.choices[0].message.content)\n\n # score = re.search(r\"effectiveness of this tactic as \\*\\*([0-9.]+)\\*\\*\", response.choices[0].message.content)\n \n # if score:\n # score = score.group(1)\n # print(\"Extracted score:\", score)\n # else:\n # print(\"Score not found.\")\n\n score = float(response.choices[0].message.content)\n\n print(f\"score: {score}\")\n\n return score\n\n except Exception as e:\n # 捕获 OpenAI 的错误和空响应错误\n logging.error(f\"Error occurred: {e}. Attempt {attempt + 1} of {attempts} failed.\")\n if attempt < attempts - 1:\n time.sleep(3) # 等待 3 秒再重新执行\n else:\n logging.error(\"Max attempts reached, exiting.\")\n return 0.1 # 达到最大重试次数后退出\n\ndef log(x):\n if x <= 0:\n return float('-inf')\n return math.log(x)\n\nalpha = 1.00\n\nclass DeepSeekProver(Prover):\n def __init__(\n self,\n actor_id: int,\n num_gpus: int,\n search_timeout: int,\n max_expansions: Optional[int],\n num_sampled_tactics: int,\n result_save_path: str,\n ):\n super().__init__(actor_id, num_gpus, search_timeout, max_expansions, num_sampled_tactics, result_save_path)\n\n async def search(self, thm: TracedTheorem) -> None:\n \"\"\"\n Search for a proof of a theorem using best-first search.\n \"\"\"\n self._prepare()\n await client._test_connection()\n\n self.thm = thm\n \n self.min_beam_size = 4 \n self.max_beam_size = self.num_sampled_tactics\n\n try:\n # initialize the current root node\n self.dojo = Dojo(\n thm.root_dir, \n thm, \n self.leandojo_tactic_timeout,\n self.leandojo_num_threads,\n self.leandojo_memory_limit\n ).__enter__()\n\n current_state = self.dojo.init_state\n\n self.root = UnsolvedNode(\n leandojo_state=current_state,\n is_terminal=False,\n priority=0.0, \n depth=0,\n )\n\n self.nodes = {current_state.id: self.root}\n self.back_edges = []\n self.success_edges = []\n\n try:\n await self._best_first_search()\n except torch.OutOfMemoryError:\n logging.error(f\"🚨OOM when sampling theorem: {self.thm.name}\")\n torch.cuda.empty_cache()\n except DojoCrashError as e:\n logging.error(f\"🚨{e}\")\n \n except (asyncio.TimeoutError, asyncio.CancelledError):\n raise\n\n except DojoInitError as e:\n logging.error(f\"🚨Failed to initialize Dojo: {e}\")\n \n except Exception as e:\n logging.error(f\"🚨{type(e).__name__}: {e}\")\n \n finally: \n if hasattr(self, \"dojo\"):\n self.dojo.__exit__(None, None, None)\n \n async def _best_first_search(self) -> None:\n start_time = time.time()\n priority_queue = asyncio.PriorityQueue() # lowest priority first\n priority_queue.put_nowait((-self.root.priority / ((self.root.depth + 1) ** alpha), self.root))\n\n while True:\n if priority_queue.empty():\n logging.info(\"Search queue is empty.\")\n break\n try:\n await self._step(priority_queue)\n except (asyncio.TimeoutError, asyncio.CancelledError, DojoCrashError):\n raise\n except Exception as e:\n logging.error(f\"🚨{type(e).__name__}: {e}\")\n\n self.elapsed_time = time.time() - start_time\n if (self.elapsed_time > self.search_timeout) or (self.max_expansions and self.num_expansions >= self.max_expansions):\n if self.root.status == Status.SOLVED:\n logging.info(\"Search complete: proof found.\")\n else:\n logging.info(\"Exceeded time limit or max expansions.\")\n break\n\n if self.root.status == Status.INVALID:\n logging.info(\"Invalid tactic generated.\")\n break\n\n if self.root.status == Status.SOLVED:\n logging.info(\"Search complete: proof found.\")\n break\n\n async def _step(self, priority_queue: asyncio.PriorityQueue[Tuple[float, UnsolvedNode]]) -> None:\n _, search_node = priority_queue.get_nowait()\n \n logging.info(f\"Expanding node: {search_node}\")\n\n if isinstance(search_node.leandojo_state, TacticState):\n tactic_state_str = self._build_prompt(search_node)\n else:\n raise ValueError(f\"Invalid leandojo_state for search_node: {type(search_node.leandojo_state)}\")\n \n # generate tactics by LLM\n suggestions = await self._generate_tactics(tactic_state_str, search_node.depth)\n \n # try all the tactics\n results = []\n\n for tactic in suggestions:\n edge, proof_finished = await self._run_tactic(\n search_node, tactic, priority_queue\n )\n \n if edge is not None:\n results.append(edge)\n\n if proof_finished:\n break\n \n # expand the search node\n search_node.out_edges = results\n self.num_expansions += 1\n priority_queue.task_done()\n \n def _count_current_beam_size(\n self,\n depth: int,\n ) -> int:\n \n progress = self.num_expansions / self.max_expansions if self.max_expansions else 0\n factor = max(0, 1 - 15 * progress)\n beam = self.min_beam_size + (self.max_beam_size - self.min_beam_size) * factor\n \n return int(beam)\n\n @torch.no_grad()\n async def _generate_tactics(self, tactic_state_str: str, depth: int) -> List[Tuple[str, float]]:\n start_time = time.time()\n suggestions = await client.async_generate_tactic_sampling(\n self.select_tac_gen(),\n state=tactic_state_str,\n num_samples=self._count_current_beam_size(depth),\n # num_samples=self.num_sampled_tactics,\n )\n\n suggestions = [item['tactic'] for item in suggestions['suggestions']]\n \n if len(suggestions) == 0:\n raise ModelEmptyOutputError(\"No tactic generated.\")\n\n self.model_elapsed_time += time.time() - start_time\n return suggestions\n \n async def _run_tactic(\n self, \n search_node: UnsolvedNode, \n tactic: str, \n priority_queue: asyncio.PriorityQueue\n ) -> Tuple[Optional[Edge], bool]:\n # threshold = log(0.10)\n\n # run tactic in Lean server\n start_time = time.time()\n\n leandojo_new_state = await asyncio.to_thread(\n self.dojo.run_tac, search_node.leandojo_state, tactic\n )\n\n self.dojo_elapsed_time += time.time() - start_time # Tactics triggering timeout will not be counted\n\n # create a new child node\n depth = search_node.depth + 1\n if leandojo_new_state.id not in self.nodes:\n # proof finished\n if isinstance(leandojo_new_state, ProofFinished):\n score = log(1.00)\n child_node = SolvedNode(leandojo_state=leandojo_new_state, depth=depth)\n # invalid tactic\n elif type(leandojo_new_state) in (LeanError, ProofGivenUp):\n score = log(0.00)\n child_node = InvalidNode(leandojo_state=leandojo_new_state, depth=depth)\n # unsolved proof\n else:\n assert isinstance(leandojo_new_state, TacticState), f\"Expected TacticState, got{type(leandojo_new_state)}\"\n score = log(await evaluate_func(\n current_state=search_node.leandojo_state.pp,\n tactic=tactic,\n new_state=leandojo_new_state.pp\n ))\n # if score <= threshold: \n # return None, False # Skip tactics with low scores\n child_node = UnsolvedNode(\n leandojo_state=leandojo_new_state,\n is_terminal=False,\n priority=search_node.priority + score,\n depth=depth,\n )\n priority_queue.put_nowait((-child_node.priority / ((child_node.depth + 1) ** alpha), child_node))\n self.nodes[leandojo_new_state.id] = child_node\n edge = Edge(src=search_node, dst=child_node, tactic=tactic, score=score)\n else:\n assert isinstance(leandojo_new_state, TacticState), f\"Expected TacticState, got{type(leandojo_new_state)}\"\n child_node = self.nodes[leandojo_new_state.id]\n assert isinstance(child_node, UnsolvedNode), f\"Expected UnsolvedNode, got {type(child_node)}\"\n child_node.depth = min(child_node.depth, depth)\n if await asyncio.to_thread(child_node.is_descendant, search_node):\n edge = None\n else:\n score = log(await evaluate_func(\n current_state=search_node.leandojo_state.pp,\n tactic=tactic,\n new_state=leandojo_new_state.pp\n ))\n edge = Edge(src=search_node, dst=child_node, tactic=tactic, score=score)\n self.back_edges.append(edge)\n\n if isinstance(child_node, UnsolvedNode) and edge is not None:\n child_node.in_edges.append(edge)\n elif isinstance(child_node, SolvedNode) or isinstance(child_node, InvalidNode):\n child_node.in_edge = edge\n\n return edge, isinstance(leandojo_new_state, ProofFinished)\n \n def _build_prompt(\n self,\n search_node: UnsolvedNode,\n ) -> str:\n input_template = \"[GOAL]\\n{state}\\n[PROOFSTEP]\\n\"\n return input_template.format(state=search_node.leandojo_state.pp)\n \n def _normalize_scores(\n self,\n suggestions: List[Tuple[str, float]]\n ) -> List[Tuple[str, float, float]]:\n\n exps = [math.exp(score) for _, score in suggestions]\n total = sum(exps)\n \n return [\n (tactic, score, exp_score / total)\n for (tactic, score), exp_score in zip(suggestions, exps)\n ]\n \n async def get_result(self, visualize: bool = True) -> Optional[SearchResult]:\n if not self.thm or not self.root:\n return None\n \n # check if the search was successful\n if self.root.status == Status.SOLVED:\n proof = [e.tactic for e in await asyncio.to_thread(self.root.extract_proof)]\n if visualize:\n self.success_edges = await asyncio.to_thread(collect_success_edges, self.root)\n else:\n proof = None\n \n if visualize:\n await asyncio.to_thread(\n visualize_proof_tree,\n list(self.nodes.values()),\n self.success_edges, \n self.back_edges,\n self.result_save_path + \"/visual\",\n self.thm.name,\n ['simple','detail']\n )\n\n # box the result\n result = SearchResult(\n theorem=self.thm,\n status=self.root.status,\n proof=proof,\n num_total_nodes=len(self.nodes),\n num_expansions=self.num_expansions,\n elapsed_time=self.elapsed_time,\n dojo_elapsed_time=self.dojo_elapsed_time,\n model_elapsed_time=self.model_elapsed_time,\n )\n\n return result"}}}} |