Upload folder using huggingface_hub

requirements.txt

@@ -7,7 +7,7 @@ scipy>=0.19
 openai==0.27.7
 grpcio-tools==1.54.2
 gpt_index==0.4.24
-langchain==0.0.190
+langchain==0.0.236
 environs==9.5.0
 pypdf==3.9.1
 pypdfium2==4.18.0

src/doc_loading.py

@@ -0,0 +1,37 @@
+from langchain.document_loaders import PyPDFLoader, TextLoader
+from langchain.docstore.document import Document
+from typing import List
+
+
+def load_docs(file_path: str, with_pageinfo: bool = True) -> List[Document]:
+    """Load a file and return the text.
+
+    Args:
+        file_path (str): Path to the pdf file. This can either be a local path or a tempfile.TemporaryFileWrapper_.
+        with_pageinfo (bool, optional): If True the page information is added to the document. Defaults to True.
+
+    Raises:
+        ValueError: If the file type is not supported.
+
+    Returns:
+        List[Document]: List of documents.
+    """
+    if file_path.endswith(".pdf"):
+        loader = PyPDFLoader(file_path)
+        docs = loader.load()
+    elif file_path.endswith(".txt"):
+        loader = TextLoader(file_path)
+        docs = loader.load()
+    else:
+        raise ValueError(
+            f"File type ({file_path.split('.')[1]}) not supported. Please upload a pdf or txt file."
+        )
+    for doc in docs:
+        doc.page_content = doc.page_content.replace("\n", " \n ")
+        # if doc contains a page append it to the text
+        if with_pageinfo and hasattr(doc, "metadata"):
+            doc.page_content = f"(Quelle Seite: {doc.metadata.get('page')+1}) .".join(
+                doc.page_content.split(" .")
+            )
+
+    return docs

src/gradio_app.py

@@ -4,11 +4,15 @@ import pypdfium2 as pdfium
 import gradio as gr
 
 from langchain.chat_models import ChatOpenAI
-from src.summarization import summarize_wrapper, parallel_summarization
+from src.summarization import (
+    parallel_summarization,
+    parallel_legal_implications,
+    PARALLEL_SUMMARIZATION_MAPPING,
+)
 from src.mailing import send_email
 
-# Function to render a specific page of a PDF file as an image
-def render_file(file):
+
+def _file_render_helper(file):
     pdf = pdfium.PdfDocument(file.name)
     images = []
     for page_index in range(len(pdf)):
@@ -18,13 +22,31 @@ def render_file(file):
             rotation=0,  # no additional rotation
             # ... further rendering options
         )
-        images.append((bitmap.to_pil(), f"Seite {page_index+1}"))
+        images.append(
+            (bitmap.to_pil(), f"{file.name.split('/')[-1]} Seite {page_index+1}")
+        )
+    return images
+
 
+# Function to render a specific page of a PDF file as an image
+def render_file(file):
+    images = _file_render_helper(file)
     return gr.update(
         value=images,
     )
 
 
+def render_files(files):
+    all_images = []
+    for file in files:
+        images = _file_render_helper(file)
+        all_images.extend(images)
+
+    return gr.update(
+        value=all_images,
+    )
+
+
 def switch_buttons(interactive: bool):
     """This switches the buttons to interactive or not interactive.
 
@@ -44,59 +66,63 @@ def switch_buttons(interactive: bool):
     )
 
 
-def run_summarization_model_gradio(
-    llm: ChatOpenAI,
-    share_gradio_via_link: bool = False,
-    summarization_kwargs: dict = {},
-    run_local: bool = False,
-):
-    """Run the Summarization assistant with gradio
+def load_summary_section(llm: ChatOpenAI):
+    """Load the summary section
 
     Args:
         llm (ChatOpenAI): Language model.
-        share_gradio_via_link (bool, optional): Whether to launch the gradio app via a public link. Defaults to False.
-        summarization_kwargs (dict, optional): Keyword arguments for the summarization. Defaults to {}.
-        run_local (bool, optional): Whether to run the gradio app locally. Defaults to False.
 
+    Returns:
+        gr.Blocks: The summarization section
     """
-    title = "Summarization of Legal Documents"
-    description = f"Upload a document and get a summarization."
 
     with gr.Blocks(
         theme="soft",
-        title=title,
-    ) as webui:
-        with gr.Row().style(equal_height=True):
-            Header_box = generate_title(title=title, description=description)
-        with gr.Row().style(equal_height=True):
-            clear = gr.Button("Clear All Components")
-            file_upload = gr.File(
-                file_count="single",
-                file_types=[".pdf", ".txt"],
-                label="Upload PDF",
-            )
+    ) as summary_section:
         with gr.Row().style(equal_height=True):
-            summary_short = gr.Button("Kurze Zusammenfassung", interactive=False)
-            summary_middle = gr.Button("Mittlere Zusammenfassung", interactive=False)
-            summary_long = gr.Button("Lange Zusammenfassung", interactive=False)
-            summary_parallel = gr.Button("Parallele Zusammenfassung", interactive=False)
+            with gr.Column(scale=1):
+                file_upload_summary = gr.File(
+                    file_count="single",
+                    file_types=[".pdf", ".txt"],
+                    label="Upload PDF",
+                )
+                summary_parallel_button = gr.Button(
+                    "Parallel Summary", interactive=False
+                )
+                clear = gr.Button("Clear All Components")
+            with gr.Column(scale=2):
+                sections_to_select = [
+                    i for i in PARALLEL_SUMMARIZATION_MAPPING.keys() if "I." not in i
+                ]
+                summary_sections_dropdown = gr.Dropdown(
+                    sections_to_select,
+                    value=sections_to_select,
+                    interactive=True,
+                    multiselect=True,
+                    label="Sections for Summarization",
+                    info="Select the sections you want to include in the summarization.",
+                )
         with gr.Row().style(equal_height=True):
             with gr.Column(scale=1):
-                summary_output = gr.Textbox(label="Zusammenfassung", lines=9).style(
+                summary_output = gr.Textbox(label="Summary", lines=9).style(
                     show_copy_button=True
                 )
             with gr.Column(scale=1):
-                show_pdf = gr.Gallery(label="Uploaded PDF").style(object_fit="contain")
+                summary_show_pdf = gr.Gallery(label="Uploaded PDF").style(
+                    object_fit="contain"
+                )
 
         with gr.Row().style(equal_height=True):
             with gr.Column(scale=1):
-                recipiant_email = gr.Textbox(
+                recipiant_email_summary = gr.Textbox(
                     label="Recipiant Email", placeholder="Enter Email"
                 )
-                subject_email = gr.Textbox(label="Subject", placeholder="Enter Subject")
+                subject_email_summary = gr.Textbox(
+                    label="Subject", placeholder="Enter Subject"
+                )
                 send_email_button = gr.Button("Open Email", interactive=False)
             with gr.Column(scale=3):
-                email_instructions = gr.Textbox(
+                email_instructions_summary = gr.Textbox(
                     label="Email Instructions",
                     placeholder="Write Email Instructions here.",
                     value=(
@@ -110,91 +136,39 @@ def run_summarization_model_gradio(
                 )
 
         # Once a file is uploaded, enable the summarization buttons and visualize the uploaded file
-        file_upload.upload(
+        file_upload_summary.upload(
             switch_buttons,
             [gr.State(True)],
-            [summary_short, summary_middle, summary_long],
+            [summary_parallel_button, gr.State(None), gr.State(None)],
             queue=False,
-        ).then(
-            switch_buttons,
-            [gr.State(True)],
-            [summary_parallel, gr.State(None), gr.State(None)],
-            queue=False,
-        ).then(
-            fn=render_file, inputs=[file_upload], outputs=[show_pdf]
-        )
+        ).then(fn=render_file, inputs=[file_upload_summary], outputs=[summary_show_pdf])
 
-        # If you click any button first disable all buttons, then summarzize and then enable the clicked button
-        for s, summarization_type in [
-            (summary_short, "short"),
-            (summary_middle, "middle"),
-            (summary_long, "long"),
-        ]:
-            s.click(
-                switch_buttons,
-                [gr.State(False)],
-                [summary_short, summary_middle, summary_long],
-                queue=False,
-            ).then(
-                summarize_wrapper,
-                [
-                    file_upload,
-                    gr.State([llm]),
-                    gr.State(summarization_type),
-                    gr.State(summarization_kwargs),
-                ],
-                [summary_output],
-                queue=False,
-            ).then(
-                switch_buttons,
-                [gr.State(True)],
-                [summary_short, summary_middle, summary_long],
-                queue=False,
-            ).then(
-                switch_buttons,
-                [gr.State(True)],
-                [send_email_button, gr.State(None), gr.State(None)],
-                queue=False,
-            )
-
-        summary_parallel.click(
+        summary_parallel_button.click(
             switch_buttons,
             [gr.State(False)],
-            [summary_short, summary_middle, summary_long],
+            [summary_parallel_button, gr.State(None), gr.State(None)],
             queue=False,
         ).then(
             parallel_summarization,
-            [file_upload, gr.State([llm]), gr.State(summarization_kwargs)],
+            [file_upload_summary, summary_sections_dropdown, gr.State([llm])],
             [summary_output],
             queue=False,
         ).then(
             switch_buttons,
             [gr.State(True)],
-            [summary_short, summary_middle, summary_long],
+            [summary_parallel_button, gr.State(None), gr.State(None)],
             queue=False,
-        ).then(
-            switch_buttons,
-            [gr.State(True)],
-            [send_email_button, gr.State(None), gr.State(None)],
-            queue=False,
-        )
+        ).then
 
         # The clear button clears the dashboard
         clear.click(lambda: None, None, summary_output, queue=False).then(
-            lambda: None, None, file_upload, queue=False
-        ).then(
-            switch_buttons,
-            [gr.State(False)],
-            [summary_short, summary_middle, summary_long],
-            queue=False,
-        ).then(
-            lambda: None, None, show_pdf, queue=False
-        ).then(
+            lambda: None, None, file_upload_summary, queue=False
+        ).then(lambda: None, None, summary_show_pdf, queue=False).then(
             lambda: None, None, send_email_button, queue=False
         ).then(
-            lambda: None, None, email_instructions, queue=False
+            lambda: None, None, email_instructions_summary, queue=False
         ).then(
-            lambda: None, None, recipiant_email, queue=False
+            lambda: None, None, recipiant_email_summary, queue=False
         )
 
         # Email button click opens the default email client and fills in the email instructions
@@ -202,12 +176,129 @@ def run_summarization_model_gradio(
             send_email,
             [
                 summary_output,
-                recipiant_email,
-                subject_email,
-                email_instructions,
+                recipiant_email_summary,
+                subject_email_summary,
+                email_instructions_summary,
             ],
             queue=False,
         )
+    return summary_section
+
+
+def load_legal_implications_section(llm: ChatOpenAI):
+    """Load the legal implications section
+
+    Args:
+        llm (ChatOpenAI): Language model.
+
+    Returns:
+        gr.Block: Legal Implications Section
+    """
+    with gr.Blocks(theme="soft") as legal_implications_section:
+
+        with gr.Row().style(equal_height=True):
+            with gr.Column(scale=3):
+                file_upload_legal_implications = gr.File(
+                    file_count="multiple",
+                    file_types=[".pdf", ".txt"],
+                    label="Upload PDF",
+                )
+            with gr.Column(scale=1):
+                extract_legal_implications_button = gr.Button(
+                    "Extract Legal Implications", interactive=False
+                )
+                clear_legal_implications_button = gr.Button("Clear All Components")
+
+        with gr.Row().style(equal_height=True):
+            with gr.Column(scale=1):
+                legal_implications_output = gr.Textbox(
+                    label="Legal Implications", lines=9
+                ).style(show_copy_button=True)
+            with gr.Column(scale=1):
+                legal_implications_show_pdf = gr.Gallery(label="Uploaded PDF").style(
+                    object_fit="contain"
+                )
+
+        with gr.Row().style(equal_height=True):
+            with gr.Column(scale=1):
+                recipiant_email_legal_implications = gr.Textbox(
+                    label="Recipiant Email", placeholder="Enter Email"
+                )
+                subject_email_legal_implications = gr.Textbox(
+                    label="Subject", placeholder="Enter Subject"
+                )
+                send_email_button = gr.Button("Open Email", interactive=False)
+            with gr.Column(scale=3):
+                email_instructions_legal_implications = gr.Textbox(
+                    label="Email Instructions",
+                    placeholder="Write Email Instructions here.",
+                    value=(
+                        "Dear Recipient\n\n"
+                        "Please find the Legal Implications of the uploaded documents below.\n\n"
+                        "<TEXT_FROM_LLM>\n\n"
+                        "Kind regards,\n"
+                        "Your Legal Assistant"
+                    ),
+                    lines=9,
+                )
+        # Once a file is uploaded, enable the summarization buttons and visualize the uploaded file
+        file_upload_legal_implications.upload(
+            switch_buttons,
+            [gr.State(True)],
+            [extract_legal_implications_button, gr.State(None), gr.State(None)],
+            queue=False,
+        ).then(
+            fn=render_files,
+            inputs=[file_upload_legal_implications],
+            outputs=[legal_implications_show_pdf],
+        )
+
+        extract_legal_implications_button.click(
+            switch_buttons,
+            [gr.State(False)],
+            [extract_legal_implications_button, gr.State(None), gr.State(None)],
+            queue=False,
+        ).then(
+            parallel_legal_implications,
+            [file_upload_legal_implications, gr.State([llm])],
+            [legal_implications_output],
+            queue=False,
+        ).then(
+            switch_buttons,
+            [gr.State(True)],
+            [extract_legal_implications_button, gr.State(None), gr.State(None)],
+            queue=False,
+        )
+
+
+def run_summarization_model_gradio(
+    llm: ChatOpenAI,
+    share_gradio_via_link: bool = False,
+    summarization_kwargs: dict = {},
+    run_local: bool = False,
+):
+    """Run the Summarization assistant with gradio
+
+    Args:
+        llm (ChatOpenAI): Language model.
+        share_gradio_via_link (bool, optional): Whether to launch the gradio app via a public link. Defaults to False.
+        summarization_kwargs (dict, optional): Keyword arguments for the summarization. Defaults to {}.
+        run_local (bool, optional): Whether to run the gradio app locally. Defaults to False.
+
+    """
+    title = "Summarization of Legal Documents"
+    description = f"Upload a document and get a summarization."
+
+    with gr.Blocks(
+        theme="soft",
+        title=title,
+    ) as webui:
+        with gr.Row().style(equal_height=True):
+            Header_box = generate_title(title=title, description=description)
+        with gr.Tab("Summarize Verdict"):
+            load_summary_section(llm=llm)
+        with gr.Tab("Legal Implications"):
+            load_legal_implications_section(llm=llm)
 
     webui.queue()
 

src/llm_utils.py

@@ -0,0 +1,28 @@
+from langchain.chains.llm import LLMChain
+from langchain.chat_models import ChatOpenAI
+from langchain.docstore.document import Document
+import time
+from typing import List
+
+
+async def async_generate(
+    llm: ChatOpenAI, docs: List[Document], llm_kwargs: dict, k: str
+) -> dict:
+    """Asyncronous LLMChain function.
+
+    Args:
+        llm (ChatOpenAI): Language model to use.
+        docs (List[Document]): List of documents.
+        llm_kwargs (dict): Keyword arguments for the LLMChain.
+        k (str): Key for a dictionary under which the output is returned.
+
+    Returns:
+        dict: Dictionary with the summarization.
+    """
+    print(f"Starting summarization for {k}")
+    now = time.time()
+    chain = LLMChain(llm=llm, **llm_kwargs)
+
+    resp = await chain.arun(text=docs)
+    print(f"Time taken for {k}: ", time.time() - now)
+    return {k: resp}

src/prompts.py

@@ -1,6 +1,10 @@
 from langchain.prompts.prompt import PromptTemplate
 
 
+#########################################
+###### SUMMARIZATION CHAIN PROMPTS ######
+#########################################
+
 prompts = {
     ############## SHORT DE
     "short_de": {
@@ -152,12 +156,14 @@ Die Teile der Zusammenfassung mit Angabe der Seitenzahlen:
 }
 
 
-def get_template_mp(name: str, headline: str, additional_text: str = ""):
+#########################################
+#### PARALLEL SUMMARIZATION PROMPTS #####
+#########################################
+def get_template_parallel(name: str, headline: str, additional_text: str = ""):
     base_multi = (
         "Schreibe, ein/e <KEY> des Urteils, das durch dreifache Anführungszeichen begrenzt ist, in maximal einem Paragraphen.\n"
         "<ADDITIONAL_TEXT>\n"
         'Als Überschrift muss "<HEAD_LINE>" angegeben werden. \n'
-        # "Nach dem Paragraph müssen die Seiten angegeben werden die genutzt wurden."
         "Urteil:\n"
         "```{text}```\n"
         "\n"
@@ -173,18 +179,18 @@ def get_template_mp(name: str, headline: str, additional_text: str = ""):
 prompts_parallel = {
     "intro": PromptTemplate(
         input_variables=["text"],
-        template=get_template_mp(name="Einleitung", headline="I. Einleitung"),
+        template=get_template_parallel(name="Einleitung", headline="I. Einleitung"),
     ),
     "darstellung_des_rechtsproblems": PromptTemplate(
         input_variables=["text"],
-        template=get_template_mp(
+        template=get_template_parallel(
             name="Darstellung des Rechtsproblems",
             headline="Darstellung des Rechtsproblems",
         ),
     ),
     "angaben_ueber_das_urteil": PromptTemplate(
         input_variables=["text"],
-        template=get_template_mp(
+        template=get_template_parallel(
             name="Angaben über das Urteil",
             headline="Angaben über das Urteil",
             additional_text="Gib die folgenden Informationen an: Gericht, Datum, Aktenzeichen (AZ: ...), Fundstelle(n)",
@@ -192,7 +198,7 @@ prompts_parallel = {
     ),
     "sachverhalt": PromptTemplate(
         input_variables=["text"],
-        template=get_template_mp(
+        template=get_template_parallel(
             name="Sachverhalt",
             headline="Sachverhalt (unter Rückgriff auf Instanzentscheidung)",
             additional_text=(
@@ -203,13 +209,13 @@ prompts_parallel = {
     ),
     "prozessgeschichte": PromptTemplate(
         input_variables=["text"],
-        template=get_template_mp(
+        template=get_template_parallel(
             name="Prozessgeschichte", headline="3. Prozessgeschichte"
         ),
     ),
     "rechtsproblem": PromptTemplate(
         input_variables=["text"],
-        template=get_template_mp(
+        template=get_template_parallel(
             name="Rechtsproblem",
             headline="Rechtsproblem",
             additional_text="Das Problem des Falles ist genau herauszuarbeiten und im rechtlichen Kontext zu verankern.",
@@ -217,13 +223,13 @@ prompts_parallel = {
     ),
     "loesung_des_gerichts": PromptTemplate(
         input_variables=["text"],
-        template=get_template_mp(
+        template=get_template_parallel(
             name="Lösung des Gerichts", headline="Lösung des Gerichts"
         ),
     ),
     "loesungsansaetze_zum_problem": PromptTemplate(
         input_variables=["text"],
-        template=get_template_mp(
+        template=get_template_parallel(
             name="Lösungsansätze zum Problem",
             headline="Lösungsansätze zum Problem",
             additional_text="Knappe, aber möglichst vollständige Übersicht der vertretenen Ansichten bzw. der Lösungsvorschläge im Urteil.",
@@ -231,7 +237,7 @@ prompts_parallel = {
     ),
     "analyse_und_einordnung_der_entscheidung": PromptTemplate(
         input_variables=["text"],
-        template=get_template_mp(
+        template=get_template_parallel(
             name="Analyse und Einordnung der Entscheidung",
             headline="Analyse und Einordnung der Entscheidung",
             additional_text="Es soll nur der Inhalt des Urteils wiedergegeben werden.",
@@ -239,7 +245,7 @@ prompts_parallel = {
     ),
     "bewertung_und_kritik_der_entscheidung": PromptTemplate(
         input_variables=["text"],
-        template=get_template_mp(
+        template=get_template_parallel(
             name="Bewertung und Kritik der Entscheidung",
             headline="Bewertung und Kritik der Entscheidung",
             additional_text="Verwende ausschließlich den Kontext des Urteils und schreib keinen neuen Text. Wenn keine Bewertung oder Kritik vorhanden ist, antworte mit 'Keine Bewertung oder Kritik vorhanden.'",
@@ -247,7 +253,7 @@ prompts_parallel = {
     ),
     "eigener_loesungsvorschlag": PromptTemplate(
         input_variables=["text"],
-        template=get_template_mp(
+        template=get_template_parallel(
             name="Eigener Lösungsvorschlag",
             headline="Eigener Lösungsvorschlag",
             additional_text="Es soll nur der Inhalt des Urteils wiedergegeben werden. Wenn das Urteil keinen eigenen Lösungsvorschlag hat schreib: 'Keine Informationen zum eigenen Lösungsvorschlag vorhanden'",
@@ -255,7 +261,7 @@ prompts_parallel = {
     ),
     "ausblick": PromptTemplate(
         input_variables=["text"],
-        template=get_template_mp(
+        template=get_template_parallel(
             name="Ausblick",
             headline="Ausblick",
             additional_text="Es soll nur der Inhalt des Urteils wiedergegeben werden. Wenn das Urteil keinen Ausblick gibt schreib: 'Keine Informationen zum Auslbick vorhanden'.",

src/summarization.py

@@ -1,49 +1,13 @@
-from langchain.document_loaders import PyPDFLoader, TextLoader
 from langchain.chains.summarize import load_summarize_chain
-from langchain.chains.llm import LLMChain
-from langchain.chains.combine_documents.stuff import StuffDocumentsChain
 from langchain.chat_models import ChatOpenAI
-from langchain.docstore.document import Document
 from src.prompts import prompts, prompts_parallel
+from src.doc_loading import load_docs
+from src.llm_utils import async_generate
 import time
 from typing import Dict, List
 import asyncio
 
 
-def load_docs(file_path: str, with_pageinfo: bool = True) -> List[Document]:
-    """Load a file and return the text.
-
-    Args:
-        file_path (str): Path to the pdf file. This can either be a local path or a tempfile.TemporaryFileWrapper_.
-        with_pageinfo (bool, optional): If True the page information is added to the document. Defaults to True.
-
-    Raises:
-        ValueError: If the file type is not supported.
-
-    Returns:
-        List[Document]: List of documents.
-    """
-    if file_path.endswith(".pdf"):
-        loader = PyPDFLoader(file_path)
-        docs = loader.load()
-    elif file_path.endswith(".txt"):
-        loader = TextLoader(file_path)
-        docs = loader.load()
-    else:
-        raise ValueError(
-            f"File type ({file_path.split('.')[1]}) not supported. Please upload a pdf or txt file."
-        )
-    for doc in docs:
-        doc.page_content = doc.page_content.replace("\n", " \n ")
-        # if doc contains a page append it to the text
-        if with_pageinfo and hasattr(doc, "metadata"):
-            doc.page_content = f"(Quelle Seite: {doc.metadata.get('page')+1}) .".join(
-                doc.page_content.split(" .")
-            )
-
-    return docs
-
-
 def summarize_chain(
     file_path: str, llm: ChatOpenAI, summarization_kwargs: Dict[str, str]
 ) -> str:
@@ -61,13 +25,6 @@ def summarize_chain(
         llm=llm,
         **summarization_kwargs,
     )
-    # del summarization_kwargs["map_prompt"]
-    # summarization_kwargs["prompt"] = summarization_kwargs["combine_prompt"]
-    # del summarization_kwargs["combine_prompt"]
-    # llm_chain = LLMChain(llm=llm, **summarization_kwargs)
-    # chain = StuffDocumentsChain(
-    #     llm_chain=llm_chain, document_variable_name="text"
-    # )
     summary = chain.run(docs)
     return summary
 
@@ -75,7 +32,8 @@ def summarize_chain(
 def summarize_wrapper(
     file: str, llm: ChatOpenAI, summarization_type: str, summarization_kwargs: dict
 ) -> str:
-    """Wrapper for the summarization function to make it compatible with gradio.
+    """Wrapper for the summarization function to make it compatible with gradio. This function uses a
+        single summarization chain.
 
     Args:
         file (str): Path to the file. This can either be a local path or a tempfile.TemporaryFileWrapper_.
@@ -115,34 +73,14 @@ def summarize_wrapper(
     )
 
 
-async def async_generate(
-    llm: ChatOpenAI, docs: List[Document], summarization_kwargs: dict, k: str
-) -> dict:
-    """Asyncronous summarization.
-
-    Args:
-        llm (ChatOpenAI): Language model to use for the summarization.
-        docs (List[Document]): List of documents.
-        summarization_kwargs (dict): Keyword arguments for the summarization.
-        k (str): Key for the summarization.
-
-    Returns:
-        dict: Dictionary with the summarization.
-    """
-    print(f"Starting summarization for {k}")
-    now = time.time()
-    # chain = load_summarize_chain(llm=llm, **summarization_kwargs)
-    chain = LLMChain(llm=llm, **summarization_kwargs)
-    resp = await chain.arun(text=docs)
-    print(f"Time taken for {k}: ", time.time() - now)
-    return {k: resp}
-
-
-async def generate_concurrently(file_path: str, llm: ChatOpenAI) -> List[dict]:
-    """Parallel summarization.
+async def generate_summary_concurrently(
+    file_path: str, sections: List[str], llm: ChatOpenAI
+) -> List[dict]:
+    """Parallel summarization. This function is used to run different prompts for the same docs in parallel.
 
     Args:
         file_path (str): Path to the pdf file. This can either be a local path or a tempfile.TemporaryFileWrapper_.
+        sections (List[str]): List of sections to summarize.
         llm (ChatOpenAI): Language model to use for the summarization.
 
     Returns:
@@ -154,12 +92,12 @@ async def generate_concurrently(file_path: str, llm: ChatOpenAI) -> List[dict]:
 
     # create parallel tasks
     tasks = []
-    i = 0
-    for k, pt in prompts_parallel.items():
-        sk = summarization_kwargs.copy()
-        sk["prompt"] = pt
-        print(f"Appending task for {k}")
-        tasks.append(async_generate(llm=llm, docs=docs, summarization_kwargs=sk, k=k))
+    for k in PARALLEL_SUMMARIZATION_ORDER:
+        if PARALLEL_SUMMARIZATION_MAPPING_INVERSE.get(k, k) in sections:
+            sk = summarization_kwargs.copy()
+            sk["prompt"] = prompts_parallel[k]
+            print(f"Appending task for {k}")
+            tasks.append(async_generate(llm=llm, docs=docs, llm_kwargs=sk, k=k))
     print("-------------------")
     # execute all coroutines concurrently
     values = await asyncio.gather(*tasks)
@@ -171,55 +109,80 @@ async def generate_concurrently(file_path: str, llm: ChatOpenAI) -> List[dict]:
     return values_flattened
 
 
-def parallel_summarization(
-    file: str, llm: ChatOpenAI, summarization_kwargs: dict
-) -> str:
-    """Wrapper for the summarization function to make it compatible with gradio.
+PARALLEL_SUMMARIZATION_ORDER = [
+    "intro",
+    "darstellung_des_rechtsproblems",
+    "II.  Die Entscheidung",
+    "angaben_ueber_das_urteil",
+    "sachverhalt",
+    "prozessgeschichte",
+    "rechtsproblem",
+    "loesung_des_gerichts",
+    "III.  Analyse",
+    "loesungsansaetze_zum_problem",
+    "analyse_und_einordnung_der_entscheidung",
+    "bewertung_und_kritik_der_entscheidung",
+    "eigener_loesungsvorschlag",
+    "ausblick",
+]
+PARALLEL_SUMMARIZATION_MAPPING = {
+    "I.  Einleitung": "intro",
+    "Darstellung des Rechtsproblems": "darstellung_des_rechtsproblems",
+    "Angaben über das Urteil": "angaben_ueber_das_urteil",
+    "Sachverhalt": "sachverhalt",
+    "Prozessgeschichte": "prozessgeschichte",
+    "Rechtsproblem": "rechtsproblem",
+    "Lösung des Gerichts": "loesung_des_gerichts",
+    "Lösungsansätze zum Problem": "loesungsansaetze_zum_problem",
+    "Analyse und Einordnung der Entscheidung": "analyse_und_einordnung_der_entscheidung",
+    "Bewertung und Kritik der Entscheidung": "bewertung_und_kritik_der_entscheidung",
+    "Eigener Lösungsvorschlag": "eigener_loesungsvorschlag",
+    "Ausblick": "ausblick",
+}
+PARALLEL_SUMMARIZATION_MAPPING_INVERSE = {
+    v: k for k, v in PARALLEL_SUMMARIZATION_MAPPING.items()
+}
+
+
+def parallel_summarization(file: str, sections: List[str], llm: ChatOpenAI) -> str:
+    """Wrapper for the parallel summarization function to make it compatible with gradio.
 
     Args:
         file (str): Path to the file. This can either be a local path or a tempfile.TemporaryFileWrapper_.
+        sections (List[str]): List of sections to summarize.
         llm (ChatOpenAI): Language model.
-        summarization_kwargs (dict): Keyword arguments for the summarization.
 
     Returns:
         str: Summarization of the file.
     """
     now = time.time()
     values_flattened = asyncio.run(
-        generate_concurrently(file_path=file.name, llm=llm[0])
+        generate_summary_concurrently(
+            file_path=file.name, sections=sections, llm=llm[0]
+        )
     )
-    print("Time taken: ", time.time() - now)
-
-    output = f"""
-    
-{values_flattened["intro"]}
-
-{values_flattened["darstellung_des_rechtsproblems"]}
-
-
-II.  Die Entscheidung
-
-{values_flattened["angaben_ueber_das_urteil"]}
-
-{values_flattened["sachverhalt"]}
-
-{values_flattened["prozessgeschichte"]}
-
-{values_flattened["rechtsproblem"]}
-
-{values_flattened["loesung_des_gerichts"]}
-
-III.  Analyse
-
-{values_flattened["loesungsansaetze_zum_problem"]}
+    print("Time taken for complete parallel summarization: ", time.time() - now)
+    order = PARALLEL_SUMMARIZATION_ORDER
+    output = ""
+    for section in order:
+        output += (
+            values_flattened.get(
+                section, PARALLEL_SUMMARIZATION_MAPPING_INVERSE.get(section, section)
+            )
+            + "\n\n"
+        )
 
-{values_flattened["analyse_und_einordnung_der_entscheidung"]}
+    return output
 
-{values_flattened["bewertung_und_kritik_der_entscheidung"]}
 
-{values_flattened["eigener_loesungsvorschlag"]}
+def parallel_legal_implications(file: str, llm: ChatOpenAI) -> str:
+    """Wrapper for the parallel legal implication extraction function to make it compatible with gradio.
 
-{values_flattened["ausblick"]}
-"""
+    Args:
+        file (str): Path to the file. This can either be a local path or a tempfile.TemporaryFileWrapper_.
+        llm (ChatOpenAI): Language model.
 
-    return output
+    Returns:
+        str: Legal Implications of the file.
+    """
+    return "TBD"