Home

Awesome

Official Repo of Tree of Thoughts (ToT)

<p> <a href="https://badge.fury.io/py/tree-of-thoughts-llm"> <img src="https://badge.fury.io/py/tree-of-thoughts-llm.svg"> </a> <a href="https://www.python.org/"> <img alt="Build" src="https://img.shields.io/badge/Python-3.7+-1f425f.svg?color=purple"> </a> <a href="https://copyright.princeton.edu/policy"> <img alt="License" src="https://img.shields.io/badge/License-MIT-blue"> </a> <a href="https://zenodo.org/badge/latestdoi/642099326"> <img src="https://zenodo.org/badge/642099326.svg"> </a> </p>

teaser

Official implementation for paper Tree of Thoughts: Deliberate Problem Solving with Large Language Models with code, prompts, model outputs. Also check its tweet thread in 1min.

Setup

  1. Set up OpenAI API key and store in environment variable OPENAI_API_KEY (see here).

  2. Install tot package in two ways:

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

Quick Start

The following minimal script will attempt to solve the game of 24 with 4 5 6 10 (might be a bit slow as it's using GPT-4):

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])

And the output would be something like (note it's not deterministic, and sometimes the output can be wrong):

10 - 4 = 6 (left: 5 6 6)
5 * 6 = 30 (left: 6 30)
30 - 6 = 24 (left: 24)
Answer: (5 * (10 - 4)) - 6 = 24

Paper Experiments

Run experiments via sh scripts/{game24, text, crosswords}/{standard_sampling, cot_sampling, bfs}.sh, except in crosswords we use a DFS algorithm for ToT, which can be run via scripts/crosswords/search_crosswords-dfs.ipynb.

The very simple run.py implements the ToT + BFS algorithm, as well as the naive IO/CoT sampling. Some key arguments:

Paper Trajectories

logs/ contains all the trajectories from the paper's experiments, except for logs/game24/gpt-4_0.7_propose1_value3_greedy5_start900_end1000.json which was reproduced after the paper (as the original experiment was done in a notebook) and achieved a 69% score instead of the original 74% score due to randomness in GPT decoding. We hope to aggregate multiple runs in the future to account for sampling randomness and update the paper, but this shouldn't affect the main conclusions of the paper.

How to Add A New Task

Setting up a new task is easy, and mainly involves two steps.

Citations

Please cite the paper and star this repo if you use ToT and find it interesting/useful, thanks! Feel free to contact shunyuyao.cs@gmail.com or open an issue if you have any questions.

@misc{yao2023tree,
      title={{Tree of Thoughts}: Deliberate Problem Solving with Large Language Models}, 
      author={Shunyu Yao and Dian Yu and Jeffrey Zhao and Izhak Shafran and Thomas L. Griffiths and Yuan Cao and Karthik Narasimhan},
      year={2023},
      eprint={2305.10601},
      archivePrefix={arXiv},
      primaryClass={cs.CL}
}