How to create a biomedical RAG application using Snowflake Arctic for PubMed paper understanding

Question: 

Autophagy is the process where a virus obtains nutrients from it's host, yes or no?

Ground Truth: 

No, autophagy is important in cellular homeostasis for the cell survival mechanism and is involved apoptosis.

Prediction:

No, autophagy is not the process where a virus obtains nutrients from its host. Autophagy is a cellular process that recycles or eliminates cell components and damaged organelles through lysosomal degradation. Some viruses, like Herpes simplex virus type I (HSV-1), can manipulate the autophagic process for their replication and survival, but it is not a means for the virus to obtain nutrients from its host.

Context:

Autophagy is a homeostatic process involved in the turnover or elimination of
cytoplasmic components, damaged organelles, and protein aggregates via a
lysosomal degradation mechanism. Autophagy also provides a mechanism of innate
immunity, known as xenophagy, designed to protect cells from intracellular
pathogens, but it may unfortunately be subverted to act as a pro-viral pathway
facilitating the replication of certain viruses. Herpes simplex virus type I
(HSV-1) is a neurotropic virus that remains latent in host neurons; it is the
most common cause of sporadic viral encephalitis. Moreover, HSV-1 has been
related to the pathogenesis of Alzheimer's disease. HSV-1 can modulate the
autophagic process through a mechanism mediated by the viral protein ICP34.5.
Here we report that HSV-1 induces a strong increase in GFP-LC3 and endogenous
LC3 lipidation, and triggers the accumulation of intracellular autophagic
compartments (mainly autophagosomes) without enhancing autophagic long-lived
protein degradation in the late stages of infection. Autophagy inhibition
mediated by ATG5 gene silencing had no effect on viral growth. The present
results suggest that HSV-1 infection activates the host autophagic machinery and
strongly controls the autophagic process, blocking the fusion of autophagosomes
with lysosomes. These events might be important in the neurodegenerative process
associated with HSV-1 infection. (Score: 0.5636628997173461)

class GenericLLMModel(weave.Model):
    model_name: str = "replicate/snowflake/snowflake-arctic-instruct"
    prompt_template: PromptTemplate
    temperature: float = 0.0
    name: str = "GenericLLMModel"

    def __init__(
        self,
        system_prompt: Optional[str] = None,
        human_prompt: Optional[str] = None,
        model_name: str = "gpt-3.5-turbo",
        temperature: float = 0.0,
    ):
        super().__init__(
            model_name=model_name,
            prompt_template=PromptTemplate(
                system_prompt=system_prompt, human_prompt=human_prompt
            ),
            temperature=temperature,
        )

    @weave.op()
    def predict(
        self,
        human_prompt_args: Optional[dict] = {},
        system_prompt_args: Optional[dict] = {},
    ) -> dict:
        messages = self.prompt_template.format_prompt(
            human_prompt_args=human_prompt_args, system_prompt_args=system_prompt_args
        )
        # ...
        response = completion(**completion_args)
        answer = response.choices[0].message.content
        return {"answer": answer}


question_2_query_model = GenericLLMModel(
    system_prompt=question_2_query_system_prompt,
    human_prompt=question_2_query_human_prompt
)

transformed_query = question_2_query_model.predict(human_prompt_args={"question": question})['answer']

question_2_query_system_prompt = """
### Instruction ###
You are an expert biomedical researcher tasked with converting biomedical questions into optimized semantic search queries. Your goal is to generate queries that will retrieve the most relevant documents from the BioASQ dataset to answer the given question.

### Process ###
Follow these steps to create the semantic search query:
1. Carefully analyze the biomedical question to identify the most important keywords, concepts, and entities
2. Construct a search query using those keywords, aiming to retrieve all potentially relevant documents
3. Optimize the query by incorporating synonyms, related terms, and expanding acronyms if applicable
4. Double check that the query captures the core intent of the question and will match pertinent documents
5. Provide only the final semantic search query in your response, without any additional commentary

### Context ###
The BioASQ dataset consists of biomedical questions along with relevant documents. Your semantic search queries will be used to find the most relevant documents from this dataset to answer each question. The ideal answers have been removed, so your query should focus solely on the question text.

### Examples ###
Question: Is Hirschsprung disease a mendelian or a multifactorial disorder?
Semantic Search Query: Hirschsprung disease AND (mendelian OR multifactorial OR complex) AND (inheritance OR genetics OR genes)

Question: List signaling molecules (ligands) that interact with the receptor EGFR?  
Semantic Search Query: EGFR AND (ligands OR "signaling molecules") AND (EGF OR BTC OR EPR OR HB-EGF OR TGF-α OR AREG OR EPG)

Question: Is the protein Papilin secreted?
Semantic Search Query: Papilin AND (secreted OR extracellular OR "secretory pathway")

### Evaluation ###
Your performance will be evaluated on:  
- Inclusion of the most salient keywords, concepts and entities from the biomedical question
- Appropriate use of synonyms and related terms to improve retrieval
- Ability of the query to capture the full scope and intent of the question
- Overall likelihood of the query retrieving documents that can answer the question
- Adherence to the response format instructions

You MUST provide a well-constructed query that fulfills the given criteria. You will be penalized for queries that are too narrow, off-topic, or poorly formulated.
"""

Question: List signaling molecules (ligands) that interact with the receptor EGFR?  
Semantic Search Query: EGFR AND (ligands OR "signaling molecules") AND (EGF OR BTC OR EPR OR HB-EGF OR TGF-α OR AREG OR EPG)

vector_store = weave.ref('VectorStore:latest').get()
embedding_model = weave.ref('SentenceTransformersModel:latest').get()
vector_store.set_embedding_model(embedding_model)

relevant_docs = vector_store.get_most_relevant_documents(
    query=transformed_query,
    n=5,
    ranking_method="cosine"
)

vector_store = weave.ref('VectorStore:latest').get()
embedding_model = weave.ref('SentenceTransformersModel:latest').get()
qap = weave.ref('QuestionAnswerPairsTrainFiltered:latest').get()

article_relevance_model = GenericLLMModel(
    system_prompt=article_relevance_system_prompt,
    human_prompt=article_relevance_human_prompt
)

for doc in _context:
    doc["relevance"] = article_relevance_model.predict(
        human_prompt_args={
            "question": question,
            "article_text": doc["document"]["passage"]
        }
    )['answer']

relevant_context = [doc for doc in _context if doc["relevance"].lower() == "yes"]

article_relevance_system_prompt = """
### Instruction ###
You are an expert medical researcher librarian. Your task is to determine whether articles from the BioASQ dataset may be relevant to questions from clinicians based on the articles' abstracts. You MUST provide a yes or no answer. You will be penalized for answers that are not a clear yes or no.

### Process ###
1. Carefully read the provided clinical question.
2. Analyze the given article abstract in the context of the question.
3. Determine if the abstract contains information potentially relevant to answering the question.
4. Provide a definitive yes or no answer. Do not hedge or equivocate.

### Evaluation ###
Your performance will be evaluated on:
- Ability to identify abstracts with information relevant to the clinical question
- Providing a clear, unambiguous yes or no answer
- Avoiding reliance on stereotypes or biases in your determination
- Adherence to the required answer format

You MUST provide a yes or no answer. Any other response will be penalized.
"""

summarization_model = GenericLLMModel(
    system_prompt=summarization_system_prompt,
    human_prompt=summarization_human_prompt
)

context_str = "\\\\n\\\\n".join([f"{doc['document']['passage']} (Score: {doc['score']})" for doc in relevant_context])

summary = summarization_model.predict(human_prompt_args={"question": question, "context_str": context_str})['answer']

summarization_system_prompt = """
### Instruction ###
You are an expert medical researcher tasked with summarizing relevant excerpts from biomedical literature to provide background information necessary to answer clinicians' questions. Your summary should be concise yet informative, capturing the key points from the provided context.

### Process ###
1. Read the provided clinical question to understand the information needed.
2. Analyze the given context, which includes excerpts from biomedical literature along with relevance scores.
3. Identify the most pertinent information from the context in relation to the question.
4. Summarize the key points from the relevant excerpts, considering their relevance scores.
5. Synthesize the individual summaries into a coherent overview addressing the question.
6. If the context is not sufficient to answer the question, indicate that more information is needed.

### Format ###
Question: <question>
Summary: <summary_of_relevant_information>
Relevant Excerpts: <excerpts_in_order_of_relevance>

### Evaluation ###
Your performance will be evaluated on:
- Ability to identify and summarize relevant information from the provided context
- Synthesis of individual excerpt summaries into a coherent overview
- Consideration of excerpt relevance scores in the final summary
- Clarity and conciseness of the summary
- Adherence to the specified response format

Provide a summary that directly addresses the given question using the most relevant excerpts from the context. If the provided context is insufficient to answer the question, state "Insufficient information to answer the question."
"""


@weave.op()
def predict(self, question: str, context_str: str) -> str:
    return self.model.predict(human_prompt_args={"question": question, "context_str": context_str})['answer']

synthesis_model = GenericLLMModel(
    system_prompt=synthesis_system_prompt,
    human_prompt=synthesis_human_prompt
)

answer = synthesis_model.predict(human_prompt_args={"question": question, "summary": summary})['answer']

synthesis_system_prompt = """
### Instruction ###
You are an expert medical assistant. Your task is to provide accurate, concise answers to medical questions based on summaries of relevant biomedical literature. Ensure responses are clear, informative, unbiased, and avoid stereotypes. Answer in a natural, human-like manner.

### Process ###
1. Analyze the provided question to understand the key information needed.
2. Review the summary of relevant excerpts from biomedical literature.
3. Identify the most pertinent information in the summary for answering the question.
4. Synthesize the key points into a coherent, concise answer.
5. If the summary lacks sufficient information to conclusively answer the question, state "There is insufficient information provided to conclusively answer the question."

### Format ###
Question: <question>
Answer: <final_answer_based_on_summary>

### Evaluation ###
Your performance will be evaluated on:
- Accuracy and relevance of the answer based on the provided summary
- Clarity and conciseness of the response
- Ability to identify when the summary is insufficient to conclusively answer the question
- Avoidance of bias and stereotyping
- Adherence to the specified format

Provide an answer that directly addresses the question using only the information in the summary. If the summary is insufficient, state that conclusively answering is not possible. Produce the answer in a clear, natural style.
"""


class BioASQAdvancedRAGModel(RAGModel):
    def __init__(self, vector_store, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.vector_store = vector_store

    @weave.op()
    def score_context(self, _context, question) -> str:
        for doc in _context:
            doc["relevance"] = article_relevance_model.predict(
                human_prompt_args={"question": question, "article_text": doc["document"]["passage"]}
            )['answer']

    @weave.op()
    def predict(self, question: str, n_documents: int = 5) -> str:
        # Query transformation
        transformed_query = question_2_query_model.predict(
            human_prompt_args={"question": question}
        )['answer']

        # Document retrieval
        _context = self.vector_store.get_most_relevant_documents(query=transformed_query, n=n_documents)

        # Context scoring
        self.score_context(_context, question)
        relevant_context = [doc for doc in _context if doc["relevance"].lower() == "yes"]
        if len(relevant_context) == 0:
            relevant_context = [_context[0]]

        # Summarization
        context_str = "\\\\n\\\\n".join([f"{doc['document']['passage']} (Score: {doc['score']})" for doc in relevant_context])
        summary = summarization_model.predict(
            human_prompt_args={"question": question, "context_str": context_str}
        )['answer']

        # Final answer synthesis
        answer = synthesis_model.predict(
            human_prompt_args={"question": question, "summary": summary}
        )['answer']

        return {
            "answer": answer,
            "context": [doc["document"]["passage"] for doc in relevant_context],
            "all_context": _context
        }

# Create an instance of our RAG model
rag_model = BioASQAdvancedRAGModel(vector_store=vector_store)

# Load the evaluation dataset
qap = weave.ref('QuestionAnswerPairsTrainFiltered:latest').get()
sub_qap = qap.rows[:10]  # Using first 10 questions for this example

# Define evaluation metrics
from weave_example_demo.scorers.llm_guard_scorer import LLMGuardScorer
from weave_example_demo.scorers.tonic_validate_scorer import TonicValidateScorer

scorers = [
    TonicValidateScorer(
        metrics=[
            "AnswerSimilarityMetric",
            "AugmentationPrecisionMetric",
            "AnswerConsistencyMetric",
        ]
    ),
    LLMGuardScorer(
        metrics=["NoRefusal", "Relevance", "Sensitive"]),
]


import streamlit as st
import weave

# Load the RAG model
@st.cache_resource
def load_rag_model():
    return weave.ref('BioASQAdvancedRAGModel:latest').get()

rag_model = load_rag_model()

st.title("Biomedical Question Answering")

# User input
question = st.text_input("Enter your biomedical question:")

if question:
    with st.spinner("Generating answer..."):
        # Get response from the model
        response = rag_model.predict(question)

    # Display the answer
    st.subheader("Answer:")
    st.write(response['answer'])

    # Display relevant context
    st.subheader("Relevant Context:")
    for context in response['context']:
        st.write(context)


Question:

Which animal bite can cause Capnocytophaga canimorsus infection?

Ground Truth:

Capnocytophaga canimorsus infection is typically associated with dog bites, especially in asplenic or immunocompromised patients, and typically manifest as sepsis and/or bacteremia.

Response:

The animal bite that can cause Capnocytophaga canimorsus infection is from dogs. Capnocytophaga canimorsus is a commensal bacterium found in dogs' mouths, and it can lead to septicemia or meningitis in humans through bites or scratches.

Context:

Capnocytophaga canimorsus, a commensal bacterium from dogs' mouths, can cause
septicemia or meningitis in humans through bites or scratches. Here, we describe
and characterize the inflammatory response of human and mouse macrophages on C.
canimorsus infection. Macrophages infected with 10 different strains failed to
release tumor necrosis factor (TNF)- alpha and interleukin (IL)-1 alpha .
Macrophages infected with live and heat-killed (HK) C. canimorsus 5 (Cc5), a
strain isolated from a patient with fatal septicemia, did not release IL-6,
IL-8, interferon- gamma , macrophage inflammatory protein-1 beta , and nitric
oxide (NO). This absence of a proinflammatory response was characterized by the
inability of Toll-like receptor (TLR) 4 to respond to Cc5. Moreover, live but
not HK Cc5 blocked the release of TNF- alpha and NO induced by HK Yersinia
enterocolitica. In addition, live Cc5 down-regulated the expression of TLR4 and
dephosphorylated p38 mitogen-activated protein kinase. These results highlight
passive and active mechanisms of immune evasion by C. canimorsus, which may
explain its capacity to escape from the host immune system.