separate the llms

utils/haystack.py

@@ -33,6 +33,8 @@ def start_haystack(huggingface_token):
     llm.warm_up()
     # start_keyword_pipeline(llm)
     # start_qa_pipeline(llm)
+    keyword_llm = llm
+    answer_llm = llm
     keyword_prompt_template = """
 Your task is to convert the follwing question into 3 keywords that can be used to find relevant medical research papers on PubMed.
 Here is an examples:
@@ -65,10 +67,10 @@ Articles:
     pipe = Pipeline()
 
     pipe.add_component("keyword_prompt_builder", keyword_prompt_builder)
-    pipe.add_component("keyword_llm", llm)
+    pipe.add_component("keyword_llm", keyword_llm)
     pipe.add_component("pubmed_fetcher", fetcher)
     pipe.add_component("prompt_builder", prompt_builder)
-    pipe.add_component("llm", llm)
+    pipe.add_component("llm", answer_llm)
 
     pipe.connect("keyword_prompt_builder.prompt", "keyword_llm.prompt")
     pipe.connect("keyword_llm.replies", "pubmed_fetcher.queries")
@@ -83,7 +85,7 @@ def query(query, _pipeline):
     try:
         result = _pipeline.run(data={"keyword_prompt_builder":{"question":query},
                           "prompt_builder":{"question": query},
-                          "llm":{"generation_kwargs": {"max_new_tokens": 500}}})
+                          "answer_llm":{"generation_kwargs": {"max_new_tokens": 500}}})
     except Exception as e:
         result = ["Please make sure you are providing a correct, public Mastodon account"]
     return result