reasonbench 0.0.1__py3-none-any.whl

reasonbench/__init__.py

@@ -0,0 +1,126 @@
+from typing import TypedDict
+from .typedefs import DecodingParameters
+from .datasets import get_dataset_path
+
+class BenchmarkFactory:
+    registry = {}
+
+    @classmethod
+    def register(cls, benchmark_cls):
+        cls.registry[benchmark_cls.__name__.lower()] = benchmark_cls
+        return benchmark_cls
+
+    @classmethod
+    def get(cls, task: str, *args, **kwargs):
+        key = f"benchmark{task}".lower()
+        try:
+            path = get_dataset_path(task)
+            return cls.registry[key](path=path, *args, **kwargs)
+        except KeyError:
+            raise ValueError(f"No benchmark found for task={task}")
+        
+class EnvironmentFactory:
+    registry = {}
+
+    @classmethod
+    def register(cls, env_cls):
+        cls.registry[env_cls.__name__.lower()] = env_cls
+        return env_cls
+    
+    @classmethod
+    def get(cls, task: str, *args, **kwargs):
+        key = f"environment{task}".lower()
+        try:
+            return cls.registry[key](*args, **kwargs)
+        except KeyError:
+            raise ValueError(f"No environment found for task={task}")
+    
+class AgentFactory:
+    registry = {}
+
+    @classmethod
+    def register(cls, agent_cls):
+        cls.registry[agent_cls.__name__.lower()] = agent_cls
+        return agent_cls
+
+    @classmethod
+    def get(cls, agent_type: str, benchmark: str, *args, **kwargs):
+        key = f"agent{agent_type}{benchmark}".lower()
+        try:
+            return cls.registry[key]#(*args, **kwargs) : Not initialized
+        except KeyError:
+            raise ValueError(f"No agent found for type={agent_type}, benchmark={benchmark}")
+
+class AgentDictFactory:
+    registry = {}
+
+    @classmethod
+    def register(cls, agent_dict_cls):
+        cls.registry[agent_dict_cls.__name__.lower()] = agent_dict_cls
+        return agent_dict_cls
+
+    @classmethod
+    def get(cls, method: str, *args, **kwargs):
+        key = f"agentdict{method}".lower()
+        try:
+            return cls.registry[key](*args, **kwargs)
+        except KeyError:
+            raise ValueError(f"No agent dict found for method={method}")
+
+
+class MethodFactory:
+    registry = {}
+
+    @classmethod
+    def register(cls, method_cls):
+        cls.registry[method_cls.__name__.lower()] = method_cls
+        return method_cls
+
+    @classmethod
+    def get(cls, method: str, benchmark: str, params: DecodingParameters, *args, **kwargs):
+        key = f"method{method}".lower()
+
+        
+        if method == "io":
+            agents = {
+                "step": AgentFactory.get("io", benchmark),
+            }
+        elif method in ["cot", "cot_sc"]:
+            agents = {
+                "step": AgentFactory.get("cot", benchmark),
+            }
+        elif method == "foa":
+            agents = {
+                "step": AgentFactory.get("act", benchmark),
+                "evaluate": AgentFactory.get("evaluate", benchmark),
+            }
+        elif method in ["tot_bfs", "tot_dfs"]:
+            agents = {
+                "step": AgentFactory.get("bfs", benchmark),
+                "evaluate": AgentFactory.get("evaluate", benchmark),
+            }
+        elif method == "got":
+            agents = {
+                "step": AgentFactory.get("act", benchmark),
+                "aggregate": AgentFactory.get("aggregate", benchmark),
+                "evaluate": AgentFactory.get("evaluate", benchmark),
+            }
+        elif method == "rap":
+            agents = {
+                "step": AgentFactory.get("react", benchmark),
+                "evaluate": AgentFactory.get("selfevaluate", benchmark),
+            }
+        elif method == "react":
+            agents = {
+                "step": AgentFactory.get("react", benchmark),
+            }
+        else:
+            raise NotImplementedError(f"Method {method} is not implemented yet.")
+        
+        # For the moment, only supporting same params for all agents
+        agents.update({k+"_params": params for k in agents.keys()})
+
+        try:
+            return cls.registry[key](agents=agents, *args, **kwargs)
+        except KeyError:
+            raise ValueError(f"No method found for name={method}")

reasonbench/datasets.py

@@ -0,0 +1,53 @@
+"""Utility for downloading datasets from HuggingFace Hub."""
+
+import os
+import logging
+from pathlib import Path
+from huggingface_hub import hf_hub_download
+
+logger = logging.getLogger(__name__)
+
+HF_REPO_ID = "potamitisn/ReasonBench"
+LOCAL_DATASETS_DIR = Path("datasets")
+
+# Maps task name to the actual filename on HF Hub
+DATASET_FILES = {
+    "game24": "dataset_game24.csv.gz",
+    "hle": "dataset_hle.jsonl.gz",
+    "hotpotqa": "dataset_hotpotqa.csv.gz",
+    "humaneval": "dataset_humaneval.csv.gz",
+    "logiqa": "dataset_logiqa.csv.gz",
+    "matharena": "dataset_matharena.jsonl.gz",
+    "scibench": "dataset_scibench.csv.gz",
+    "sonnetwriting": "dataset_sonnetwriting.jsonl.gz",
+}
+
+
+def get_dataset_path(task: str, local_dir: str | Path = LOCAL_DATASETS_DIR) -> str:
+    """Return the local path to a dataset file, downloading from HF Hub if needed.
+
+    Args:
+        task: Task name (e.g., "game24", "hle").
+        local_dir: Local directory to store/look for datasets.
+
+    Returns:
+        Path to the local dataset file.
+    """
+    if task not in DATASET_FILES:
+        raise ValueError(f"Unknown task: {task}. Available: {list(DATASET_FILES.keys())}")
+
+    filename = DATASET_FILES[task]
+    local_path = Path(local_dir) / filename
+
+    if local_path.exists():
+        return str(local_path)
+
+    logger.info(f"Dataset for '{task}' not found locally. Downloading from HuggingFace Hub...")
+    downloaded = hf_hub_download(
+        repo_id=HF_REPO_ID,
+        filename=filename,
+        repo_type="dataset",
+        local_dir=str(local_dir),
+    )
+    logger.info(f"Downloaded to {downloaded}")
+    return str(downloaded)

reasonbench/methods/__init__.py

@@ -0,0 +1,9 @@
+from .foa import AgentDictFOA, MethodFOA
+from .tot_bfs import AgentDictTOT, MethodTOT_BFS
+from .tot_dfs import AgentDictTOT, MethodTOT_DFS
+from .got import AgentDictGOT, MethodGOT
+from .rap import MethodRAP, AgentDictRAP
+from .react import AgentDictReact, MethodReact
+from .io import AgentDictIO, MethodIO
+from .cot import AgentDictCOT, MethodCOT
+from .cot_sc import AgentDictCoT, MethodCOT_SC

reasonbench/methods/cot.py

@@ -0,0 +1,55 @@
+import random
+import logging
+import asyncio
+from typing import TypedDict
+from omegaconf import OmegaConf
+from ..typedefs import Method, Model, Agent, Environment, DecodingParameters, State, Benchmark, MAX_SEED
+from .. import MethodFactory, AgentDictFactory
+from ..utils import Resampler
+logger = logging.getLogger(__name__)
+
+@AgentDictFactory.register
+class AgentDictCOT(TypedDict):
+    step: Agent # ActAgent
+    step_params: DecodingParameters
+
+
+@MethodFactory.register
+class MethodCOT(Method):
+    def __init__(self,
+                 model: Model,
+                 agents: AgentDictCOT,
+                 env: Environment,
+                 config: OmegaConf,
+                 n: int = 1
+                ):
+        super().__init__(model, agents, env, config)
+        
+        self.step_agent = agents["step"]
+        self.step_params = agents["step_params"]
+
+        self.n = config.n
+        assert self.n == 1, "CoT has only 1 output"
+
+    async def solve(self, idx: int, state: State, namespace: str, value_cache: dict=None):
+        randomness = idx
+        random.seed(randomness)
+
+        states = [state.clone(randomness=random.randint(0, MAX_SEED)) for _ in range(self.n)]
+
+        action_coroutines = [
+            self.step_agent.act(
+                model=self.model,
+                state=state,
+                n=1,
+                namespace=namespace,
+                request_id=f"idx{idx}-{hash(state)}-agent{i}",
+                params=self.step_params
+            )
+            for i, state in enumerate(states)
+        ]
+        actions = await asyncio.gather(*action_coroutines)
+
+        # Execute the actions
+        states = [self.env.step(state, action[0]) for state, action in zip(states, actions)]
+        return states

reasonbench/methods/cot_sc.py

@@ -0,0 +1,54 @@
+import random
+import logging
+import asyncio
+from typing import TypedDict
+from collections import Counter
+from omegaconf import OmegaConf
+from ..typedefs import Method, Model, Agent, Environment, DecodingParameters, State, Benchmark, MAX_SEED
+from .. import MethodFactory, AgentDictFactory
+from ..utils import Resampler
+logger = logging.getLogger(__name__)
+
+@AgentDictFactory.register
+class AgentDictCoT(TypedDict):
+    step: Agent # ActAgent
+    step_params: DecodingParameters
+
+
+@MethodFactory.register
+class MethodCOT_SC(Method):
+    def __init__(self,
+                 model: Model,
+                 agents: AgentDictCoT,
+                 env: Environment,
+                 config: OmegaConf,
+                ):
+        super().__init__(model, agents, env, config)
+        
+        self.step_agent = agents["step"]
+        self.step_params = agents["step_params"]
+
+        self.n = config.n
+        assert self.n > 1, "CoT-SC needs at least 2 outputs"
+        
+
+    async def solve(self, idx: int, state: State, namespace: str, value_cache: dict=None):
+        randomness = idx
+        random.seed(randomness)
+
+        actions = await self.step_agent.act(
+            model=self.model,
+            state=state,
+            n=self.n,
+            namespace=namespace,
+            request_id=f"idx{idx}-{hash(state)}",
+            params=self.step_params
+        )
+
+        votes = [action for action in actions]
+        counts = Counter(votes)
+        most_common_action = counts.most_common(1)[0][0]
+        state = self.env.step(state, most_common_action)
+        return [state]
+
+    

reasonbench/methods/foa.py

@@ -0,0 +1,127 @@
+import random
+import logging
+import asyncio
+from typing import TypedDict
+from omegaconf import OmegaConf
+from ..typedefs import Method, Model, Agent, Environment, DecodingParameters, State, Benchmark, MAX_SEED
+from .. import MethodFactory, AgentDictFactory
+from ..utils import Resampler
+logger = logging.getLogger(__name__)
+
+@AgentDictFactory.register
+class AgentDictFOA(TypedDict):
+    step: Agent # ActAgent
+    evaluate: Agent # EvaluateAgent
+    step_params: DecodingParameters
+    evaluate_params: DecodingParameters
+
+@MethodFactory.register
+class MethodFOA(Method):
+    def __init__(self, 
+                 agents: AgentDictFOA,
+                 model: Model, 
+                 env: Environment, 
+                 config: OmegaConf
+                 ):
+        super().__init__(model, agents, env, config)
+
+        self.step_agent = agents["step"]
+        self.eval_agent = agents["evaluate"]
+
+        self.step_params = agents["step_params"]
+        self.evaluate_params = agents["evaluate_params"]
+        
+        self.num_agents = config.num_agents
+        self.num_steps = config.num_steps
+        self.k = config.k
+        self.backtrack = config.backtrack
+        self.resampling = config.resampling
+        self.origin = config.origin
+        self.min_steps = config.min_steps
+        self.num_evaluations = config.num_evaluations
+
+
+    async def solve(self, idx: int, state: State, namespace: str, value_cache: dict = None):
+        randomness = idx
+        random.seed(randomness)
+        resampler = Resampler(randomness)
+
+        # Records of previously visited states (state_identifier, state_value, state)
+        visited_states = [("INIT", self.origin, state)]
+        initial_state = state
+
+        # Initialize state for each agent
+        states = [state.clone(randomness=random.randint(0, MAX_SEED)) for _ in range(self.num_agents)]
+
+        solved = False
+        for step in range(self.num_steps):
+
+            if solved:
+                break
+            
+            # Generate actions for each state
+            action_coroutines = [
+                self.step_agent.act(
+                    model=self.model,
+                    state=state,
+                    n=1,
+                    namespace=namespace,
+                    request_id=f"idx{idx}-step{step}-{hash(state)}-agent{i}",
+                    params=self.step_params
+                )
+                for i, state in enumerate(states)
+            ]
+            actions = await asyncio.gather(*action_coroutines)
+
+            # Execute actions
+            states = [self.env.step(state, action[0]) for state, action in zip(states, actions)]
+
+            # Early stop in case any state is solved
+            if any(self.env.evaluate(state)[1] == 1 for state in states):
+                solved = True
+                break
+
+            # Filter previously visited states records
+            remaining_steps = self.num_steps - (step + 1)
+            visited_states = [(identifier, value*self.backtrack, state) for identifier, value, state in visited_states]
+            visited_states = [state for state in visited_states if  remaining_steps >= self.min_steps - len(state[2].steps)]
+
+            # Pruning : Failed = Finished not correctly
+            failed = [i for i, state in enumerate(states) if self.env.is_final(state)]
+            if visited_states != []:
+                replacements, _ = resampler.resample(visited_states.copy(), len(failed), self.resampling)
+            else:
+                replacements, _ = resampler.resample([("", 1, state) for state in states], len(failed), resampling_method="linear")
+            states = [replacements.pop(0) if i in failed else state for i, state in enumerate(states)]
+
+            # Evaluation phase
+            if step < self.num_steps-1 and self.k and step % self.k == 0:
+                
+                # Evaluate the states
+                value_coroutines = [
+                    self.eval_agent.act(
+                        model=self.model,
+                        state=state,
+                        n=self.num_evaluations,
+                        namespace=namespace,
+                        request_id=f"idx{idx}-evaluation{step}-{hash(state)}-agent{i}",
+                        params=self.evaluate_params,
+                        cache=value_cache
+                    )
+                    for i, state in enumerate(states)
+                ]
+                values = await asyncio.gather(*value_coroutines)
+
+                # Update previously visited states records
+                for i, (state, value) in enumerate(zip(states, values)):
+                    if i not in failed:
+                        visited_states.append((f"{i}.{step}", value, state))
+
+                # Resampling
+                states, resampled_idxs = resampler.resample(visited_states, self.num_agents, self.resampling)
+
+        if len(states) == 0:
+            return [initial_state]
+        else:
+            return states
+    

reasonbench/methods/got.py

@@ -0,0 +1,123 @@
+import random
+import asyncio
+import logging
+from typing import TypedDict
+from omegaconf import OmegaConf
+from ..typedefs import Method, Model, Agent, Environment, DecodingParameters, State, Benchmark, MAX_SEED
+from .. import MethodFactory, AgentDictFactory
+logger = logging.getLogger(__name__)
+
+@AgentDictFactory.register
+class AgentDictGOT(TypedDict):
+    step: Agent
+    aggregate: Agent
+    evaluate: Agent
+    step_params: DecodingParameters
+    aggregate_params: DecodingParameters
+    evaluate_params: DecodingParameters
+
+@MethodFactory.register
+class MethodGOT(Method):
+    def __init__(self, 
+                 model: Model, 
+                 agents: AgentDictGOT, 
+                 env: Environment,
+                 config: OmegaConf
+                 ):
+        super().__init__(model, agents, env, config)
+
+        self.step_agent = agents["step"]
+        self.aggregate_agent = agents["aggregate"]
+        self.eval_agent = agents["evaluate"]
+
+        self.step_params = agents["step_params"]
+        self.aggregate_params = agents["aggregate_params"]
+        self.evaluate_params = agents["evaluate_params"]
+
+        self.num_selections = config.num_selections
+        self.num_steps = config.num_steps
+        self.num_generate = config.num_generate
+        self.num_best = config.num_best
+        self.num_evaluations = config.num_evaluations
+
+    async def solve(self, idx: int, state: State, namespace: str, value_cache: dict = None):
+        randomness = idx
+        random.seed(randomness)
+        states = [state.clone(randomness=random.randint(0, MAX_SEED))]
+        logger.debug(f"Solving game: {idx}")
+
+        solved = False
+        for step in range(self.num_steps):
+            if solved:
+                logger.debug(f"Task {idx} solved at step {step - 1}.")
+                break
+
+            logger.debug(f"Step: {step} ({idx})")
+            # Generate actions for each state
+            action_coroutines = [
+                self.step_agent.act(
+                    model=self.model,
+                    state=state,
+                    n=self.num_generate,
+                    namespace=namespace,
+                    request_id=f"idx{idx}-step{step}-{hash(state)}-agent{i}",
+                    params=self.step_params,
+                )
+                for i, state in enumerate(states)
+            ]
+            generated_actions = await asyncio.gather(*action_coroutines)
+            logger.debug(f"{len(generated_actions)} Actions generated for task {idx}; \n {generated_actions}")
+
+            # Aggregate actions
+            aggregate_coroutines = [
+                self.aggregate_agent.act(
+                    model=self.model,
+                    state=state,
+                    actions=action,
+                    k=self.num_selections,
+                    namespace=namespace,
+                    request_id=f"idx{idx}-aggregate{step}-{hash(state)}-agent{i}",
+                    params=self.aggregate_params,
+                )
+                for i, (state, action) in enumerate(zip(states, generated_actions))
+            ]
+
+            actions = await asyncio.gather(*aggregate_coroutines)
+            logger.debug(f"{len(actions)} Actions selected for task {idx}: \n{actions}")
+
+            # Execute actions on environment
+            proposed_states = []
+            for state, actions in zip(states, actions):
+                for action in actions:
+                    proposed_states.append(self.env.step(state, action))
+            
+            if proposed_states == []:
+                break
+            
+            # Early stop in case any state is solved
+            if any(self.env.evaluate(state)[1] == 1 for state in states):
+                solved = True
+            
+            logger.debug(f"Env step for task {idx}: \n{proposed_states}")
+            # Evaluate all proposals
+            value_coroutines = [
+                self.eval_agent.act(
+                    model=self.model,
+                    state=state,
+                    n=self.num_evaluations,
+                    namespace=namespace,
+                    request_id=f"idx{idx}-evaluation{step}-{hash(state)}-agent{i}",
+                    params=self.evaluate_params,
+                    cache=value_cache
+                )
+                for i, state in enumerate(proposed_states)
+            ]
+            values = await asyncio.gather(*value_coroutines)
+            logger.debug(f"Values given for task {idx}: \n{values}")
+
+            # Choose the best states based on their value
+            state_value_pairs = list(zip(proposed_states, values))
+            sorted_pairs = sorted(state_value_pairs, key=lambda x: x[1], reverse=True)
+            states, values = map(list, zip(*sorted_pairs[:self.num_best]))
+        
+        return states

reasonbench/methods/io.py

@@ -0,0 +1,63 @@
+import random
+import logging
+import asyncio
+from typing import TypedDict
+from omegaconf import OmegaConf
+from ..typedefs import Method, Model, Agent, Environment, DecodingParameters, State, Benchmark, MAX_SEED
+from .. import MethodFactory, AgentDictFactory
+from ..utils import Resampler
+logger = logging.getLogger(__name__)
+
+@AgentDictFactory.register
+class AgentDictIO(TypedDict):
+    step: Agent # ActAgent
+    step_params: DecodingParameters
+
+
+@MethodFactory.register
+class MethodIO(Method):
+    def __init__(self,
+                 model: Model,
+                 agents: AgentDictIO,
+                 env: Environment,
+                 config: OmegaConf,
+                ):
+        super().__init__(model, agents, env, config)
+        
+        self.step_agent = agents["step"]
+        self.step_params = agents["step_params"]
+
+        self.n = config.n
+        assert self.n == 1, "IO has only 1 output"
+
+    async def solve(self, idx: int, state: State, namespace: str, value_cache: dict=None):
+        randomness = idx
+        random.seed(randomness)
+
+        states = [state.clone(randomness=random.randint(0, MAX_SEED)) for _ in range(self.n)]
+
+        action_coroutines = [
+            self.step_agent.act(
+                model=self.model,
+                state=state,
+                n=1,
+                namespace=namespace,
+                request_id=f"idx{idx}-{hash(state)}-agent{i}",
+                params=self.step_params
+            )
+            for i, state in enumerate(states)
+        ]
+        actions = await asyncio.gather(*action_coroutines)
+
+        # Execute the actions
+        new_states = []
+
+        for state, action in zip(states, actions):
+            try:
+                new_states.append(self.env.step(state, action[0]))
+            except Exception as e:
+                print(f"Step failed: {e}")
+                new_states.append(state)  # or whatever fallback you want
+
+        states = new_states
+        return states