Search code examples
pythonhuggingface-transformersbert-language-modelhaystack

How to chain multiple PromptNodes together in a Haystack GenerativeQAPipeline


I'm trying to chain together a simple question answering prompt to an elaboration prompt using Haystack.

I had the following code working just fine:

import os

from haystack.document_stores import InMemoryDocumentStore
from haystack.nodes import BM25Retriever
from haystack.nodes import PromptNode, PromptTemplate, AnswerParser
from haystack.pipelines import Pipeline, TextIndexingPipeline


class Bert:
pipe = None

def __init__(self, data_path):
    print("Initializing model...")
    doc_dir = data_path
    document_store = InMemoryDocumentStore(use_bm25=True)

    files_to_index = [os.path.join(doc_dir, f) for f in os.listdir(doc_dir)]
    indexing_pipeline = TextIndexingPipeline(document_store)
    indexing_pipeline.run_batch(file_paths=files_to_index)

    print("Done indexing")

    retriever = BM25Retriever(document_store=document_store, top_k=2)

    lfqa_prompt = PromptTemplate(
        prompt="""Synthesize a comprehensive answer from the following text for the given 
question.
                                 Provide a clear and concise response that summarizes the key 
points and information presented in the text.
                                 Your answer should be in your own words and be no longer than 
50 words.
                                 \n\n Related text: {join(documents)} \n\n Question: {query} 
\n\n Answer:""",
        output_parser=AnswerParser(),
    )

    prompt_node = PromptNode(model_name_or_path="google/flan-t5-large", 
default_prompt_template=lfqa_prompt)

    elaboration_prompt = PromptTemplate(
        prompt="""Elaborate on the answer to the following question given the related texts.
                                 Provide additional details to the answer in your own words.
                                 The final response should be between 100-200 words.
                                 \n\n Related text: {join(documents)} \n\n Question: {query} 
 \n\n Answer: {prompt_node}""",
        output_parser=AnswerParser(),
    )
    elaboration_node = PromptNode(model_name_or_path="google/flan-t5-large", 
default_prompt_template=elaboration_prompt)

    self.pipe = Pipeline()
    self.pipe.add_node(component=retriever, name="retriever", inputs=["Query"])
    self.pipe.add_node(component=prompt_node, name="prompt_node", inputs=["retriever"])
    #self.pipe.add_node(component=elaboration_node, name="elaboration_node", inputs=["Query", 
"retriever", "prompt_node"])




def generate(self, query):
    prediction = self.pipe.run(query=query)

    return prediction

But when I tried to chain another PromptNode to the end of the lfqa_prompt, I ran into errors. I did some research online and saw that I may need to use Shapers and I edited my code as follows:

import os

from haystack.document_stores import InMemoryDocumentStore
from haystack.nodes import AnswerParser, BM25Retriever, BaseComponent, PromptNode, 
PromptTemplate, Shaper
from haystack.schema import Answer, Document, List
from haystack.pipelines import Pipeline, TextIndexingPipeline


class QAPromptOutputAdapter(BaseComponent):
outgoing_edges = 1

def run(self, **kwargs):
    print(kwargs)
    return {"answers": [Answer(answer=result, type="generative") for result in results]}, 
"output_1"

def run_batch(self):
    pass


class Bert:
pipe = None

def __init__(self, data_path):
    print("Initializing model...")
    doc_dir = data_path
    document_store = InMemoryDocumentStore(use_bm25=True)

    files_to_index = [os.path.join(doc_dir, f) for f in os.listdir(doc_dir)]
    indexing_pipeline = TextIndexingPipeline(document_store)
    indexing_pipeline.run_batch(file_paths=files_to_index)

    print("Done indexing")

    retriever = BM25Retriever(document_store=document_store, top_k=2)

    lfqa_prompt = PromptTemplate(
        prompt="""Synthesize a comprehensive answer from the following text for the given 
question.
                                 Provide a clear and concise response that summarizes the key 
points and information presented in the text.
                                 Your answer should be in your own words and be no longer than 
50 words.
                                 \n\n Related text: {join(documents)} \n\n Question: {query} 
\n\n Answer:""",
        #output_parser=AnswerParser(),
    )

    prompt_node = PromptNode(model_name_or_path="google/flan-t5-large", 
default_prompt_template=lfqa_prompt)

    question_shaper = Shaper(func="value_to_list", inputs={"value": "query", "target_list": 
"documents"},
                             outputs=["questions"])
    answer_shaper = Shaper(func="value_to_list",
                           inputs={"value": "prompt_node.results", 
"target_list": "documents"}, outputs=["answers"])

    elaboration_prompt = PromptTemplate(
        prompt="""Elaborate on the answer to the following question given the related texts.
                                 Provide additional details to the answer in your own words.
                                 The final response should be between 100-200 words.
                                 \n\n Related text: {join(documents)} \n\n Question: 
{questions} \n\n Answer: {outputs}""",
        output_parser=AnswerParser(),
    )
    elaboration_node = PromptNode(model_name_or_path="google/flan-t5-large",
                                  default_prompt_template=elaboration_prompt)

    self.pipe = Pipeline()
    self.pipe.add_node(component=retriever, name="retriever", inputs=["Query"])
    self.pipe.add_node(component=prompt_node, name="prompt_node", inputs=["retriever"])
    self.pipe.add_node(component=question_shaper, name="question_shaper", inputs= 
["prompt_node"])
    self.pipe.add_node(component=answer_shaper, name="answer_shaper", inputs=["prompt_node"])
    self.pipe.add_node(component=elaboration_node, name="elaboration_node",
                       inputs=["question_shaper", "retriever", "answer_shaper"])

def generate(self, query):
    prediction = self.pipe.run(query=query)

    return prediction

Now I just get:

Exception: Exception while running node 'answer_shaper': name 'results' is not defined

Is this the correct solution to chaining two prompt nodes together? Should I be using shapers or am I going about this completely wrong? I'm fairly new to Haystack and generative AI models in general, so help is greatly appreciated.


Solution

  • The answer is supposedly to set the "output_variable" parameter of the PromptNode like this:

    lfqa_node = PromptNode(
        model_name_or_path="google/flan-t5-large", 
        default_prompt_template=lfqa_prompt, 
        output_variable="my_answer"
    )
    

    And then you can use the output like:

    elaboration_prompt = PromptTemplate(
        prompt="""
             ...
             Previous answer: {my_answer} \n\n New answer: 
        """
    )
    

    However, this solution did not seem to work for me, so I simply wrote two separate pipelines, and manually parsed the response from the first pipeline and inputted the answer variable into the second pipeline like this:

    lfqa = self.pipe.run(query=query)
    lfqa_answer = lfqa['results'][0]
    elaboration = self.elaboration_pipeline.run(query=lfqa_answer)