From 50413524aedb8f716f4d7634fe2e92d2b06d4eb2 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?crypticG=C3=B8=C3=B8se?= <christogoosen@gmail.com>
Date: Wed, 20 Mar 2024 18:37:35 +0200
Subject: [PATCH] Add num_workers to raptor pack SummaryModule init, in
RaptorPack (#12106)
* Add num_workers to raptor pack SummaryModule init
* Update poetry version of package
* Fix spelling mistake
---
.../llama-index-packs-raptor/README.md | 27 +++++++++++++++++++
.../llama_index/packs/raptor/base.py | 9 +++++--
.../llama-index-packs-raptor/pyproject.toml | 2 +-
3 files changed, 35 insertions(+), 3 deletions(-)
diff --git a/llama-index-packs/llama-index-packs-raptor/README.md b/llama-index-packs/llama-index-packs-raptor/README.md
index 4b9b07fb5..761205f1e 100644
--- a/llama-index-packs/llama-index-packs-raptor/README.md
+++ b/llama-index-packs/llama-index-packs-raptor/README.md
@@ -76,3 +76,30 @@ retriever = RaptorRetriever([], ..., vector_store=vector_store)
```
Check out the [notebook here for complete details!](https://github.com/run-llama/llama_index/blob/main/llama-index-packs/llama-index-packs-raptor/examples/raptor.ipynb).
+
+## Configure Summary Module
+
+Using the SummaryModule you can configure how the Raptor Pack does summaries and how many workers are applied to summaries.
+
+You can configure the LLM.
+
+You can configure summary_prompt. This will change the prompt sent to your LLM to summarize you docs.
+
+You can configure num_workers, which will influence the number of workers or rather async semaphores allowing more summaries to process simulatneously.
+This might affect openai or other LLm provider API limits, be aware.
+
+```python
+from llama_index.packs.raptor.base import SummaryModule
+from llama_index.packs.raptor import RaptorRetriever
+
+summary_prompt = "As a professional summarizer, create a concise and comprehensive summary of the provided text, be it an article, post, conversation, or passage with as much detail as possible."
+
+# Adding SummaryModule you can configure the summary prompt and number of workers doing summaries.
+summary_module = SummaryModule(
+ llm=llm, summary_prompt=summary_prompt, num_workers=16
+)
+
+pack = RaptorPack(
+ documents, llm=llm, embed_model=embed_model, summary_module=summary_module
+)
+```
diff --git a/llama-index-packs/llama-index-packs-raptor/llama_index/packs/raptor/base.py b/llama-index-packs/llama-index-packs-raptor/llama_index/packs/raptor/base.py
index 89a539ee0..4dc1834b8 100644
--- a/llama-index-packs/llama-index-packs-raptor/llama_index/packs/raptor/base.py
+++ b/llama-index-packs/llama-index-packs-raptor/llama_index/packs/raptor/base.py
@@ -59,13 +59,18 @@ class SummaryModule(BaseModel):
arbitrary_types_allowed = True
def __init__(
- self, llm: Optional[LLM] = None, summary_prompt: str = DEFAULT_SUMMARY_PROMPT
+ self,
+ llm: Optional[LLM] = None,
+ summary_prompt: str = DEFAULT_SUMMARY_PROMPT,
+ num_workers: int = 4,
) -> None:
response_synthesizer = get_response_synthesizer(
response_mode="tree_summarize", use_async=True, llm=llm
)
super().__init__(
- response_synthesizer=response_synthesizer, summary_prompt=summary_prompt
+ response_synthesizer=response_synthesizer,
+ summary_prompt=summary_prompt,
+ num_workers=num_workers,
)
async def generate_summaries(
diff --git a/llama-index-packs/llama-index-packs-raptor/pyproject.toml b/llama-index-packs/llama-index-packs-raptor/pyproject.toml
index 3745f6e99..d786a2551 100644
--- a/llama-index-packs/llama-index-packs-raptor/pyproject.toml
+++ b/llama-index-packs/llama-index-packs-raptor/pyproject.toml
@@ -31,7 +31,7 @@ license = "MIT"
name = "llama-index-packs-raptor"
packages = [{include = "llama_index/"}]
readme = "README.md"
-version = "0.1.2"
+version = "0.1.3"
[tool.poetry.dependencies]
python = ">=3.9,<4.0"
--
GitLab