Dreamer (#16)

* added qa dreamer

* Added CoRE features

---------

Co-authored-by: Sai <[email protected]>

.gitignore

@@ -1,3 +1,4 @@
 *.DS_Store
 .streamlit/secrets.toml
-config.py
+config.py
+.history

app.py

@@ -6,9 +6,10 @@ import os
 import json
 import requests
 from llm_reasoner import LLMReasoner
-from prompts import templates
+from prompts import templates, get_examples
 from typing import Any
 from string import Template
+from qa_dreamer import get_questions
 
 def safe_parse_json(model_answer):
     """.."""
@@ -49,7 +50,7 @@ def check_password():
             del st.session_state["username"]
             return
 
-          # If authentication fails
+        # If authentication fails
         st.session_state["password_correct"] = False
 
     # Return True if the username + password is validated.
@@ -67,6 +68,7 @@ def select_models():
 
     retriever_options = ["Choose one...", "BM25 Retriever", "Off-the-shelf Retriever", "Finetuned Retriever", "No Retriever"]
     reasoner_options = ["Choose one...", "Claude Sonnet", "GPT-4o", "o3-mini"]
+    dreamer_options = ["None", "CoRE", "CoRE-Contrastive", "QA-Decomposition"]
 
     #selectboxes
 
@@ -76,6 +78,12 @@ def select_models():
         key="retriever"
     )
 
+    dreamer = st.selectbox(
+        "Select the DREAMER",
+        dreamer_options,
+        key="dreamer"
+    )
+
     reasoner = st.selectbox(
         "Select the Reasoner Model",
         reasoner_options,
@@ -88,14 +96,14 @@ def select_models():
         # Check that both selections are not the placeholder.
         if retriever == "Choose one..." or reasoner == "Choose one...":
             st.info("Please select both a retriever and a reasoner.")
-            return None, None
+            return None, None, None
         else:
             # Store the valid selections in session state
-            st.session_state["selected_models"] = (retriever, reasoner)
-            return retriever, reasoner
+            st.session_state["selected_models"] = (retriever, dreamer, reasoner)
+            return retriever, dreamer, reasoner
     else:
         st.info("Click 'Next' once you have made your selections.")
-        return None, None
+        return None, None, None
 
 def retriever(query: str, selected_retriever: str):
     """Simulate a 'retriever' step, searching for relevant information."""
@@ -144,14 +152,6 @@ def retriever(query: str, selected_retriever: str):
 
     corpus = '\n\n'.join(topk_documents)
 
-    if "last_expander_id" not in st.session_state:
-        st.session_state.last_expander_id = 0
-
-    # Function to create a unique key for each expander
-    def get_expander_key():
-        st.session_state.last_expander_id += 1
-        return f"expander_{st.session_state.last_expander_id}"
-
     print(f"Topk documents: {topk_documents}")
 
     for idx, document in enumerate(topk_documents):
@@ -197,26 +197,90 @@ def retriever(query: str, selected_retriever: str):
     # You could return retrieved info here.
     return corpus
 
-def reasoner(query: str, documents: list[str], selected_reasoner: str, llm_client: Any):
+def display_to_chat(placeholder, message):
+    """Simulate a stream of messages in the chat."""
+    text = ""
+    import re
+    for chunk in re.split(r'(\s+)', message):
+        text += chunk + " "
+        time.sleep(0.05)
+        # Add a blinking cursor to simulate typing
+        placeholder.markdown(text + "▌")
+    placeholder.markdown(text)
+
+def reasoner(query: str, documents: list[str], selected_dreamer: str, selected_reasoner: str, llm_client: Any):
     """Simulate a 'reasoner' step, thinking about how to answer."""
     with st.chat_message("assistant"):
         placeholder = st.empty()
         text = ""
+        message = ""
+
+        # find the appropriate template
+        if selected_dreamer == 'None':
+            if not documents or len(documents) == 0:
+                prompt_template = Template(templates["no_evidence"])
+                prompt = prompt_template.substitute(claim=query)
+                print(prompt)
+            else:
+                prompt_template = Template(templates["with_evidence"])
+                prompt = prompt_template.substitute(claim=query, corpus_text=documents)
+        elif selected_dreamer == 'CoRE' or selected_dreamer == 'CoRE-Contrastive':
+            conditions = [""]
+            prompt_template = Template(templates["generate_condition"])
+            prompt = prompt_template.substitute(claim=query)
+            prompt = get_examples() + prompt + " answer: " 
+            print(prompt)
+
+            message += f"Using {selected_dreamer} to analyze and verify the claim in detail..."
+            display_to_chat(placeholder, message)
+            placeholder = st.empty()
+
+            llm_response = llm_client.run_inference(prompt)
+            print(llm_response)
+            conditions = llm_response.split('\n\n')
+            print(conditions)
+            
+            condition = conditions[0] if selected_dreamer == 'CoRE' else conditions[1]
+
+            message = "To reason about the claim, CoRE is using the condition: " + condition + "\n\n\n\n"
+
+            if not documents or len(documents) == 0:
+                prompt_template = Template(templates["with_condition"])
+                prompt = prompt_template.substitute(claim=query, condition=condition)
+            else:
+                prompt_template = Template(templates["with_evidence_condition"])
+                prompt = prompt_template.substitute(claim=query, corpus_text=documents, condition=condition)
+        elif selected_dreamer == 'QA-Decomposition':
+            message += "Decomposing into atomic questions...\n"
+            display_to_chat(placeholder, message)
+            placeholder = st.empty()
+            questions = get_questions(query)
+            message = questions + "\n\n\n\n"
+
+            message += "Now using the question decomposition to reason and verify the claim...\n\n"
+
+            if not documents or len(documents) == 0:
+                prompt_template = Template(templates["no_evidence_questions"])
+                prompt = prompt_template.substitute(claim=query, questions=questions)
+                print(prompt)
+            else:
+                prompt_template = Template(templates["with_evidence_questions"])
+                prompt = prompt_template.substitute(claim=query, corpus_text=documents, questions=questions)
+        
+        else:
+            if not documents or len(documents) == 0:
+                prompt_template = Template(templates["no_evidence"])
+                prompt = prompt_template.substitute(claim=query)
+            else:
+                prompt_template = Template(templates["with_evidence"])
+                prompt = prompt_template.substitute(claim=query, corpus_text=documents)
 
         if selected_reasoner == "Claude Sonnet":
-            message = "Using Claude Sonnet to reason and verify the claim..."
+            message += "Using Claude Sonnet to reason and verify the claim..."
         elif selected_reasoner == "GPT-4o":
-            message = "Using GPT-4o to analyze and verify the claim in detail..."
+            message += "Using GPT-4o to analyze and verify the claim in detail..."
         elif selected_reasoner == "o3-mini":
-            message = "Using o3-mini to quickly analyze the claim..."
-
-        if not documents or len(documents) == 0:
-            prompt_template = Template(templates["no_evidence"])
-            prompt = prompt_template.substitute(claim=query)
-            print(prompt)
-        else:
-            prompt_template = Template(templates["with_evidence"])
-            prompt = prompt_template.substitute(claim=query, corpus_text=documents)
+            message += "Using o3-mini to quickly analyze the claim..."
 
         print(prompt)
         
@@ -230,12 +294,7 @@ def reasoner(query: str, documents: list[str], selected_reasoner: str, llm_clien
             print(f"Error with parsing the returned {answer_dict}")
             decision, reasoning = "", ""
 
-        for chunk in message.split():
-            text += chunk + " "
-            time.sleep(0.05)
-            # Add a blinking cursor to simulate typing
-            placeholder.markdown(text + "▌")
-        placeholder.markdown(text)
+        display_to_chat(placeholder, message)
     # You could return reasoning info here.
     return reasoning, decision
 
@@ -247,25 +306,27 @@ def main():
         st.stop()
 
     if "selected_models" not in st.session_state:
-        selected_retriever, selected_reasoner = select_models()
+        selected_retriever, selected_dreamer, selected_reasoner = select_models()
         # If valid selections are returned, store them and reset the change flag.
         if selected_retriever is not None and selected_reasoner is not None:
-            st.session_state.selected_models = (selected_retriever, selected_reasoner)
+            st.session_state.selected_models = (selected_retriever, selected_dreamer, selected_reasoner)
             st.rerun()
         else:
             st.stop()  # Halt further execution until valid selections are made.
     else:
-        selected_retriever, selected_reasoner = st.session_state.selected_models
+        selected_retriever, selected_dreamer, selected_reasoner = st.session_state.selected_models
 
     # START OF AGENTIC DEMO
 
     column1, column2 = st.columns(2)
     column1.caption(f"Retriever Selected: {selected_retriever}")
+    column1.caption(f"Dreamer Selected: {selected_dreamer}")
     column2.caption(f"Reasoner Selected: {selected_reasoner}")
 
     if st.button("Change Selection", key="change_selection_btn"):
         st.session_state.pop("selected_models", None)
         st.session_state.pop("retriever", None)
+        st.session_state.pop("dreamer", None)
         st.session_state.pop("reasoner", None)
         st.session_state.messages = [{"role": "assistant", "content": "Let's start verifying the claims here! 👇"}]
         st.rerun()
@@ -340,7 +401,7 @@ def main():
     if "auto_submit" in st.session_state and st.session_state.auto_submit:
         prompt = st.session_state.auto_submit_text
 
-        display_message = prompt + "  \n" + "  \n" + f"Retriever: {selected_retriever}, Reasoner: {selected_reasoner}"
+        display_message = prompt + "  \n" + "  \n" + f"Retriever: {selected_retriever}, Dreamer: {selected_dreamer}, Reasoner: {selected_reasoner}"
         st.session_state.messages.append({"role": "user", "content": prompt})
         st.session_state.messages.append({"role": "summary", "content": display_message})
 
@@ -349,7 +410,7 @@ def main():
             st.markdown(display_message)
 
         retrieved_documents = retriever(prompt, selected_retriever)
-        reasoning, decision = reasoner(prompt, retrieved_documents, selected_reasoner, llm_client)
+        reasoning, decision = reasoner(prompt, retrieved_documents, selected_dreamer, selected_reasoner, llm_client)
 
         # Display assistant response in chat message container
         with st.chat_message("assistant"):
@@ -360,7 +421,7 @@ def main():
             elif decision.lower() == 'contradict':
                 assistant_response = f'The claim is INCORRECT because {reasoning}'
             else:
-                assistant_response = f'Sorry, the query failed due to an issue with connecting to the LLM service.'
+                assistant_response = 'Sorry, the query failed due to an issue with connecting to the LLM service.'
 
             # Simulate stream of response with milliseconds delay
             for chunk in assistant_response.split():
@@ -381,7 +442,7 @@ def main():
     prompt = st.chat_input("Type your claim here")
     if prompt:
         # Add user message to chat history
-        display_message = prompt + "  \n"+ "  \n"+ f"Retriever: {selected_retriever}, Reasoner: {selected_reasoner}"
+        display_message = prompt + "  \n"+ "  \n"+ f"Retriever: {selected_retriever}, Dreamer: {selected_dreamer}, Reasoner: {selected_reasoner}"
         st.session_state.messages.append({"role": "user", "content": prompt})
         st.session_state.messages.append({"role": "summary", "content": display_message})
         # Display user message in chat message container
@@ -389,7 +450,7 @@ def main():
             st.markdown(display_message)
 
         retrieved_documents = retriever(prompt, selected_retriever)
-        reasoning, decision = reasoner(prompt, retrieved_documents, selected_reasoner, llm_client)
+        reasoning, decision = reasoner(prompt, retrieved_documents, selected_dreamer, selected_reasoner, llm_client)
 
         # Display assistant response in chat message container
         with st.chat_message("assistant"):

prompts.py

@@ -1,3 +1,11 @@
+def get_examples():
+    examples = [('A: Dan got a Sega Genesis.\n\nGenerate a pair of constrastive conditions which can help infer if the context in A is true.\nMake sure the first condition in the pair is supported by the context in A and the second contrastive condition is not supported by the context in A.\nThe conditions should not be stated or paraphrased from the context in A and you are allowed to use your internal knowledge.\nThe conditions should be stative expressions.\nDo not describe activities.\nMake sure the conditions in a pair are constrastive.\nDo not give an explanation.', "Rob was a child in the early 200's. \n\n Rob was a child in the late 80's or early 90's."), ("A:  Kelsey's dog rested on the couch.\n\nGenerate a pair of constrastive conditions which can help infer if the context in A is true.\nMake sure the first condition in the pair is supported by the context in A and the second contrastive condition is not supported by the context in A.\nThe conditions should not be stated or paraphrased from the context in A and you are allowed to use your internal knowledge.\nThe conditions should be stative expressions.\nDo not describe activities.\nMake sure the conditions in a pair are constrastive.\nDo not give an explanation.", 'The couch cover was not durable. \n\n The couch cover was durable.')]
+    ret_str = ""
+    for prompt, value in examples:
+        ret_str += "example: " + prompt + " answer: " + str(value) + " "
+    return ret_str
+
+
 templates = {
 	"no_evidence": (
                 "You are an AI model tasked with verifying claims related to medical and health topics using zero-shot learning. "
@@ -45,5 +53,140 @@ templates = {
                 '  "decision": "SUPPORT or CONTRADICT"\n'
                 "}\n\n"
                 "Now, please evaluate the claim above."
-            )
-}
+            ),
+    
+    "with_evidence_questions":  """\
+You are an AI model tasked with verifying claims related to medical and health topics using zero-shot learning. Your job is to analyze a given claim along with provided supporting evidence (i.e. corpus articles) and decide whether the available evidence and your general medical knowledge would likely SUPPORT or CONTRADICT the claim. You are also given some questions that can help you analyze the claim and evidence.
+
+Claim to evaluate:
+<claim>
+$claim
+</claim>
+
+Additional evidence provided:
+<corpus_text>
+$corpus_text
+</corpus_text>
+
+Questions to consider:
+<questions>
+$questions
+</questions>
+
+Guidelines:
+1. Evaluate the claim's plausibility based on general medical knowledge.
+2. Consider the quality and relevance of the provided evidence.
+3. Analyze the context and determine if the evidence supports or contradicts the claim.
+4. Assess any biases or limitations in the evidence.
+5. Use the provided questions to guide your analysis.
+
+After your analysis, output your response strictly as a JSON object with exactly two keys: "reasoning" and "decision". Do not include any additional commentary or keys.
+
+Example output format:
+{
+  "reasoning": "Your brief explanation here (one or two sentences).",
+  "decision": "SUPPORT or CONTRADICT"
+}
+
+Now, please evaluate the claim above.
+""",
+
+    "no_evidence_questions": """\
+You are an AI model tasked with verifying claims related to medical and health topics using zero-shot learning. Your job is to analyze a given claim and decide whether the available evidence and your general medical knowledge would likely SUPPORT or CONTRADICT the claim. You are also given some questions that can help you analyze the claim and evidence.
+
+Claim to evaluate:
+<claim>
+$claim
+</claim>
+
+Questions to consider:
+<questions>
+$questions
+</questions>
+
+Guidelines:
+1. Evaluate the claim's plausibility based on general medical knowledge.
+2. Consider the specificity and credibility of any numbers or percentages.
+3. Analyze the context and scope of the claim.
+4. Assess any potential biases or limitations.
+5. Use the provided questions to guide your analysis.
+
+After your analysis, output your response strictly as a JSON object with exactly two keys: "reasoning" and "decision". Do not include any additional commentary or keys.
+
+Example output format:
+{
+  "reasoning": "Your brief explanation here (one or two sentences).",
+  "decision": "SUPPORT or CONTRADICT"
+}
+
+Now, please evaluate the claim above.
+""",
+
+	"with_condition": (
+                "You are an AI model tasked with verifying claims related to medical and health topics using zero-shot learning. "
+                "Your job is to analyze a given claim and decide whether the known condition, available evidence and your general medical knowledge would likely SUPPORT or CONTRADICT the claim.\n\n"
+                "Claim to evaluate:\n"
+                "<claim>\n"
+                "$claim\n"
+                "</claim>\n\n"
+                "Known Condition:\n"
+                "<condition>\n"
+                "$condition\n"
+                "</condition>\n\n"
+                "Guidelines:\n"
+                "1. Evaluate the claim's plausibility based on general medical knowledge.\n"
+                "2. Consider the specificity and credibility of any numbers or percentages.\n"
+                "3. Analyze the context and scope of the claim.\n"
+                "4. Examine the provided condition and use it as additional help for assessing the claim's plausibility.\n"
+                "5. Assess any potential biases or limitations.\n\n"
+                "After your analysis, output your response strictly as a JSON object with exactly two keys: \"reasoning\" and \"decision\". "
+                "Do not include any additional commentary or keys.\n\n"
+                "Example output format:\n"
+                "{\n"
+                '  "reasoning": "Your brief explanation here (one or two sentences).",\n'
+                '  "decision": "SUPPORT or CONTRADICT"\n'
+                "}\n\n"
+                "Now, please evaluate the claim above."
+            ),
+	"with_evidence_condition": (
+                "You are an AI model tasked with verifying claims related to medical and health topics using zero-shot learning. "
+                "Your job is to analyze a given claim along with provided supporting evidence (i.e. corpus articles) and decide whether "
+                "the known condition, available evidence and your general medical knowledge would likely SUPPORT or CONTRADICT the claim.\n\n"
+                "Claim to evaluate:\n"
+                "<claim>\n"
+                "$claim\n"
+                "</claim>\n\n"
+                "Additional evidence provided:\n"
+                "<corpus_text>\n"
+                "$corpus_text\n"
+                "</corpus_text>\n\n"
+                "Known Condition:\n"
+                "<condition>\n"
+                "$condition\n"
+                "</condition>\n\n"
+                "Guidelines:\n"
+                "1. Evaluate the claim's plausibility based on general medical knowledge.\n"
+                "2. Consider the quality and relevance of the provided evidence.\n"
+                "3. Analyze the context and determine if the evidence supports or contradicts the claim.\n"
+                "4. Assess any biases or limitations in the evidence.\n"
+                "5. Examine the provided condition and use it as additional help for assessing the claim's plausibility.\n\n"
+                "After your analysis, output your response strictly as a JSON object with exactly two keys: \"reasoning\" and \"decision\". "
+                "Do not include any additional commentary or keys.\n\n"
+                "Example output format:\n"
+                "{\n"
+                '  "reasoning": "Your brief explanation here (one or two sentences).",\n'
+                '  "decision": "SUPPORT or CONTRADICT"\n'
+                "}\n\n"
+                "Now, please evaluate the claim above."
+            ),         
+    "generate_condition": (
+            "A: $claim\n\n"
+            "Generate a pair of constrastive conditions which can help infer if the context in A is true.\n"
+            "Make sure the first condition in the pair is supported by the context in A and the second contrastive condition is not supported by the context in A.\n" 
+            "The conditions should not be stated or paraphrased from the context in A and you are allowed to use your internal knowledge.\n"  
+            "The conditions should be stative expressions.\n"
+            "Do not describe activities.\n"  
+            "Make sure the conditions in a pair are constrastive.\n"
+            "Do not give an explanation."
+        )
+}

qa_dreamer.py

@@ -0,0 +1,62 @@
+from openai import OpenAI
+from tenacity import (
+    retry,
+    stop_after_attempt,
+    wait_random_exponential,
+)
+
+
+PQA_QUESTION2QUESTION_PROMPT = """\
+Given a question, decompose it into multiple atomic questions.
+
+### Instructions:
+- The questions MUST be atomic, i.e., they MUST be answerable by only a single piece of information.
+- The questions MUST be standalone, i.e., they MUST NOT reference any other question or the given question.
+- The questions can be both open-ended or yes/no questions.
+- The questions should be decomposed only from the main question.
+- Each question should be on a new line and start with `**** `.
+
+### Input:
+{text}
+
+### Output:
+"""
+
+qa_client = OpenAI(
+    base_url="http://130.85.37.21:4774/v1",
+    api_key="EMPTY"
+)
+
+@retry(wait=wait_random_exponential(min=1, max=60), stop=stop_after_attempt(7))
+def completion_with_backoff(client, **kwargs):
+    return client.chat.completions.create(**kwargs)
+
+def generation_to_questions(generated_text, header, numbered=True):
+    try:
+        lines = generated_text.split("\n")
+        lines = [line.strip() for line in lines]
+        lines = [line for line in lines if line.startswith(header)]
+        lines = [line.replace(header, "").strip() for line in lines]
+    except:
+        lines = []
+        print("Error in processing generated text")
+    return lines
+
+def get_questions(claim):
+    prompt = PQA_QUESTION2QUESTION_PROMPT.format(text=claim)
+    print(f"Question Generation Prompt: {prompt}")
+    response = completion_with_backoff(
+        client=qa_client,
+        model="dipta007/Llama-3.1-8B-Instruct-finetuned-pqa",
+        messages=[
+            {"role": "user", "content": prompt},
+        ],
+        max_tokens=2048,
+        top_p=1.0,
+        temperature=0.0,
+    )
+    print(f"Questions: {response.choices[0].message}")
+    generation = response.choices[0].message.content
+    questions = generation_to_questions(generation, "****")
+    questions = [f"- {q}" for q in questions]
+    return "\n".join(questions)