from langchain_core.output_parsers import JsonOutputParser
from langchain_core.prompts import PromptTemplate

qa_prompt_template = """
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 abstract: ####{paper}####
"""

qa_prompt = PromptTemplate(
    template=qa_prompt_template,
    input_variables=["paper"],
)
auto_qa_chain = lambda model: qa_prompt | model | JsonOutputParser()