File size: 3,369 Bytes
af1cf81
 
041fc49
af1cf81
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
bab66cf
 
 
af1cf81
 
bab66cf
af1cf81
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
970c4fe
af1cf81
 
bab66cf
af1cf81
 
 
 
 
3e8bc94
af1cf81
 
 
 
 
 
3e8bc94
af1cf81
970c4fe
3e8bc94
af1cf81
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
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_keyword_pipeline(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:
# question: "What are the latest treatments for major depressive disorder?"
# keywords:
# Antidepressive Agents
# Depressive Disorder, Major
# Treatment-Resistant depression
# ---
# question: {{ question }}
# keywords:
# """
#     keyword_prompt_builder = PromptBuilder(template=keyword_prompt_template)

#     keyword_pipeline = Pipeline()
#     keyword_pipeline.add_component("keyword_prompt_builder", keyword_prompt_builder)
#     keyword_pipeline.add_component("keyword_llm", llm)
#     return keyword_pipeline

# def start_qa_pipeline(llm):
#     return qa_pipeline

def start_haystack(huggingface_token):
    #Use this function to contruct a pipeline
    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