Spaces:
Runtime error
Runtime error
"""Chain for summarization with self-verification.""" | |
from pathlib import Path | |
from typing import Dict, List | |
from pydantic import BaseModel, Extra | |
from langchain.chains.base import Chain | |
from langchain.chains.llm import LLMChain | |
from langchain.chains.sequential import SequentialChain | |
from langchain.llms.base import BaseLLM | |
from langchain.prompts.prompt import PromptTemplate | |
PROMPTS_DIR = Path(__file__).parent / "prompts" | |
CREATE_ASSERTIONS_PROMPT = PromptTemplate.from_file( | |
PROMPTS_DIR / "create_facts.txt", ["summary"] | |
) | |
CHECK_ASSERTIONS_PROMPT = PromptTemplate.from_file( | |
PROMPTS_DIR / "check_facts.txt", ["assertions"] | |
) | |
REVISED_SUMMARY_PROMPT = PromptTemplate.from_file( | |
PROMPTS_DIR / "revise_summary.txt", ["checked_assertions", "summary"] | |
) | |
ARE_ALL_TRUE_PROMPT = PromptTemplate.from_file( | |
PROMPTS_DIR / "are_all_true_prompt.txt", ["checked_assertions"] | |
) | |
class LLMSummarizationCheckerChain(Chain, BaseModel): | |
"""Chain for question-answering with self-verification. | |
Example: | |
.. code-block:: python | |
from langchain import OpenAI, LLMSummarizationCheckerChain | |
llm = OpenAI(temperature=0.0) | |
checker_chain = LLMSummarizationCheckerChain(llm=llm) | |
""" | |
llm: BaseLLM | |
"""LLM wrapper to use.""" | |
create_assertions_prompt: PromptTemplate = CREATE_ASSERTIONS_PROMPT | |
check_assertions_prompt: PromptTemplate = CHECK_ASSERTIONS_PROMPT | |
revised_summary_prompt: PromptTemplate = REVISED_SUMMARY_PROMPT | |
are_all_true_prompt: PromptTemplate = ARE_ALL_TRUE_PROMPT | |
input_key: str = "query" #: :meta private: | |
output_key: str = "result" #: :meta private: | |
max_checks: int = 2 | |
"""Maximum number of times to check the assertions. Default to double-checking.""" | |
class Config: | |
"""Configuration for this pydantic object.""" | |
extra = Extra.forbid | |
arbitrary_types_allowed = True | |
def input_keys(self) -> List[str]: | |
"""Return the singular input key. | |
:meta private: | |
""" | |
return [self.input_key] | |
def output_keys(self) -> List[str]: | |
"""Return the singular output key. | |
:meta private: | |
""" | |
return [self.output_key] | |
def _call(self, inputs: Dict[str, str]) -> Dict[str, str]: | |
all_true = False | |
count = 0 | |
output = None | |
original_input = inputs[self.input_key] | |
chain_input = original_input | |
while not all_true and count < self.max_checks: | |
chain = SequentialChain( | |
chains=[ | |
LLMChain( | |
llm=self.llm, | |
prompt=self.create_assertions_prompt, | |
output_key="assertions", | |
verbose=self.verbose, | |
), | |
LLMChain( | |
llm=self.llm, | |
prompt=self.check_assertions_prompt, | |
output_key="checked_assertions", | |
verbose=self.verbose, | |
), | |
LLMChain( | |
llm=self.llm, | |
prompt=self.revised_summary_prompt, | |
output_key="revised_summary", | |
verbose=self.verbose, | |
), | |
LLMChain( | |
llm=self.llm, | |
output_key="all_true", | |
prompt=self.are_all_true_prompt, | |
verbose=self.verbose, | |
), | |
], | |
input_variables=["summary"], | |
output_variables=["all_true", "revised_summary"], | |
verbose=self.verbose, | |
) | |
output = chain({"summary": chain_input}) | |
count += 1 | |
if output["all_true"].strip() == "True": | |
break | |
if self.verbose: | |
print(output["revised_summary"]) | |
chain_input = output["revised_summary"] | |
if not output: | |
raise ValueError("No output from chain") | |
return {self.output_key: output["revised_summary"].strip()} | |
def _chain_type(self) -> str: | |
return "llm_summarization_checker_chain" | |