langchain/libs/experimental/langchain_experimental/tot/thought_generation.py
Vadim Gubergrits e7e5cb9d08
Tree of Thought introducing a new ToTChain. (#5167)
# [WIP] Tree of Thought introducing a new ToTChain.

This PR adds a new chain called ToTChain that implements the ["Large
Language Model Guided
Tree-of-Though"](https://arxiv.org/pdf/2305.08291.pdf) paper.

There's a notebook example `docs/modules/chains/examples/tot.ipynb` that
shows how to use it.


Implements #4975


## Who can review?

Community members can review the PR once tests pass. Tag
maintainers/contributors who might be interested:

- @hwchase17
- @vowelparrot

---------

Co-authored-by: Vadim Gubergrits <vgubergrits@outbox.com>
Co-authored-by: Harrison Chase <hw.chase.17@gmail.com>
2023-07-26 21:29:39 -07:00

95 lines
3.0 KiB
Python

"""
We provide two strategies for generating thoughts in the Tree of Thoughts (ToT)
framework to avoid repetition:
These strategies ensure that the language model generates diverse and
non-repeating thoughts, which are crucial for problem-solving tasks that require
exploration.
"""
from abc import abstractmethod
from typing import Any, Dict, List, Tuple
from langchain.chains.llm import LLMChain
from langchain.prompts.base import BasePromptTemplate
from pydantic import Field
from langchain_experimental.tot.prompts import COT_PROMPT, PROPOSE_PROMPT
class BaseThoughtGenerationStrategy(LLMChain):
"""
Base class for a thought generation strategy.
"""
c: int = 3
"""The number of children thoughts to propose at each step."""
@abstractmethod
def next_thought(
self,
problem_description: str,
thoughts_path: Tuple[str, ...] = (),
**kwargs: Any
) -> str:
"""
Generate the next thought given the problem description and the thoughts
generated so far.
"""
class SampleCoTStrategy(BaseThoughtGenerationStrategy):
"""
Sample thoughts from a Chain-of-Thought (CoT) prompt.
This strategy works better when the thought space is rich, such as when each
thought is a paragraph. Independent and identically distributed samples
lead to diversity, which helps to avoid repetition.
"""
prompt: BasePromptTemplate = COT_PROMPT
def next_thought(
self,
problem_description: str,
thoughts_path: Tuple[str, ...] = (),
**kwargs: Any
) -> str:
response_text = self.predict_and_parse(
problem_description=problem_description, thoughts=thoughts_path, **kwargs
)
return response_text if isinstance(response_text, str) else ""
class ProposePromptStrategy(BaseThoughtGenerationStrategy):
"""
Propose thoughts sequentially using a "propose prompt".
This strategy works better when the thought space is more constrained, such
as when each thought is just a word or a line. Proposing different thoughts
in the same prompt completion helps to avoid duplication.
"""
prompt: BasePromptTemplate = PROPOSE_PROMPT
tot_memory: Dict[Tuple[str, ...], List[str]] = Field(default_factory=dict)
def next_thought(
self,
problem_description: str,
thoughts_path: Tuple[str, ...] = (),
**kwargs: Any
) -> str:
if thoughts_path not in self.tot_memory or not self.tot_memory[thoughts_path]:
new_thoughts = self.predict_and_parse(
problem_description=problem_description,
thoughts=thoughts_path,
n=self.c,
**kwargs
)
if not new_thoughts:
return ""
if isinstance(new_thoughts, list):
self.tot_memory[thoughts_path] = new_thoughts[::-1]
else:
return ""
return self.tot_memory[thoughts_path].pop()