refine summarize chain

app_modules/llm_summarize_chain.py

@@ -1,5 +1,6 @@
 import os
 from typing import List, Optional
+from langchain import PromptTemplate
 
 from langchain.chains.base import Chain
 from langchain.chains.summarize import load_summarize_chain
@@ -7,12 +8,58 @@ from langchain.chains.summarize import load_summarize_chain
 from app_modules.llm_inference import LLMInference
 
 
+def get_llama_2_prompt_template(instruction):
+    B_INST, E_INST = "[INST]", "[/INST]"
+    B_SYS, E_SYS = "<<SYS>>\n", "\n<</SYS>>\n\n"
+
+    system_prompt = "You are a helpful assistant, you always only answer for the assistant then you stop. Read the text to get context"
+
+    SYSTEM_PROMPT = B_SYS + system_prompt + E_SYS
+    prompt_template = B_INST + SYSTEM_PROMPT + instruction + E_INST
+    return prompt_template
+
+
 class SummarizeChain(LLMInference):
     def __init__(self, llm_loader):
         super().__init__(llm_loader)
 
     def create_chain(self) -> Chain:
-        chain = load_summarize_chain(self.llm_loader.llm, chain_type="refine")
+        use_llama_2_prompt_template = (
+            os.environ.get("USE_LLAMA_2_PROMPT_TEMPLATE") == "true"
+        )
+        prompt_template = """Write a concise summary of the following:
+{text}
+CONCISE SUMMARY:"""
+
+        if use_llama_2_prompt_template:
+            prompt_template = get_llama_2_prompt_template(prompt_template)
+        prompt = PromptTemplate.from_template(prompt_template)
+
+        refine_template = (
+            "Your job is to produce a final summary\n"
+            "We have provided an existing summary up to a certain point: {existing_answer}\n"
+            "We have the opportunity to refine the existing summary"
+            "(only if needed) with some more context below.\n"
+            "------------\n"
+            "{text}\n"
+            "------------\n"
+            "Given the new context, refine the original summary."
+            "If the context isn't useful, return the original summary."
+        )
+
+        if use_llama_2_prompt_template:
+            refine_template = get_llama_2_prompt_template(refine_template)
+        refine_prompt = PromptTemplate.from_template(refine_template)
+
+        chain = load_summarize_chain(
+            llm=self.llm_loader.llm,
+            chain_type="refine",
+            question_prompt=prompt,
+            refine_prompt=refine_prompt,
+            return_intermediate_steps=True,
+            input_key="input_documents",
+            output_key="output_text",
+        )
         return chain
 
     def run_chain(self, chain, inputs, callbacks: Optional[List] = []):

summarize.py

@@ -15,17 +15,16 @@ from app_modules.init import app_init, get_device_types
 from app_modules.llm_summarize_chain import SummarizeChain
 
 
-def load_documents(source_pdfs_path, urls) -> List:
+def load_documents(source_pdfs_path, keep_page_info) -> List:
     loader = PyPDFDirectoryLoader(source_pdfs_path, silent_errors=True)
     documents = loader.load()
-    if urls is not None and len(urls) > 0:
+    if not keep_page_info:
         for doc in documents:
-            source = doc.metadata["source"]
-            filename = source.split("/")[-1]
-            for url in urls:
-                if url.endswith(filename):
-                    doc.metadata["url"] = url
-                    break
+            if doc is not documents[0]:
+                documents[0].page_content = (
+                    documents[0].page_content + "\n" + doc.page_content
+                )
+        documents = [documents[0]]
     return documents
 
 
@@ -43,8 +42,11 @@ source_pdfs_path = (
 )
 chunk_size = sys.argv[2] if len(sys.argv) > 2 else os.environ.get("CHUNCK_SIZE")
 chunk_overlap = sys.argv[3] if len(sys.argv) > 3 else os.environ.get("CHUNK_OVERLAP")
+keep_page_info = (
+    sys.argv[3] if len(sys.argv) > 3 else os.environ.get("KEEP_PAGE_INFO")
+) == "true"
 
-sources = load_documents(source_pdfs_path, None)
+sources = load_documents(source_pdfs_path, keep_page_info)
 
 print(f"Splitting {len(sources)} PDF pages in to chunks ...")