From e15956a98c138b8901853e4b7eb1b863d516ad7d Mon Sep 17 00:00:00 2001
From: Johannes Wesch <johannes.wesch@aleph-alpha.com>
Date: Tue, 6 Feb 2024 09:28:03 +0100
Subject: [PATCH] added break for recursive summary when number of partial
 summary doesn't change

---
 src/intelligence_layer/core/chunk.py          |  4 +---
 .../summarize/recursive_summarize.py          |  8 +++++---
 tests/core/test_chunk.py                      |  5 +++--
 .../summarize/test_recursive_summarize.py     | 20 +++++++++++++++++++
 4 files changed, 29 insertions(+), 8 deletions(-)

diff --git a/src/intelligence_layer/core/chunk.py b/src/intelligence_layer/core/chunk.py
index 0bdd6b566..92b3d8c2f 100644
--- a/src/intelligence_layer/core/chunk.py
+++ b/src/intelligence_layer/core/chunk.py
@@ -96,9 +96,7 @@ def __init__(
                     overlap_length_tokens, max_tokens_per_chunk
                 )
             )
-        self.chunk_task = ChunkTask(
-            client, model, overlap_length_tokens // 2
-        )
+        self.chunk_task = ChunkTask(client, model, overlap_length_tokens // 2)
         self.tokenizer = client.tokenizer(model)
         self.max_tokens_per_chunk = max_tokens_per_chunk
         self.overlap_length_tokens = overlap_length_tokens
diff --git a/src/intelligence_layer/use_cases/summarize/recursive_summarize.py b/src/intelligence_layer/use_cases/summarize/recursive_summarize.py
index 80cebabdd..0c2e4e7ab 100644
--- a/src/intelligence_layer/use_cases/summarize/recursive_summarize.py
+++ b/src/intelligence_layer/use_cases/summarize/recursive_summarize.py
@@ -27,20 +27,22 @@ def __init__(
     def do_run(
         self, input: LongContextSummarizeInput, task_span: TaskSpan
     ) -> SummarizeOutput:
+        num_partial_summaries = 0
         text = input.text
-        loop_count = 0
         while True:
             summarize_output = self.long_context_summarize_task.run(
                 LongContextSummarizeInput(text=text, language=input.language), task_span
             )
+            if num_partial_summaries == len(summarize_output.partial_summaries):
+                break
+            num_partial_summaries = len(summarize_output.partial_summaries)
+
             num_generated_tokens = 0
             text = ""
             for partial_summary in summarize_output.partial_summaries:
                 num_generated_tokens += partial_summary.generated_tokens
                 text += partial_summary.summary + "\n"
 
-            loop_count += 1
-
             if len(summarize_output.partial_summaries) == 1:
                 break
 
diff --git a/tests/core/test_chunk.py b/tests/core/test_chunk.py
index b2c6ca1bc..888e5d912 100644
--- a/tests/core/test_chunk.py
+++ b/tests/core/test_chunk.py
@@ -37,14 +37,15 @@ def test_overlapped_chunking(
         print(first)
 
         assert (
-            len(first) <= MAX_TOKENS + 2
+            len(first)
+            <= MAX_TOKENS + 2
             # `+2` because re-tokenizing the chunk can add a few extra tokens at
             # the beginning or end of each chunk. This is a hack.
         )
         next = output_tokenized[chunk_index + 1].tokens
 
         found = False
-        for offset in range(len(next)-OVERLAP//2):
+        for offset in range(len(next) - OVERLAP // 2):
             if first[-OVERLAP // 2 :] != next[offset : offset + OVERLAP // 2]:
                 continue
             found = True
diff --git a/tests/use_cases/summarize/test_recursive_summarize.py b/tests/use_cases/summarize/test_recursive_summarize.py
index b0e95df1e..b1fd7412f 100644
--- a/tests/use_cases/summarize/test_recursive_summarize.py
+++ b/tests/use_cases/summarize/test_recursive_summarize.py
@@ -4,12 +4,18 @@
 from aleph_alpha_client import Client, CompletionRequest, CompletionResponse
 from pytest import fixture
 
+from intelligence_layer.connectors.limited_concurrency_client import (
+    AlephAlphaClientProtocol,
+)
 from intelligence_layer.core import NoOpTracer
 from intelligence_layer.use_cases import (
     LongContextHighCompressionSummarize,
     LongContextSummarizeInput,
     RecursiveSummarize,
 )
+from intelligence_layer.use_cases.summarize.steerable_long_context_summarize import (
+    SteerableLongContextSummarize,
+)
 
 
 class RecursiveCountingClient(Client):
@@ -52,6 +58,20 @@ def test_recursive_summarize_stops_when_hitting_max_tokens(
     assert "new orleans" in output.summary.lower()
 
 
+def test_recursive_summarize_stops_when_num_partial_summaries_stays_same(
+    client: AlephAlphaClientProtocol,
+) -> None:
+    max_tokens = None
+    slcs = SteerableLongContextSummarize(
+        client, model="luminous-base", max_generated_tokens=75, max_tokens_per_chunk=145
+    )
+    input = LongContextSummarizeInput(text=short_text, max_tokens=max_tokens)
+    task = RecursiveSummarize(slcs)
+    output = task.run(input, NoOpTracer())
+
+    assert output.generated_tokens > 145
+
+
 def test_recursive_summarize_stops_after_one_chunk(
     recursive_counting_client: RecursiveCountingClient,
 ) -> None: