Upload 2 files

app.py

@@ -15,7 +15,7 @@ from chat_chains import (
     parse_context_and_question,
     ai_response_format,
 )
-from autoqa_chains import auto_qa_chain, auto_qa_output_parser, followup_qa_chain
+from autoqa_chains import auto_qa_chain
 from chain_of_density import chain_of_density_chain
 from insights_bullet_chain import insights_bullet_chain
 from insights_mind_map_chain import insights_mind_map_chain
@@ -292,16 +292,34 @@ def auto_qa_chain_wrapper(inputs):
         raise InvalidArgumentError("Please provide snippet ids")
     document = "\n\n".join([st.session_state.documents[c].page_content for c in inputs])
     llm = ChatOpenAI(model=st.session_state.model, temperature=0)
+    retriever = st.session_state.retriever
+    formatted_response = ""
     with get_openai_callback() as cb:
-        auto_qa_response = auto_qa_output_parser.invoke(
-            auto_qa_chain(llm).invoke({"paper": document})
-        )["questions"]
-        formated_response = "\n\n".join(
-            f"#### {qa['question']}\n\n{qa['answer']}" for qa in auto_qa_response
-        )
+        auto_qa_response = auto_qa_chain(llm).invoke({"paper": document})
         stats = cb
+        for section in auto_qa_response:
+            section_name = section["section_name"]
+            formatted_response += f"# {section_name}\n"
+            for question in section["questions"]:
+                response = (
+                    qa_chain(ChatOpenAI(model=st.session_state.model, temperature=0))
+                    .invoke(
+                        {
+                            "context": format_docs(
+                                retriever.get_relevant_documents(question)
+                            ),
+                            "question": question,
+                        }
+                    )
+                    .content
+                )
+                answer = parse_model_response(response)["answer"]
+                formatted_response += f"## {question}\n"
+                formatted_response += f"* {answer}\n"
+    formatted_response = "```\n" + formatted_response + "\n```"
+
     st.session_state.messages.append(
-        (f"/auto-insight {inputs}", formated_response, "identity")
+        (f"/auto-insight {inputs}", formatted_response, "identity")
     )
     st.session_state.costing.append(
         {
@@ -311,7 +329,7 @@ def auto_qa_chain_wrapper(inputs):
         }
     )
     return (
-        formated_response,
+        formatted_response,
         "identity",
     )
 

autoqa_chains.py

@@ -1,57 +1,15 @@
-from langchain_core.pydantic_v1 import BaseModel, Field
-from typing import List
 from langchain_core.output_parsers import JsonOutputParser
 from langchain_core.prompts import PromptTemplate
 
-
-class QA(BaseModel):
-    question: str = Field(description="question")
-    answer: str = Field(description="answer")
-
-
-class AutoQA(BaseModel):
-    questions: List[QA] = Field(description="list of question and answers")
-
-
 qa_prompt_template = """
-Generate 10 insightful questions and their corresponding detailed answers about the key aspects of a specific machine learning research paper. 
-The focus should be on the paper's objectives, methodology, key findings, and implications for future research or application. 
-The answers must be based on the content of the research paper, offering clear and comprehensive insights for readers.
-Ensure that the questions cover a broad range of topics related to the paper, including but not limited to the introduction, literature review, \
-methodology, results, discussion, and conclusions. 
-The goal is to capture the essence of the paper in a way that is accessible to an expert audience.
-Your response should be recorded in the following json format: {format_instructions}.
+Create a mind map of questions (based on the given abstract) that will help understand a machine learning research paper.
+Ensure that the outline is structured in the following JSON array for clarity, such that each section should have two keys: "section_name" and "questions"
 
-here is the research paper: ####{paper}####
+Here is the research paper abstract: ####{paper}####
 """
 
-auto_qa_output_parser = JsonOutputParser(pydantic_object=AutoQA)
 qa_prompt = PromptTemplate(
     template=qa_prompt_template,
     input_variables=["paper"],
-    partial_variables={
-        "format_instructions": auto_qa_output_parser.get_format_instructions()
-    },
-)
-auto_qa_chain = lambda model: qa_prompt | model
-
-
-followup_prompt_template = """
-Question: {question}
-Answer: {answer}
-Based on the above question and answer and the research paper as your context, come up with a followup question and its answer.
-The answer should be a bit detailed and strictly based on the research paper.
-Your response should be recorded in the following json format: {format_instructions}.
-
-here is the research paper: ####{paper}####
-"""
-
-followup_prompt = PromptTemplate(
-    template=followup_prompt_template,
-    input_variables=["paper", "question", "answer"],
-    partial_variables={
-        "format_instructions": auto_qa_output_parser.get_format_instructions()
-    },
 )
-
-followup_qa_chain = lambda model: followup_prompt | model
+auto_qa_chain = lambda model: qa_prompt | model | JsonOutputParser()