DOCSEARCHbyAlgolia
Submit Docs

Tree of Thoughts (ToT) Framework

repository·master·Indexed 26 days ago

https://github.com/princeton-nlp/tree-of-thought-llm

Official implementation of the Tree of Thoughts (ToT) framework, enabling Large Language Models to perform deliberate problem-solving through structured thought generation, evaluation, and selection. The library provides tools for implementing search algorithms like BFS and DFS to solve complex tasks such as the Game of 24, creative writing, and mini crosswords. It includes a task registry, prompt wrapping utilities, and integration with OpenAI API backends.

Tokens
3.6K
Snippets
5
Records
23
Agent score
91%

What's inside tree-of-thoughts-llm

  2. Run paper experiments via CLI

    master

    Experiments for different task types can be run using the provided shell scripts located in the scripts/ directory:

    • Game24: sh scripts/game24/{standard_sampling, cot_sampling, bfs}.sh
    • Text: sh scripts/text/{standard_sampling, cot_sampling, bfs}.sh
    • Crosswords: sh scripts/crosswords/{standard_sampling, cot_sampling, bfs}.sh (Note: Crosswords use a DFS algorithm via scripts/crosswords/search_crosswords-dfs.ipynb instead of the standard shell scripts).

    The run.py script implements the ToT + BFS algorithm and supports the following arguments:

  3. Add a new task to Tree of Thoughts

    master

    To implement a new task, follow these two steps:

    1. Define Task Logic: Create a new task class in tot/tasks/ and define the necessary task files in tot/data/. You must also register the task in tot/tasks/__init__.py. Refer to tot/tasks/game24.py for an example.
    2. Define Prompts: Set up task-specific prompts in tot/prompts/. Choose the appropriate --method_generate (sample or propose) and --method_evaluate (value or vote) strategies and provide corresponding prompts. Refer to tot/prompts/game24.py for an example.

  4. Install the tree-of-thoughts-llm package

    master

    You can install the tot package either via PyPI or from the source repository.

    Option 1: Install from PyPI

    pip install tree-of-thoughts-llm

    Option 2: Install from source

    git clone https://github.com/princeton-nlp/tree-of-thought-llm
cd tree-of-thought-llm
pip install -r requirements.txt
pip install -e .
    pip install tree-of-thoughts-llm

  7. Quick Start with Game24Task

    master

    Use the following minimal script to solve the game of 24 using the solve method from tot.methods.bfs and the Game24Task class. Note that the args object (an argparse.Namespace) controls the backend, temperature, and ToT algorithm parameters.

    import argparse
from tot.methods.bfs import solve
from tot.tasks.game24 import Game24Task

args = argparse.Namespace(backend='gpt-4', temperature=0.7, task='game24', naive_run=False, prompt_sample=None, method_generate='propose', method_evaluate='value', method_select='greedy', n_generate_sample=1, n_evaluate_sample=3, n_select_sample=5)

task = Game24Task()
ys, infos = solve(args, task, 900)
print(ys[0])

  9. Reference: run.py command line arguments

    master

    Arguments for the run.py script:

    • --naive_run: If True, run naive IO/CoT sampling instead of ToT + BFS.
    • --prompt_sample: Sampling prompt. Choices: standard, cot.
    • --method_generate: Thought generator. Choices: sample (independent thoughts, e.g., Creative Writing) or propose (sequential thoughts, e.g., Game of 24).
    • --method_evaluate: State evaluator. Choices: value (independent value states, e.g., Game of 24) or vote (voting on states together, e.g., Creative Writing).
    • --n_generate_sample: Number of times to prompt for thought generation.
    • --n_evaluate_sample: Number of times to prompt for state evaluation.
    • --n_select_sample: Number of states to keep from each step (parameter b in the ToT + BFS algorithm).

  10. Implement Depth-First Search (DFS) for crossword solving

    master

    You can implement a DFS algorithm to explore possible crossword solutions. The dfs function requires the following parameters:

    • env: The MiniCrosswordsEnv instance.
    • actions: A list to track the current path of actions.
    • infos: A list to collect results/information from completed paths.
    • time_limit: Maximum number of info entries to collect.
    • prune: Boolean indicating whether to stop exploring if the current state is marked as impossible via env.prompt_status().
    • max_per_state: Limits the number of candidate actions explored at each state to control branching factor.

    Example usage for DFS with pruning:

    # dfs with pruning
infoss = []
for i in range(0, 100, 5):
    env.reset(i)
    infos = []
    actions = []
    dfs(env, actions, infos, 100, prune=True, max_per_state=3)
    infoss.append(infos)
    with open('logs/crosswords/infoss_dfs_prune.json', 'w') as fout:
        json.dump(infoss, fout)

  11. Solve tasks using BFS search via `solve()`

    master

    The solve function implements a Breadth-First Search (BFS) strategy to solve tasks using Tree-of-Thought reasoning. It iterates through a fixed number of steps defined by task.steps, performing generation, evaluation, and selection at each step.

    Arguments:

    • args: A configuration object (typically from CLI arguments) containing:
      • backend: The LLM backend to use.
      • temperature: The sampling temperature.
      • method_generate: Strategy for generating new candidates ('sample' or 'propose').
      • method_evaluate: Strategy for evaluating candidates ('vote' or 'value').
      • method_select: Strategy for selecting the next candidates ('sample' or 'greedy').
      • n_generate_sample: Number of samples to generate per candidate.
      • n_evaluate_sample: Number of samples used for evaluation.
      • n_select_sample: Number of candidates to carry forward to the next step.
      • prompt_sample: The prompt type for sampling ('standard' or 'cot').
    • task: A task object that must implement:
      • steps: Integer number of reasoning steps.
      • get_input(idx): Returns the input for a given index.
      • stops: A list of stop sequences for each step.
      • value_prompt_wrap(x, y): Wraps input and candidate in a value prompt.
      • value_outputs_unwrap(x, y, outputs): Parses LLM outputs into a value.
      • vote_prompt_wrap(x, ys): Wraps input and candidates in a voting prompt.
      • vote_outputs_unwrap(outputs, len_ys): Parses voting outputs.
      • propose_prompt_wrap(x, y): Wraps input and candidate in a proposal prompt.
      • standard_prompt_wrap(x, y): Wraps input and candidate in a standard prompt.
      • cot_prompt_wrap(x, y): Wraps input and candidate in a Chain-of-Thought prompt.
    • idx: The index of the specific task instance to solve.
    • to_print: Boolean flag to enable console logging of the search process.

    Returns:

    • A tuple containing (ys, {'steps': infos}), where ys is the list of final output candidates and infos contains the history of the search process.

  12. Use Game24Task prompt wrappers

    master

    The Game24Task class provides several static methods to wrap prompts for different stages of the Tree of Thoughts process:

    • standard_prompt_wrap(x, y=''): Wraps the standard prompt with input x and optional suffix y.
    • cot_prompt_wrap(x, y=''): Wraps the Chain of Thought prompt with input x and optional suffix y.
    • propose_prompt_wrap(x, y=''): Generates a prompt to propose the next step. If y (the current trajectory) indicates the numbers left are already '24', it returns a CoT prompt; otherwise, it returns a prompt based on the current_numbers extracted from y.
    • value_prompt_wrap(x, y): Generates a prompt to evaluate the current state. If y is the last step (no 'left: ' in the last line), it uses value_last_step_prompt. Otherwise, it uses value_prompt based on the numbers remaining.
    • value_outputs_unwrap(x, y, value_outputs): Processes the model's value evaluations. It maps specific keywords to numerical values: impossible $\rightarrow$ 0.001, likely $\rightarrow$ 1, and sure $\rightarrow$ 20.