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