princeton-nlp/tree-of-thought-llm

This is the official code repository for a research paper called "Tree of Thoughts," published at NeurIPS 2023 by researchers at Princeton. The paper introduces a technique for getting AI language models to solve hard problems more reliably by having them explore multiple reasoning paths at once, rather than committing to one answer in a single pass. The core idea is inspired by how humans think through difficult problems: instead of picking the first plausible answer, you might sketch several approaches, evaluate which looks most promising, and continue down that branch while discarding less promising ones. The code implements this using large language models (in practice, GPT-4 via the OpenAI API) as both the idea generator and the evaluator. The model generates candidate "thoughts" (partial solutions or reasoning steps), evaluates how good each one is, and uses a search strategy (breadth-first or depth-first search) to find a complete solution. The repository includes experiments on three specific tasks from the paper: the Game of 24 (a math puzzle using arithmetic), creative writing, and crossword puzzle solving. Shell scripts and a Jupyter notebook reproduce each experiment. Saved logs from the original paper runs are included so you can inspect the model's step-by-step reasoning without re-running anything. Adding a new task to the framework is documented in the README and involves writing two small Python files: one defining the task and one defining the prompts the model should use. The library is also available as a pip package. The code requires an OpenAI API key and Python 3.7 or later. It is licensed under MIT.