treequest 0.3.1__tar.gz → 0.3.2__tar.gz

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: treequest
-Version: 0.3.1
+Version: 0.3.2
 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
@@ -237,6 +237,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
 

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

@@ -201,6 +201,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
 

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

@@ -1,6 +1,6 @@
 [project]
 name = "treequest"
-version = "0.3.1"
+version = "0.3.2"
 description = "A Tree Search Library with Flexible API for LLM Inference-Time Scaling"
 readme = "README.md"
 authors = [

treequest-0.3.2/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.2}/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.1 → treequest-0.3.2}/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.2}/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.2}/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",
-]