treequest 0.3.1__tar.gz → 0.3.3__tar.gz

{treequest-0.3.1 → treequest-0.3.3}/PKG-INFO

@@ -1,6 +1,6 @@
-Metadata-Version: 2.4
+Metadata-Version: 2.5
 Name: treequest
-Version: 0.3.1
+Version: 0.3.3
 Summary: A Tree Search Library with Flexible API for LLM Inference-Time Scaling
 Author-email: Kou Misaki <kou.misaki@sakana.ai>, Takuya Akiba <takiba@sakana.ai>
 License-File: LICENSE
@@ -57,24 +57,27 @@ import treequest as tq
 # Each node is associated with a user-definable `state`.
 State = str
 
+
 # 1. Define a function to be used for node generation.
 def generate(parent_state: State | None) -> tuple[State, float]:
     """Generates new states and scores based on the parent state."""
-    if parent_state is None: # None represents the expansion from root.
+    if parent_state is None:  # None represents the expansion from root.
         new_state = "Initial state"
     else:
         new_state = f"State after {parent_state}"
 
-    score = random.random() # A score for the new state; It should be normalized to the [0, 1] range.
+    # A score for the new state; It should be normalized to the [0, 1] range.
+    score = random.random()
     return new_state, score
 
+
 # 2. Instantiate the algorithm and a search tree object.
 algo = tq.ABMCTSA()
 search_tree = algo.init_tree()
 
 # 3. Run the search with a generation budget (10 in this case).
 for _ in range(10):
-    search_tree = algo.step(search_tree, {'Action A': generate})
+    search_tree = algo.step(search_tree, {"Action A": generate})
 
 # 4. Extract the best score and state.
 best_state, best_node_score = tq.top_k(search_tree, algo, k=1)[0]
@@ -82,7 +85,8 @@ print(f"Best state: {best_state}, Score: {best_node_score}")
 
 # 5. Visualize the search tree.
 output_file_basename = Path("ab_mcts_a_search_tree")
-tq.render(search_tree, output_file_basename, format="html")  # Generates `ab_mcts_a_search_tree.html`
+# Generates `ab_mcts_a_search_tree.html`
+tq.render(search_tree, output_file_basename, format="html")
 ```
 
 Alternatively, you can use an ask–tell interface with batched AB-MCTS sampling steps:
@@ -92,8 +96,9 @@ import treequest as tq
 
 State = str
 
-def generate(parent_state: State | None) -> tuple[State, float]:
-    ...
+
+def generate(parent_state: State | None) -> tuple[State, float]: ...
+
 
 generate_fns = {"Action A": generate}
 actions = list(generate_fns.keys())
@@ -160,11 +165,13 @@ import dataclasses
 
 import treequest as tq
 
+
 @dataclasses.dataclass
 class State:
     llm_answer: str
     score: float
 
+
 def generate(parent_state: State | None) -> tuple[State, float]:
     """Generate a new node by calling an LLM."""
     if parent_state is None:
@@ -174,12 +181,14 @@ def generate(parent_state: State | None) -> tuple[State, float]:
 
     return state, state.score
 
+
 def initial_generation() -> State:
     """
     Call LLM API to generate an initial answer.
     """
     ...
 
+
 def refine_answer(llm_answer: str, score: float) -> State:
     """
     Call LLM API to refine an answer.
@@ -190,11 +199,11 @@ def refine_answer(llm_answer: str, score: float) -> State:
 algo = tq.ABMCTSM()
 search_tree = algo.init_tree()
 for i in range(20):
-    search_tree = algo.step(search_tree, {'Action Label': generate})
+    search_tree = algo.step(search_tree, {"Action Label": generate})
     # Logging best node during the search.
     if (i + 1) % 5 == 0:
         best_interim_state, _ = tq.top_k(search_tree, algo, k=1)[0]
-        print(f"Iteration {i+1}: Best state so far = {best_interim_state}")
+        print(f"Iteration {i + 1}: Best state so far = {best_interim_state}")
 
 best_state, _ = tq.top_k(search_tree, algo, k=1)[0]
 print(f"Best Answer: {best_state.llm_answer}, Best Score: {best_state.score}")
@@ -209,6 +218,7 @@ from functools import partial
 
 import treequest as tq
 
+
 def generate(llm_name: str, parent_state=None):
     """
     Call LLM API using litellm, vllm, etc., to generate a new node
@@ -216,9 +226,12 @@ def generate(llm_name: str, parent_state=None):
     ...
     return new_state, new_score
 
+
 llm_names = ["o4-mini", "gemini-2.5-pro"]
 # Create dict of different actions backed by different LLMs.
-generate_fns = {llm_name: partial(generate, llm_name=llm_name) for llm_name in llm_names}
+generate_fns = {
+    llm_name: partial(generate, llm_name=llm_name) for llm_name in llm_names
+}
 
 algo = tq.StandardMCTS()
 search_tree = algo.init_tree()
@@ -237,6 +250,7 @@ The variation is not limited to LLM types; you can use different prompts, action
   - Idempotent: calling `tell` twice on the same `trial_id` does not add extra nodes.
   - For queue-based algorithms, over-told Trials beyond possible number of childs from a parent node (e.g., `(# actions)*samples_per_action` for StandardMCTS) become `INVALID` and are not reflected.
 - Scores are expected to be normalized to the `[0, 1]` range.
+  - `ABMCTSM` exposes an internal `_advanced_config` for advanced knobs (e.g., bypassing the runtime check or tuning Bayesian priors). In most cases, the defaults are sufficient; only change it if you understand the impact.
 
 ## Algorithms
 
@@ -285,10 +299,12 @@ from pathlib import Path
 import treequest as tq
 from PIL import Image
 
+
 class State:
     text: str
     image: Image.Image
 
+
 def state_formatter_html(state: State) -> str:
     """Formats the state for HTML visualization."""
     buffer = BytesIO()
@@ -298,13 +314,15 @@ def state_formatter_html(state: State) -> str:
     # In case state.text cannot be trusted, we recommend e.g., html.escape(state.text) to avoid html injection.
     return f'<p>{state.text}</p><br/><img src="data:image/webp;base64,{img_str}" width=100% />'
 
+
 generate_fns = {
     "action_1": ...  # Define your node generation functions (actions) here
 }
 
 algo = tq.ABMCTSA()
 search_tree = algo.init_tree()
-output_dir = Path("search_tree"); output_dir.mkdir(exist_ok=True)
+output_dir = Path("search_tree")
+output_dir.mkdir(exist_ok=True)
 for i in range(10):
     search_tree = algo.step(search_tree, generate_fns)
     # Rendering "latest.html" for checking search progress on the fly
@@ -313,11 +331,16 @@ for i in range(10):
             search_tree,
             output_basename=output_dir / "latest",
             format="html",
-            state_formatter=state_formatter_html, # Use html custom formetter
+            state_formatter=state_formatter_html,  # Use html custom formetter
         )
 
 tq.render(search_tree, output_basename=output_dir / "final", format="pdf")
-tq.render(search_tree, output_basename=output_dir / "final", format="html", state_formatter=state_formatter_html)
+tq.render(
+    search_tree,
+    output_basename=output_dir / "final",
+    format="html",
+    state_formatter=state_formatter_html,
+)
 ```
 
 > **IMPORTANT:** When using HTML format, ensure that the HTML file is securely handled, especially if the state formatter includes raw HTML content. Avoid opening untrusted HTML files in your browser. For example, XSS (cross site scripting) attacks can occur if the state includes malicious HTML/JavaScript code.

{treequest-0.3.1 → treequest-0.3.3}/README.md

@@ -21,24 +21,27 @@ import treequest as tq
 # Each node is associated with a user-definable `state`.
 State = str
 
+
 # 1. Define a function to be used for node generation.
 def generate(parent_state: State | None) -> tuple[State, float]:
     """Generates new states and scores based on the parent state."""
-    if parent_state is None: # None represents the expansion from root.
+    if parent_state is None:  # None represents the expansion from root.
         new_state = "Initial state"
     else:
         new_state = f"State after {parent_state}"
 
-    score = random.random() # A score for the new state; It should be normalized to the [0, 1] range.
+    # A score for the new state; It should be normalized to the [0, 1] range.
+    score = random.random()
     return new_state, score
 
+
 # 2. Instantiate the algorithm and a search tree object.
 algo = tq.ABMCTSA()
 search_tree = algo.init_tree()
 
 # 3. Run the search with a generation budget (10 in this case).
 for _ in range(10):
-    search_tree = algo.step(search_tree, {'Action A': generate})
+    search_tree = algo.step(search_tree, {"Action A": generate})
 
 # 4. Extract the best score and state.
 best_state, best_node_score = tq.top_k(search_tree, algo, k=1)[0]
@@ -46,7 +49,8 @@ print(f"Best state: {best_state}, Score: {best_node_score}")
 
 # 5. Visualize the search tree.
 output_file_basename = Path("ab_mcts_a_search_tree")
-tq.render(search_tree, output_file_basename, format="html")  # Generates `ab_mcts_a_search_tree.html`
+# Generates `ab_mcts_a_search_tree.html`
+tq.render(search_tree, output_file_basename, format="html")
 ```
 
 Alternatively, you can use an ask–tell interface with batched AB-MCTS sampling steps:
@@ -56,8 +60,9 @@ import treequest as tq
 
 State = str
 
-def generate(parent_state: State | None) -> tuple[State, float]:
-    ...
+
+def generate(parent_state: State | None) -> tuple[State, float]: ...
+
 
 generate_fns = {"Action A": generate}
 actions = list(generate_fns.keys())
@@ -124,11 +129,13 @@ import dataclasses
 
 import treequest as tq
 
+
 @dataclasses.dataclass
 class State:
     llm_answer: str
     score: float
 
+
 def generate(parent_state: State | None) -> tuple[State, float]:
     """Generate a new node by calling an LLM."""
     if parent_state is None:
@@ -138,12 +145,14 @@ def generate(parent_state: State | None) -> tuple[State, float]:
 
     return state, state.score
 
+
 def initial_generation() -> State:
     """
     Call LLM API to generate an initial answer.
     """
     ...
 
+
 def refine_answer(llm_answer: str, score: float) -> State:
     """
     Call LLM API to refine an answer.
@@ -154,11 +163,11 @@ def refine_answer(llm_answer: str, score: float) -> State:
 algo = tq.ABMCTSM()
 search_tree = algo.init_tree()
 for i in range(20):
-    search_tree = algo.step(search_tree, {'Action Label': generate})
+    search_tree = algo.step(search_tree, {"Action Label": generate})
     # Logging best node during the search.
     if (i + 1) % 5 == 0:
         best_interim_state, _ = tq.top_k(search_tree, algo, k=1)[0]
-        print(f"Iteration {i+1}: Best state so far = {best_interim_state}")
+        print(f"Iteration {i + 1}: Best state so far = {best_interim_state}")
 
 best_state, _ = tq.top_k(search_tree, algo, k=1)[0]
 print(f"Best Answer: {best_state.llm_answer}, Best Score: {best_state.score}")
@@ -173,6 +182,7 @@ from functools import partial
 
 import treequest as tq
 
+
 def generate(llm_name: str, parent_state=None):
     """
     Call LLM API using litellm, vllm, etc., to generate a new node
@@ -180,9 +190,12 @@ def generate(llm_name: str, parent_state=None):
     ...
     return new_state, new_score
 
+
 llm_names = ["o4-mini", "gemini-2.5-pro"]
 # Create dict of different actions backed by different LLMs.
-generate_fns = {llm_name: partial(generate, llm_name=llm_name) for llm_name in llm_names}
+generate_fns = {
+    llm_name: partial(generate, llm_name=llm_name) for llm_name in llm_names
+}
 
 algo = tq.StandardMCTS()
 search_tree = algo.init_tree()
@@ -201,6 +214,7 @@ The variation is not limited to LLM types; you can use different prompts, action
   - Idempotent: calling `tell` twice on the same `trial_id` does not add extra nodes.
   - For queue-based algorithms, over-told Trials beyond possible number of childs from a parent node (e.g., `(# actions)*samples_per_action` for StandardMCTS) become `INVALID` and are not reflected.
 - Scores are expected to be normalized to the `[0, 1]` range.
+  - `ABMCTSM` exposes an internal `_advanced_config` for advanced knobs (e.g., bypassing the runtime check or tuning Bayesian priors). In most cases, the defaults are sufficient; only change it if you understand the impact.
 
 ## Algorithms
 
@@ -249,10 +263,12 @@ from pathlib import Path
 import treequest as tq
 from PIL import Image
 
+
 class State:
     text: str
     image: Image.Image
 
+
 def state_formatter_html(state: State) -> str:
     """Formats the state for HTML visualization."""
     buffer = BytesIO()
@@ -262,13 +278,15 @@ def state_formatter_html(state: State) -> str:
     # In case state.text cannot be trusted, we recommend e.g., html.escape(state.text) to avoid html injection.
     return f'<p>{state.text}</p><br/><img src="data:image/webp;base64,{img_str}" width=100% />'
 
+
 generate_fns = {
     "action_1": ...  # Define your node generation functions (actions) here
 }
 
 algo = tq.ABMCTSA()
 search_tree = algo.init_tree()
-output_dir = Path("search_tree"); output_dir.mkdir(exist_ok=True)
+output_dir = Path("search_tree")
+output_dir.mkdir(exist_ok=True)
 for i in range(10):
     search_tree = algo.step(search_tree, generate_fns)
     # Rendering "latest.html" for checking search progress on the fly
@@ -277,11 +295,16 @@ for i in range(10):
             search_tree,
             output_basename=output_dir / "latest",
             format="html",
-            state_formatter=state_formatter_html, # Use html custom formetter
+            state_formatter=state_formatter_html,  # Use html custom formetter
         )
 
 tq.render(search_tree, output_basename=output_dir / "final", format="pdf")
-tq.render(search_tree, output_basename=output_dir / "final", format="html", state_formatter=state_formatter_html)
+tq.render(
+    search_tree,
+    output_basename=output_dir / "final",
+    format="html",
+    state_formatter=state_formatter_html,
+)
 ```
 
 > **IMPORTANT:** When using HTML format, ensure that the HTML file is securely handled, especially if the state formatter includes raw HTML content. Avoid opening untrusted HTML files in your browser. For example, XSS (cross site scripting) attacks can occur if the state includes malicious HTML/JavaScript code.

{treequest-0.3.1 → treequest-0.3.3}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "treequest"
-version = "0.3.1"
+version = "0.3.3"
 description = "A Tree Search Library with Flexible API for LLM Inference-Time Scaling"
 readme = "README.md"
 authors = [
@@ -44,16 +44,22 @@ build-backend = "hatchling.build"
 [dependency-groups]
 dev = [
     "line-profiler>=4.2.0",
-    "mypy>=1.15.0",
+    # Pin ruff and mypy so that their new releases do not change CI results.
+    # Keep ruff in sync with .pre-commit-config.yaml.
+    "mypy==2.4.0",
     "pre-commit>=4.2.0",
     "pytest>=8.3.4",
     "pytest-xdist>=3.6.1",
-    "ruff>=0.9.9",
+    "ruff==0.17.1",
     "tqdm>=4.64.0",
     "types-tqdm>=4.64.0",
     "types-PyYAML>=6.0.0",
 ]
 
+[tool.ruff.lint]
+# The rules Ruff enabled by default before 0.16, which expanded its default rule set.
+select = ["E4", "E7", "E9", "F"]
+
 [tool.hatch.build.targets.wheel]
 packages = ["src/treequest"]
 include = [

treequest-0.3.3/src/treequest/__init__.py

@@ -0,0 +1,40 @@
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    # Expose the real symbol to static type checkers / IDEs even when optional
+    # dependencies for ABMCTSM are not installed at runtime.
+    from .algos.ab_mcts_m.algo import ABMCTSM
+else:
+    from .algos.ab_mcts_m._ab_mcts_m_imports import _import as _ab_mcts_m_import
+
+    if not _ab_mcts_m_import.is_successful():
+        # Create a placeholder that raises an informative error when accessed
+        class _ABMCTSMPlaceholder:
+            def __getattr__(self, name):  # type: ignore
+                _ab_mcts_m_import.check()
+                raise ImportError("ABMCTSM import failed.")
+
+            def __call__(self, *args, **kwargs):  # type: ignore
+                _ab_mcts_m_import.check()
+                raise ImportError("ABMCTSM import failed.")
+
+        ABMCTSM = _ABMCTSMPlaceholder()  # type: ignore
+    else:
+        from .algos.ab_mcts_m.algo import ABMCTSM  # type: ignore[assignment]
+
+from .algos.ab_mcts_a.algo import ABMCTSA
+from .algos.base import Algorithm
+from .algos.standard_mcts import StandardMCTS
+from .algos.tree_of_thought_bfs import TreeOfThoughtsBFSAlgo
+from .ranker import top_k
+from .vis import render
+
+__all__ = [
+    "StandardMCTS",
+    "top_k",
+    "TreeOfThoughtsBFSAlgo",
+    "ABMCTSA",
+    "ABMCTSM",
+    "Algorithm",
+    "render",
+]

{treequest-0.3.1 → treequest-0.3.3}/src/treequest/algos/ab_mcts_a/prob_state.py

@@ -13,6 +13,24 @@ from treequest.algos.tree import Node
 T = TypeVar("T")
 
 
+def _child_index_by_identity(parent: Node, child: Node) -> int:
+    """
+    Return the index of `child` in `parent.children` using identity comparison.
+
+    NOTE: Avoid `list.index()` here because it relies on `__eq__`, and user-defined
+    states may include objects (e.g., numpy arrays / torch tensors) whose equality
+    returns non-bool values and can raise when coerced to `bool`.
+    """
+    for i, n in enumerate(parent.children):
+        if n is child:
+            return i
+    raise ValueError(
+        "Child node not found in parent.children by identity. "
+        f"child.expand_idx={getattr(child, 'expand_idx', None)}, "
+        f"parent.expand_idx={getattr(parent, 'expand_idx', None)}"
+    )
+
+
 @dataclasses.dataclass
 class BetaPrior:
     """
@@ -278,18 +296,15 @@ class NodeProbState:
         assert node.parent is not None
 
         # Update for continuing with existing nodes
+        child_idx = _child_index_by_identity(node.parent, node)
         if self.model_selection_strategy.startswith("multiarm_bandit_"):
             # Update the node's probability
-            self.node_probas["shared"][
-                node.parent.children.index(node)
-            ].tell_observation(reward)
+            self.node_probas["shared"][child_idx].tell_observation(reward)
             # For multiarm bandit strategies, update shared CONT distribution
             self.gen_vs_cont_probas["shared"]["CONT"].tell_observation(reward)
         else:
             # Update the node's probability
-            self.node_probas[action][node.parent.children.index(node)].tell_observation(
-                reward
-            )
+            self.node_probas[action][child_idx].tell_observation(reward)
             # For stack strategy, update action-specific CONT distribution
             self.gen_vs_cont_probas[action]["CONT"].tell_observation(reward)
 
@@ -444,14 +459,11 @@ class NodeProbState:
 
         # Initialize the node's probability distribution with its score
         assert node.parent is not None
+        child_idx = _child_index_by_identity(node.parent, node)
         if model_selection_strategy == "stack":
-            self.node_probas[action][node.parent.children.index(node)].tell_observation(
-                node.score
-            )
+            self.node_probas[action][child_idx].tell_observation(node.score)
         else:
-            self.node_probas["shared"][
-                node.parent.children.index(node)
-            ].tell_observation(node.score)
+            self.node_probas["shared"][child_idx].tell_observation(node.score)
 
     def get_action_for_child(self, child: Node) -> str:
         """

treequest-0.3.3/src/treequest/algos/ab_mcts_m/_ab_mcts_m_imports.py

@@ -0,0 +1,58 @@
+import importlib
+from pathlib import Path
+
+from treequest.imports import try_import
+
+_ARVIZ_IMPORT_ATTEMPTS = 3
+
+
+def _import_arviz() -> None:
+    # TODO: Remove this workaround after arviz fixes the race in its daily warning
+    # https://github.com/arviz-devs/arviz/issues/2530
+    # arviz 0.23.3 and 0.23.4 update the warning's stamp file through a fixed
+    # temporary file, so importing arviz from several processes at once can raise
+    # FileNotFoundError. The process that won has already updated the stamp, so
+    # importing again normally succeeds.
+    for attempt in range(1, _ARVIZ_IMPORT_ATTEMPTS + 1):
+        try:
+            importlib.import_module("arviz")
+            return
+        except FileNotFoundError as e:
+            # The process that lost the race fails to rename the temporary file.
+            lost_race = Path(e.filename or "").name == "daily_warning.tmp"
+            if not lost_race or attempt == _ARVIZ_IMPORT_ATTEMPTS:
+                raise
+            # Leave sys.modules as it is: Python has already dropped the failed
+            # package, the leftover arviz.preview is safe to reuse, and another
+            # thread may be importing arviz by now.
+
+
+with try_import() as _import:
+    import jax
+    from packaging.version import Version
+
+    # TODO: Remove this hotfix after numpyro fixes incompatibility with jax>=0.7
+    # https://github.com/pyro-ppl/numpyro/issues/2051
+    if Version(jax.__version__) >= Version("0.7.0"):
+        import jax.experimental.pjit as _pjit
+        from jax.extend.core.primitives import jit_p
+
+        _pjit.pjit_p = jit_p  # type: ignore
+    import numpy as np
+    import numpyro  # type: ignore
+    import pandas as pd  # type: ignore
+
+    _import_arviz()  # Import arviz before pymc does; see _import_arviz.
+    import pymc as pm  # type: ignore
+    from pymc.sampling.jax import sample_numpyro_nuts  # type: ignore
+    from xarray import DataArray
+
+__all__ = [
+    "jax",
+    "np",
+    "numpyro",
+    "pd",
+    "pm",
+    "sample_numpyro_nuts",
+    "DataArray",
+]

{treequest-0.3.1 → treequest-0.3.3}/src/treequest/algos/ab_mcts_m/algo.py

@@ -47,6 +47,37 @@ class ABMCTSMState(Generic[StateT]):
     trial_store: TrialStore[StateT] = field(default_factory=TrialStore[StateT])
 
 
+@dataclass(frozen=True)
+class ABMCTSMAdvancedConfig:
+    """
+    Advanced/experimental ABMCTSM configuration.
+
+    Most users should rely on the defaults.
+    """
+
+    # If True, validate that rewards/scores are in [0, 1].
+    validate_reward_range: bool = True
+
+    # Prior sigmas for the underlying Bayesian model.
+    prior_mu_alpha_sigma: float = 0.2
+    prior_sigma_alpha_sigma: float = 0.2
+    prior_sigma_y_sigma: float = 0.3
+
+    def __post_init__(self) -> None:
+        if self.prior_mu_alpha_sigma <= 0:
+            raise ValueError(
+                f"prior_mu_alpha_sigma must be > 0, got {self.prior_mu_alpha_sigma}"
+            )
+        if self.prior_sigma_alpha_sigma <= 0:
+            raise ValueError(
+                f"prior_sigma_alpha_sigma must be > 0, got {self.prior_sigma_alpha_sigma}"
+            )
+        if self.prior_sigma_y_sigma <= 0:
+            raise ValueError(
+                f"prior_sigma_y_sigma must be > 0, got {self.prior_sigma_y_sigma}"
+            )
+
+
 class ABMCTSM(Algorithm[StateT, ABMCTSMState[StateT]]):
     """
     Monte Carlo Tree Search algorithm using AB-MCTS-M algorithm.
@@ -63,6 +94,7 @@ class ABMCTSM(Algorithm[StateT, ABMCTSMState[StateT]]):
         min_subtree_size_for_pruning: int = 4,
         same_score_proportion_threshold: float = 0.75,
         max_process_workers: int = os.cpu_count() or 1,
+        _advanced_config: ABMCTSMAdvancedConfig | None = None,
         is_worker: bool = False,
     ):
         """
@@ -79,7 +111,13 @@ class ABMCTSM(Algorithm[StateT, ABMCTSMState[StateT]]):
             min_subtree_size_for_pruning: Pruning Config, see PruningConfig class.
             same_score_proportion_threshold: Pruning Config, see PruningConfig class.
             max_process_workers: Maximum number of parallel processes used for running PyMC sampling.
+            _advanced_config: Advanced/experimental configuration. Defaults are recommended.
         """
+        self._advanced_config = (
+            _advanced_config
+            if _advanced_config is not None
+            else ABMCTSMAdvancedConfig()
+        )
         self.enable_pruning = enable_pruning
         self.pruning_config = PruningConfig(
             min_subtree_size_for_pruning=min_subtree_size_for_pruning,
@@ -94,6 +132,9 @@ class ABMCTSM(Algorithm[StateT, ABMCTSMState[StateT]]):
             enable_pruning=self.enable_pruning,
             pruning_config=self.pruning_config,
             reward_average_priors=self.reward_average_priors,
+            prior_mu_alpha_sigma=self._advanced_config.prior_mu_alpha_sigma,
+            prior_sigma_alpha_sigma=self._advanced_config.prior_sigma_alpha_sigma,
+            prior_sigma_y_sigma=self._advanced_config.prior_sigma_y_sigma,
             model_selection_strategy=self.model_selection_strategy,
         )
 
@@ -195,6 +236,9 @@ class ABMCTSM(Algorithm[StateT, ABMCTSMState[StateT]]):
     def _get_generation_action(
         self, state: ABMCTSMState[StateT], node: Node[StateT], actions: list[str]
     ) -> str:
+        if len(actions) == 1:
+            return actions[0]
+
         observations = Observation.collect_all_observations_of_descendant(
             node, state.all_observations
         )
@@ -255,6 +299,7 @@ class ABMCTSM(Algorithm[StateT, ABMCTSMState[StateT]]):
             min_subtree_size_for_pruning=self.pruning_config.min_subtree_size_for_pruning,
             same_score_proportion_threshold=self.pruning_config.same_score_proportion_threshold,
             max_process_workers=1,
+            _advanced_config=self._advanced_config,
             is_worker=True,
         )
 
@@ -298,7 +343,10 @@ class ABMCTSM(Algorithm[StateT, ABMCTSMState[StateT]]):
 
         # Add new node to the tree
         new_node = state.tree.add_node(
-            result, parent_node, trial_id=finished_trial.trial_id
+            result,
+            parent_node,
+            trial_id=finished_trial.trial_id,
+            validate_score_range=self._advanced_config.validate_reward_range,
         )
 
         # Record observation

{treequest-0.3.1 → treequest-0.3.3}/src/treequest/algos/ab_mcts_m/pymc_interface.py

@@ -148,6 +148,9 @@ class PyMCInterface:
         enable_pruning: bool = True,
         pruning_config: Optional[PruningConfig] = None,
         reward_average_priors: Optional[float | Dict[str, float]] = None,
+        prior_mu_alpha_sigma: float = 0.2,
+        prior_sigma_alpha_sigma: float = 0.2,
+        prior_sigma_y_sigma: float = 0.3,
         model_selection_strategy: str = "multiarm_bandit_thompson",
     ):
         _import.check()
@@ -161,6 +164,22 @@ class PyMCInterface:
         self.reward_average_priors = (
             reward_average_priors if reward_average_priors is not None else dict()
         )
+        if prior_mu_alpha_sigma <= 0:
+            raise ValueError(
+                f"prior_mu_alpha_sigma must be > 0, got {prior_mu_alpha_sigma}"
+            )
+        if prior_sigma_alpha_sigma <= 0:
+            raise ValueError(
+                f"prior_sigma_alpha_sigma must be > 0, got {prior_sigma_alpha_sigma}"
+            )
+        if prior_sigma_y_sigma <= 0:
+            raise ValueError(
+                f"prior_sigma_y_sigma must be > 0, got {prior_sigma_y_sigma}"
+            )
+
+        self.prior_mu_alpha_sigma = prior_mu_alpha_sigma
+        self.prior_sigma_alpha_sigma = prior_sigma_alpha_sigma
+        self.prior_sigma_y_sigma = prior_sigma_y_sigma
 
         # Strategy for model selection:
         # "stack": Perform separate fits for each model (traditional approach)
@@ -491,13 +510,19 @@ class PyMCInterface:
         with pm.Model(coords=coords if not is_prior_model else None) as model:
             # Priors START
             # Overall Goodness of the model itself; mu is set to be 0.5 (50% prob of solving the problem)
-            mu_alpha = pm.Normal("mu_alpha", mu=0.5, sigma=0.2)
+            mu_alpha = pm.Normal(
+                "mu_alpha",
+                mu=self.get_reward_average_prior(action=None),
+                sigma=self.prior_mu_alpha_sigma,
+            )
 
             # expresses the strength of score fluctuation across answers
-            sigma_alpha = pm.HalfNormal("sigma_alpha", sigma=0.2)
+            sigma_alpha = pm.HalfNormal(
+                "sigma_alpha", sigma=self.prior_sigma_alpha_sigma
+            )
 
             # expresses the strength of score fluctuation inside answers
-            sigma_y = pm.HalfNormal("sigma_y", sigma=0.3)
+            sigma_y = pm.HalfNormal("sigma_y", sigma=self.prior_sigma_y_sigma)
             # Priors END
 
             group_dims = "action" if not is_prior_model else None
@@ -543,14 +568,18 @@ class PyMCInterface:
             # Priors START
             # Overall Goodness of the model itself; mu is set to be 0.5 (50% prob of solving the problem)
             mu_alpha = pm.Normal(
-                "mu_alpha", mu=self.get_reward_average_prior(action), sigma=0.2
+                "mu_alpha",
+                mu=self.get_reward_average_prior(action),
+                sigma=self.prior_mu_alpha_sigma,
             )
 
             # expresses the strength of score fluctuation across answers
-            sigma_alpha = pm.HalfNormal("sigma_alpha", sigma=0.2)
+            sigma_alpha = pm.HalfNormal(
+                "sigma_alpha", sigma=self.prior_sigma_alpha_sigma
+            )
 
             # expresses the strength of score fluctuation inside answers
-            sigma_y = pm.HalfNormal("sigma_y", sigma=0.3)
+            sigma_y = pm.HalfNormal("sigma_y", sigma=self.prior_sigma_y_sigma)
             # Priors END
 
             group_dims = "child_idx" if not is_prior_model else None

{treequest-0.3.1 → treequest-0.3.3}/src/treequest/algos/tree.py

@@ -14,6 +14,9 @@ class Node(Generic[StateT]):
     # Scores should be in the range from 0 to 1. For root, we set it to be -1.0 as a placeholder.
     score: float = -1.0
 
+    # If True, validate that score is in [0, 1] for non-root nodes.
+    validate_score_range: bool = dataclasses.field(default=True, kw_only=True)
+
     # Reflects the order of expansion. The root has expand_idx = -1, followed by 0,1,2,...
     expand_idx: int = -1
 
@@ -23,7 +26,9 @@ class Node(Generic[StateT]):
 
     def __post_init__(self) -> None:
         if not self.is_root():
-            if not (self.score <= 1.0 and self.score >= 0.0):
+            if self.validate_score_range and not (
+                self.score <= 1.0 and self.score >= 0.0
+            ):
                 raise RuntimeError(
                     f"The score value should be between 0 and 1, while {self.score} is set."
                 )
@@ -117,11 +122,14 @@ class Tree(Generic[StateT]):
         state_score: StateScoreType[StateT],
         parent: Node[StateT],
         trial_id: Optional[TrialId] = None,
+        *,
+        validate_score_range: bool = True,
     ) -> Node[StateT]:
         state, score = state_score
         node = Node(
             state=state,
             score=score,
+            validate_score_range=validate_score_range,
             parent=parent,
             expand_idx=self.size - 1,
             trial_id=trial_id,

treequest-0.3.1/src/treequest/__init__.py

@@ -1,33 +0,0 @@
-from .algos.ab_mcts_m._ab_mcts_m_imports import _import as _ab_mcts_m_import
-
-if not _ab_mcts_m_import.is_successful():
-    # Create a placeholder that raises an informative error when accessed
-    class _ABMCTSMPlaceholder:
-        def __getattr__(self, name):  # type: ignore
-            _ab_mcts_m_import.check()
-            raise ImportError("ABMCTSM import failed.")
-
-        def __call__(self, *args, **kwargs):  # type: ignore
-            _ab_mcts_m_import.check()
-            raise ImportError("ABMCTSM import failed.")
-
-    ABMCTSM = _ABMCTSMPlaceholder()  # type: ignore
-else:
-    from .algos.ab_mcts_m.algo import ABMCTSM  # type: ignore[assignment]
-
-from .algos.ab_mcts_a.algo import ABMCTSA
-from .algos.base import Algorithm
-from .algos.standard_mcts import StandardMCTS
-from .algos.tree_of_thought_bfs import TreeOfThoughtsBFSAlgo
-from .ranker import top_k
-from .vis import render
-
-__all__ = [
-    "StandardMCTS",
-    "top_k",
-    "TreeOfThoughtsBFSAlgo",
-    "ABMCTSA",
-    "ABMCTSM",
-    "Algorithm",
-    "render",
-]

treequest-0.3.1/src/treequest/algos/ab_mcts_m/_ab_mcts_m_imports.py

@@ -1,29 +0,0 @@
-from treequest.imports import try_import
-
-with try_import() as _import:
-    import jax
-    from packaging.version import Version
-
-    # TODO: Remove this hotfix after numpyro fixes incompatibility with jax>=0.7
-    # https://github.com/pyro-ppl/numpyro/issues/2051
-    if Version(jax.__version__) >= Version("0.7.0"):
-        import jax.experimental.pjit as _pjit
-        from jax.extend.core.primitives import jit_p
-
-        _pjit.pjit_p = jit_p  # type: ignore
-    import numpy as np
-    import numpyro  # type: ignore
-    import pandas as pd  # type: ignore
-    import pymc as pm  # type: ignore
-    from pymc.sampling.jax import sample_numpyro_nuts  # type: ignore
-    from xarray import DataArray
-
-__all__ = [
-    "jax",
-    "np",
-    "numpyro",
-    "pd",
-    "pm",
-    "sample_numpyro_nuts",
-    "DataArray",
-]