|
import streamlit as st |
|
from haystack import Pipeline |
|
from utils.pubmed_fetcher import PubMedFetcher |
|
from haystack.components.generators import HuggingFaceTGIGenerator |
|
from haystack.components.builders.prompt_builder import PromptBuilder |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
def start_haystack(huggingface_token): |
|
|
|
keyword_llm = HuggingFaceTGIGenerator("mistralai/Mixtral-8x7B-Instruct-v0.1", token=huggingface_token) |
|
keyword_llm.warm_up() |
|
|
|
llm = HuggingFaceTGIGenerator("mistralai/Mixtral-8x7B-Instruct-v0.1", token=huggingface_token) |
|
llm.warm_up() |
|
|
|
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: |
|
question: "What are the latest treatments for major depressive disorder?" |
|
keywords: |
|
Antidepressive Agents |
|
Depressive Disorder, Major |
|
Treatment-Resistant depression |
|
--- |
|
question: {{ question }} |
|
keywords: |
|
""" |
|
prompt_template = """ |
|
Answer the question truthfully based on the given documents. |
|
If the documents don't contain an answer, use your existing knowledge base. |
|
|
|
q: {{ question }} |
|
Articles: |
|
{% for article in articles %} |
|
{{article.content}} |
|
keywords: {{article.meta['keywords']}} |
|
title: {{article.meta['title']}} |
|
{% endfor %} |
|
|
|
""" |
|
keyword_prompt_builder = PromptBuilder(template=keyword_prompt_template) |
|
prompt_builder = PromptBuilder(template=prompt_template) |
|
fetcher = PubMedFetcher() |
|
|
|
pipe = Pipeline() |
|
pipe.add_component("keyword_prompt_builder", keyword_prompt_builder) |
|
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.connect("keyword_prompt_builder.prompt", "keyword_llm.prompt") |
|
pipe.connect("keyword_llm.replies", "pubmed_fetcher.queries") |
|
|
|
pipe.connect("pubmed_fetcher.articles", "prompt_builder.articles") |
|
pipe.connect("prompt_builder.prompt", "llm.prompt") |
|
return pipe |
|
|
|
|
|
@st.cache_data(show_spinner=True) |
|
def query(query, _pipeline): |
|
try: |
|
replies = _pipeline.run(data={"keyword_prompt_builder":{"question":query}, |
|
"prompt_builder":{"question": query}, |
|
"answer_llm":{"generation_kwargs": {"max_new_tokens": 500}}}) |
|
result = replies['llm']['replies'] |
|
except Exception as e: |
|
result = ["Please make sure you are providing a correct, public Mastodon account"] |
|
return result |
