|
|
from __future__ import annotations |
|
|
|
|
|
import re |
|
|
from typing import Any, Dict, List, Optional, Sequence, Tuple |
|
|
|
|
|
import numpy as np |
|
|
from langchain_core.callbacks import ( |
|
|
CallbackManagerForChainRun, |
|
|
) |
|
|
from langchain_core.language_models import BaseLanguageModel |
|
|
from langchain_core.messages import AIMessage |
|
|
from langchain_core.output_parsers import StrOutputParser |
|
|
from langchain_core.prompts import BasePromptTemplate |
|
|
from langchain_core.retrievers import BaseRetriever |
|
|
from langchain_core.runnables import Runnable |
|
|
from pydantic import Field |
|
|
|
|
|
from langchain.chains.base import Chain |
|
|
from langchain.chains.flare.prompts import ( |
|
|
PROMPT, |
|
|
QUESTION_GENERATOR_PROMPT, |
|
|
FinishedOutputParser, |
|
|
) |
|
|
from langchain.chains.llm import LLMChain |
|
|
|
|
|
|
|
|
def _extract_tokens_and_log_probs(response: AIMessage) -> Tuple[List[str], List[float]]: |
|
|
"""Extract tokens and log probabilities from chat model response.""" |
|
|
tokens = [] |
|
|
log_probs = [] |
|
|
for token in response.response_metadata["logprobs"]["content"]: |
|
|
tokens.append(token["token"]) |
|
|
log_probs.append(token["logprob"]) |
|
|
return tokens, log_probs |
|
|
|
|
|
|
|
|
class QuestionGeneratorChain(LLMChain): |
|
|
"""Chain that generates questions from uncertain spans.""" |
|
|
|
|
|
prompt: BasePromptTemplate = QUESTION_GENERATOR_PROMPT |
|
|
"""Prompt template for the chain.""" |
|
|
|
|
|
@classmethod |
|
|
def is_lc_serializable(cls) -> bool: |
|
|
return False |
|
|
|
|
|
@property |
|
|
def input_keys(self) -> List[str]: |
|
|
"""Input keys for the chain.""" |
|
|
return ["user_input", "context", "response"] |
|
|
|
|
|
|
|
|
def _low_confidence_spans( |
|
|
tokens: Sequence[str], |
|
|
log_probs: Sequence[float], |
|
|
min_prob: float, |
|
|
min_token_gap: int, |
|
|
num_pad_tokens: int, |
|
|
) -> List[str]: |
|
|
_low_idx = np.where(np.exp(log_probs) < min_prob)[0] |
|
|
low_idx = [i for i in _low_idx if re.search(r"\w", tokens[i])] |
|
|
if len(low_idx) == 0: |
|
|
return [] |
|
|
spans = [[low_idx[0], low_idx[0] + num_pad_tokens + 1]] |
|
|
for i, idx in enumerate(low_idx[1:]): |
|
|
end = idx + num_pad_tokens + 1 |
|
|
if idx - low_idx[i] < min_token_gap: |
|
|
spans[-1][1] = end |
|
|
else: |
|
|
spans.append([idx, end]) |
|
|
return ["".join(tokens[start:end]) for start, end in spans] |
|
|
|
|
|
|
|
|
class FlareChain(Chain): |
|
|
"""Chain that combines a retriever, a question generator, |
|
|
and a response generator. |
|
|
|
|
|
See [Active Retrieval Augmented Generation](https://arxiv.org/abs/2305.06983) paper. |
|
|
""" |
|
|
|
|
|
question_generator_chain: Runnable |
|
|
"""Chain that generates questions from uncertain spans.""" |
|
|
response_chain: Runnable |
|
|
"""Chain that generates responses from user input and context.""" |
|
|
output_parser: FinishedOutputParser = Field(default_factory=FinishedOutputParser) |
|
|
"""Parser that determines whether the chain is finished.""" |
|
|
retriever: BaseRetriever |
|
|
"""Retriever that retrieves relevant documents from a user input.""" |
|
|
min_prob: float = 0.2 |
|
|
"""Minimum probability for a token to be considered low confidence.""" |
|
|
min_token_gap: int = 5 |
|
|
"""Minimum number of tokens between two low confidence spans.""" |
|
|
num_pad_tokens: int = 2 |
|
|
"""Number of tokens to pad around a low confidence span.""" |
|
|
max_iter: int = 10 |
|
|
"""Maximum number of iterations.""" |
|
|
start_with_retrieval: bool = True |
|
|
"""Whether to start with retrieval.""" |
|
|
|
|
|
@property |
|
|
def input_keys(self) -> List[str]: |
|
|
"""Input keys for the chain.""" |
|
|
return ["user_input"] |
|
|
|
|
|
@property |
|
|
def output_keys(self) -> List[str]: |
|
|
"""Output keys for the chain.""" |
|
|
return ["response"] |
|
|
|
|
|
def _do_generation( |
|
|
self, |
|
|
questions: List[str], |
|
|
user_input: str, |
|
|
response: str, |
|
|
_run_manager: CallbackManagerForChainRun, |
|
|
) -> Tuple[str, bool]: |
|
|
callbacks = _run_manager.get_child() |
|
|
docs = [] |
|
|
for question in questions: |
|
|
docs.extend(self.retriever.invoke(question)) |
|
|
context = "\n\n".join(d.page_content for d in docs) |
|
|
result = self.response_chain.invoke( |
|
|
{ |
|
|
"user_input": user_input, |
|
|
"context": context, |
|
|
"response": response, |
|
|
}, |
|
|
{"callbacks": callbacks}, |
|
|
) |
|
|
if isinstance(result, AIMessage): |
|
|
result = result.content |
|
|
marginal, finished = self.output_parser.parse(result) |
|
|
return marginal, finished |
|
|
|
|
|
def _do_retrieval( |
|
|
self, |
|
|
low_confidence_spans: List[str], |
|
|
_run_manager: CallbackManagerForChainRun, |
|
|
user_input: str, |
|
|
response: str, |
|
|
initial_response: str, |
|
|
) -> Tuple[str, bool]: |
|
|
question_gen_inputs = [ |
|
|
{ |
|
|
"user_input": user_input, |
|
|
"current_response": initial_response, |
|
|
"uncertain_span": span, |
|
|
} |
|
|
for span in low_confidence_spans |
|
|
] |
|
|
callbacks = _run_manager.get_child() |
|
|
if isinstance(self.question_generator_chain, LLMChain): |
|
|
question_gen_outputs = self.question_generator_chain.apply( |
|
|
question_gen_inputs, callbacks=callbacks |
|
|
) |
|
|
questions = [ |
|
|
output[self.question_generator_chain.output_keys[0]] |
|
|
for output in question_gen_outputs |
|
|
] |
|
|
else: |
|
|
questions = self.question_generator_chain.batch( |
|
|
question_gen_inputs, config={"callbacks": callbacks} |
|
|
) |
|
|
_run_manager.on_text( |
|
|
f"Generated Questions: {questions}", color="yellow", end="\n" |
|
|
) |
|
|
return self._do_generation(questions, user_input, response, _run_manager) |
|
|
|
|
|
def _call( |
|
|
self, |
|
|
inputs: Dict[str, Any], |
|
|
run_manager: Optional[CallbackManagerForChainRun] = None, |
|
|
) -> Dict[str, Any]: |
|
|
_run_manager = run_manager or CallbackManagerForChainRun.get_noop_manager() |
|
|
|
|
|
user_input = inputs[self.input_keys[0]] |
|
|
|
|
|
response = "" |
|
|
|
|
|
for i in range(self.max_iter): |
|
|
_run_manager.on_text( |
|
|
f"Current Response: {response}", color="blue", end="\n" |
|
|
) |
|
|
_input = {"user_input": user_input, "context": "", "response": response} |
|
|
tokens, log_probs = _extract_tokens_and_log_probs( |
|
|
self.response_chain.invoke( |
|
|
_input, {"callbacks": _run_manager.get_child()} |
|
|
) |
|
|
) |
|
|
low_confidence_spans = _low_confidence_spans( |
|
|
tokens, |
|
|
log_probs, |
|
|
self.min_prob, |
|
|
self.min_token_gap, |
|
|
self.num_pad_tokens, |
|
|
) |
|
|
initial_response = response.strip() + " " + "".join(tokens) |
|
|
if not low_confidence_spans: |
|
|
response = initial_response |
|
|
final_response, finished = self.output_parser.parse(response) |
|
|
if finished: |
|
|
return {self.output_keys[0]: final_response} |
|
|
continue |
|
|
|
|
|
marginal, finished = self._do_retrieval( |
|
|
low_confidence_spans, |
|
|
_run_manager, |
|
|
user_input, |
|
|
response, |
|
|
initial_response, |
|
|
) |
|
|
response = response.strip() + " " + marginal |
|
|
if finished: |
|
|
break |
|
|
return {self.output_keys[0]: response} |
|
|
|
|
|
@classmethod |
|
|
def from_llm( |
|
|
cls, llm: BaseLanguageModel, max_generation_len: int = 32, **kwargs: Any |
|
|
) -> FlareChain: |
|
|
"""Creates a FlareChain from a language model. |
|
|
|
|
|
Args: |
|
|
llm: Language model to use. |
|
|
max_generation_len: Maximum length of the generated response. |
|
|
kwargs: Additional arguments to pass to the constructor. |
|
|
|
|
|
Returns: |
|
|
FlareChain class with the given language model. |
|
|
""" |
|
|
try: |
|
|
from langchain_openai import ChatOpenAI |
|
|
except ImportError: |
|
|
raise ImportError( |
|
|
"OpenAI is required for FlareChain. " |
|
|
"Please install langchain-openai." |
|
|
"pip install langchain-openai" |
|
|
) |
|
|
llm = ChatOpenAI(max_tokens=max_generation_len, logprobs=True, temperature=0) |
|
|
response_chain = PROMPT | llm |
|
|
question_gen_chain = QUESTION_GENERATOR_PROMPT | llm | StrOutputParser() |
|
|
return cls( |
|
|
question_generator_chain=question_gen_chain, |
|
|
response_chain=response_chain, |
|
|
**kwargs, |
|
|
) |
|
|
|
