lenox-ai/prototype

1from langchain.chains.summarize import load_summarize_chain2from langchain.chat_models import ChatOpenAI3from src.prompts import (4    prompts,5    prompts_parallel_summary,6)7from src.doc_loading import load_docs8from src.llm_utils import async_generate_llmchain9import time10from typing import Dict, List11import asyncio12 13 14def summarize_chain(15    file_path: str, llm: ChatOpenAI, summarization_kwargs: Dict[str, str]16) -> str:17    """Summarize a pdf file. The summarization is done by the language model.18 19    Args:20        file_path (str): Path to the pdf file. This can either be a local path or a tempfile.TemporaryFileWrapper_.21        llm (ChatOpenAI): Language model to use for the summarization.22 23    Returns:24        str: Summarization of the pdf file.25    """26    docs = load_docs(file_path=file_path)27    chain = load_summarize_chain(28        llm=llm,29        **summarization_kwargs,30    )31    summary = chain.run(docs)32    return summary33 34 35def summarize_wrapper(36    file: str, llm: ChatOpenAI, summarization_type: str, summarization_kwargs: dict37) -> str:38    """Wrapper for the summarization function to make it compatible with gradio. This function uses a39        single summarization chain.40 41    Args:42        file (str): Path to the file. This can either be a local path or a tempfile.TemporaryFileWrapper_.43        llm (ChatOpenAI): Language model.44        summarization_type (str): Type of summarization. Can be either "short", "middle" or "long".45        summarization_kwargs (dict): Keyword arguments for the summarization.46 47    Returns:48        str: Summarization of the file.49    """50    if summarization_type == "short":51        summarization_kwargs.update(52            dict(53                map_prompt=prompts["short_de"]["map_prompt"],54                combine_prompt=prompts["short_de"]["combine_prompt"],55            )56        )57    elif summarization_type == "middle":58        summarization_kwargs.update(59            dict(60                map_prompt=prompts["middle_de"]["map_prompt"],61                combine_prompt=prompts["middle_de"]["combine_prompt"],62            )63        )64    elif summarization_type == "long":65        summarization_kwargs.update(66            dict(67                map_prompt=prompts["long_de"]["map_prompt"],68                combine_prompt=prompts["long_de"]["combine_prompt"],69            )70        )71    else:72        raise ValueError(f"Summarization type {summarization_type} is not supported.")73 74    return summarize_chain(75        file_path=file.name, llm=llm[0], summarization_kwargs=summarization_kwargs76    )77 78 79async def generate_summary_concurrently(80    file_path: str, sections: List[str], llm: ChatOpenAI81) -> List[dict]:82    """Parallel summarization. This function is used to run different prompts for the same docs in parallel.83 84    Args:85        file_path (str): Path to the pdf file. This can either be a local path or a tempfile.TemporaryFileWrapper_.86        sections (List[str]): List of sections to summarize selected by the user.87        llm (ChatOpenAI): Language model to use for the summarization.88 89    Returns:90        List: List of summarizations.91    """92 93    docs = load_docs(file_path=file_path, with_pageinfo=False)94    summarization_kwargs = dict()95 96    # create parallel tasks97    tasks = []98    for k in PARALLEL_SUMMARIZATION_ORDER:99        if PARALLEL_SUMMARIZATION_MAPPING_INVERSE.get(k, k) in sections:100            sk = summarization_kwargs.copy()101            sk["prompt"] = prompts_parallel_summary[k]102            print(f"Appending task for summary: {k}")103            tasks.append(104                async_generate_llmchain(llm=llm, docs=docs, llm_kwargs=sk, k=k)105            )106    print("-------------------")107    # execute all coroutines concurrently108    values = await asyncio.gather(*tasks)109 110    # report return values111    values_flattened = {}112    for v in values:113        values_flattened.update(v)114    return values_flattened115 116 117PARALLEL_SUMMARIZATION_ORDER = [118    "intro",119    "darstellung_des_rechtsproblems",120    "II.  Die Entscheidung",121    "angaben_ueber_das_urteil",122    "sachverhalt",123    "prozessgeschichte",124    "rechtsproblem",125    "loesung_des_gerichts",126]127PARALLEL_SUMMARIZATION_MAPPING = {128    "I.  Einleitung": "intro",129    "Darstellung des Rechtsproblems": "darstellung_des_rechtsproblems",130    "Angaben über das Urteil": "angaben_ueber_das_urteil",131    "Sachverhalt": "sachverhalt",132    "Prozessgeschichte": "prozessgeschichte",133    "Rechtsproblem": "rechtsproblem",134    "Lösung des Gerichts": "loesung_des_gerichts",135}136PARALLEL_SUMMARIZATION_MAPPING_INVERSE = {137    v: k for k, v in PARALLEL_SUMMARIZATION_MAPPING.items()138}139 140 141def parallel_summarization(file: str, sections: List[str], llm: ChatOpenAI) -> str:142    """Wrapper for the parallel summarization function to make it compatible with gradio.143 144    Args:145        file (str): Path to the file. This can either be a local path or a tempfile.TemporaryFileWrapper_.146        sections (List[str]): List of sections to summarize.147        llm (ChatOpenAI): Language model.148 149    Returns:150        str: Summarization of the file.151    """152    now = time.time()153 154    values_flattened = asyncio.run(155        generate_summary_concurrently(156            file_path=file.name, sections=sections, llm=llm[0]157        )158    )159 160    print("Time taken for complete parallel summarization: ", time.time() - now)161    output = ""162 163    for section in values_flattened.keys():164        output += (165            values_flattened.get(166                section, PARALLEL_SUMMARIZATION_MAPPING_INVERSE.get(section, section)167            )168            + "\n\n"169        )170 171    return output172