Upload 3 files

utils/__init__.py

@@ -1,6 +1,14 @@
 from .refine_summary import RefineDataSummarizer
-
+from .templates import (
+    prompt_template,
+    refine_template,
+    prompt_template_bullet_point,
+    refine_prompt_template_bullet_point
+)
 __all__ = [
-
-    "RefineDataSummarizer"
+    "RefineDataSummarizer",
+    "prompt_template",
+    "refine_template",
+    "prompt_template_bullet_point",
+    "refine_prompt_template_bullet_point"
 ]

utils/refine_summary.py

@@ -24,72 +24,22 @@ class RefineDataSummarizer:
 
     def __init__(
             self,
-            llm: BaseChatModel
+            llm: BaseChatModel,
+            prompt_template: str,
+            refine_template: str,
+
     ):
         """Initialize the data summarizer."""
         self.llm = llm
         self.llm_model = self.llm.model_name
-        prompt_template_bullet_point = (
-            "Write a summary of the following text.\n"
-            "TEXT: {text}\n"
-            "SUMMARY:\n"
-        )
-
-        prompt_bullet_point = PromptTemplate(
-            template=prompt_template_bullet_point, input_variables=["text"]
-        )
-        
-        refine_prompt_template_bullet_point = (
-            "Write a concise summary of the following text delimited by triple backquotes.\n"
-            "Return your response in bullet points which covers the key points of the text.\n"
-            " ```{text}```\n"
-            "BULLET POINT SUMMARY:\n"
-        )
-
-        refine_prompt_bullet_point = PromptTemplate(
-            template=refine_prompt_template_bullet_point, input_variables=["text"]
-        )
-
-        prompt_template = (
-            "Write a concise summary of the following:\n"
-            "{text}\n"
-            "CONCISE SUMMARY:\n"
-        )
-
-        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.\n"
-            "If the context isn't useful, return the original summary."
-        )
-        refine_prompt = PromptTemplate.from_template(refine_template)
-
-        self.prompt = prompt
-        self.refine_prompt = refine_prompt
-
-        self.prompt_bullet_point = prompt_bullet_point
-        self.refine_prompt_bullet_point = refine_prompt_bullet_point
+        self.prompt = PromptTemplate.from_template(prompt_template.strip())
+        self.refine_prompt = PromptTemplate.from_template(refine_template.strip())
 
     def get_summarization(self,
                           text: str,
                           chunk_num: int = 5,
-                          chunk_overlap: int = 30,
-                          bullet_point: bool = True) -> Dict:
+                          chunk_overlap: int = 30) -> Dict:
         """Get Summarization."""
-        if bullet_point:
-            prompt = self.prompt_bullet_point
-            refine_prompt = self.refine_prompt_bullet_point
-        else:
-            prompt = self.prompt
-            refine_prompt = self.refine_prompt
-
         text_splitter = TokenTextSplitter(
             chunk_size=self.token_limit[self.llm_model] // chunk_num,
             chunk_overlap=chunk_overlap,
@@ -98,8 +48,8 @@ class RefineDataSummarizer:
         chain = load_summarize_chain(
             llm=self.llm,
             chain_type="refine",
-            question_prompt=prompt,
-            refine_prompt=refine_prompt,
+            question_prompt=self.prompt,
+            refine_prompt=self.refine_prompt,
             return_intermediate_steps=True,
             input_key="input_documents",
             output_key="output_text",

utils/templates.py

@@ -0,0 +1,39 @@
+# regular template
+prompt_template = (
+    "Write a concise summary of the following:\n"
+    "{text}\n"
+    "CONCISE SUMMARY:\n"
+)
+
+
+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.\n"
+    "If the context isn't useful, return the original summary."
+)
+
+
+# bullet points template
+prompt_template_bullet_point = (
+    "Write a concise summary of the following in bullet points:\n"
+    "{text}\n"
+    "BULLET POINTS SUMMARY:\n"
+)
+
+refine_prompt_template_bullet_point = (
+    "Your job is to produce a final summary in bullet points\n"
+    "We have provided an existing bullet points summary up to a certain point: {existing_answer}\n"
+    "We have the opportunity to refine the existing bullet points"
+    "(only if needed) with some more context below.\n"
+    "------------\n"
+    "{text}\n"
+    "------------\n"
+    "Given the new context, refine the original bullet points summary.\n"
+    "If the context isn't useful, return the original bullet points."
+)