section summary

app.py

@@ -21,6 +21,7 @@ from insights_bullet_chain import insights_bullet_chain
 from synopsis_chain import synopsis_chain
 from custom_exceptions import InvalidArgumentError, InvalidCommandError
 from openai_configuration import openai_parser
+from summary_chain import summary_chain
 
 st.set_page_config(layout="wide")
 
@@ -44,6 +45,7 @@ Here's a quick guide to getting started with me:
 | `/insight-bullets <list of snippet ids>` | Extract and summarize key insights, methods, results, and conclusions. |
 | `/paper-synopsis <list of snippet ids>` | Generate a synopsis of the paper. |
 | `/deep-dive [<list of snippet ids>] <query>` | Query me with a specific context. |
+| `/summarise-section [<list of snippet ids>] <section name>` | Summarize a specific section of the paper. |
 
 
 <br>
@@ -182,10 +184,29 @@ def rag_llm_wrapper(inputs):
 def query_llm_wrapper(inputs):
     context, question = parse_context_and_question(inputs)
     relevant_docs = [st.session_state.documents[c] for c in context]
-    print(context, question)
     return query_llm(question, relevant_docs)
 
 
+def summarise_wrapper(inputs):
+    context, query = parse_context_and_question(inputs)
+    document = [st.session_state.documents[c] for c in context]
+    llm = ChatOpenAI(model=st.session_state.model, temperature=0)
+    with get_openai_callback() as cb:
+        summary = summary_chain(llm).invoke({"section_name": query, "paper": document})
+        stats = cb
+    st.session_state.messages.append(
+        (f"/summarise-section {query}", summary, "identity")
+    )
+    st.session_state.costing.append(
+        {
+            "prompt tokens": stats.prompt_tokens,
+            "completion tokens": stats.completion_tokens,
+            "cost": stats.total_cost,
+        }
+    )
+    return (summary, "identity")
+
+
 def chain_of_density_wrapper(inputs):
     if inputs == []:
         raise InvalidArgumentError("Please provide snippet ids")
@@ -326,6 +347,7 @@ if __name__ == "__main__":
         ("/condense-summary", list, chain_of_density_wrapper),
         ("/insight-bullets", list, insights_bullet_wrapper),
         ("/paper-synopsis", list, synopsis_wrapper),
+        ("/summarise-section", str, summarise_wrapper),
     ]
     command_center = CommandCenter(
         default_input_type=str,

summary_chain.py

@@ -0,0 +1,18 @@
+from langchain_core.prompts import PromptTemplate
+from langchain_core.output_parsers import StrOutputParser
+
+summary_prompt_template = """
+Given the {section_name} section of a machine learning research paper, produce a comprehensive summary that encompasses all vital information, \
+    and detailed explanations of any mathematical equations present. 
+The goal is for this summary to function as an autonomous document that conveys the essence and key contributions of the research succinctly. 
+Ensure that if any mathematical content is present it is not only included but also clearly elucidated, highlighting its relevance to the research's overall objectives and results. 
+Structure the summary to be easily understandable, offering readers a full grasp of the section's critical insights without the need to consult the original paper.
+
+Here is the excerpt from the research paper: {paper}
+"""
+summary_output_parser = StrOutputParser()
+summary_prompt = PromptTemplate(
+    template=summary_prompt_template,
+    input_variables=["section_name", "paper"],
+)
+summary_chain = lambda model: summary_prompt | model | summary_output_parser