Probably wrong kind of prompt template used for token count calculation in longCEPO

[CoruNethron]

Hello. Thank you for your work.

In this line:
https://github.com/codelion/optillm/blob/12ac7863cf713c5cd417f81ec1e52ed28017bc16/optillm/plugins/longcepo/mapreduce.py#L224

We use max_context_window - tok_len(collapse_prompt) - max_output_tokens to estimate how much tokens we have to fit all the answers we previously had from MAP stage. But we probably want to estimate it as: max_context_window - tok_len(reduce_prompt) - max_output_tokens because we will actually pass combined answers to reduce prompt in the next step.

Here is a "pseudopatch" to demonstrate what I mean:

     num_tokens = get_prompt_length(format_chunk_list(context_chunks), tokenizer)
     token_budget = (
         longcepo_config.max_context_window
-        - get_prompt_length(longcepo_config.collapse_prompt, tokenizer)
+        - get_prompt_length(longcepo_config.reduce_prompt, tokenizer)
         - longcepo_config.max_output_tokens
     )
     logger.info(f"Pre-collapse length of chunks {num_tokens}, allowed {token_budget}")

[codelion]

CC: @ivanl-cerebras

[CoruNethron]

Also, it's a bit more awkward case, but probably we also want to apply this:

remove-irrelevant-during-collapse.patch

diff --git a/optillm/plugins/longcepo/mapreduce.py b/optillm/plugins/longcepo/mapreduce.py
index 4afec2b..84a204c 100644
--- a/optillm/plugins/longcepo/mapreduce.py
+++ b/optillm/plugins/longcepo/mapreduce.py
@@ -159,7 +159,6 @@ def mapreduce(
         cb_log,
         longcepo_config,
     )
-    result = remove_chunks(result, irrelevance_tags)
     if not result:
         return "No information", cb_log
 
@@ -200,6 +199,7 @@ def collapse_chunks(
     tokenizer,
     cb_log: CBLog,
     longcepo_config: LongCepoConfig,
+    irrelevance_tags: Tuple[str] = ("[NO INFORMATION]",),
 ) -> Tuple[List[str], CBLog]:
     """
     Collapses context chunk pairs in sliding window until the total token count fits within the context window.
@@ -269,6 +269,7 @@ def collapse_chunks(
             system_prompt,
             cb_log,
         )
+        context_chunks = remove_chunks(context_chunks, irrelevance_tags)
         merge_pair_idx = (merge_pair_idx + 1) % max(len(context_chunks) - 1, 1)
         num_tokens = get_prompt_length(format_chunk_list(context_chunks), tokenizer)
         collapse_step += 1

Because we allow in collapse_prompt for the model to answer with "[NO INFORMATION]", so to be sure, that model will not try to collapse "NO INFORMATION" and some actual answer, as it seems to be an unnecessary inference call in best case or may affect model to likely treat actual answer data to be irrelevant as well in worst case.