Skip to content

consolidating summaries

Peter edited this page Nov 14, 2024 · 4 revisions

Consolidating Summary Batches

aka the 'reduce' step in map-reduce summarization

If you have a document longer than your token_batch_length, the document will be split into batches of N tokens, and then summarized one by one. The model doesn't "know" about the summaries of the other batches when it is summarizing a given batch (they are completed independently to limit memory usage), so reading these multiple 'independent' summaries can be a bit repetitive.

This is where a secondary step can be used to consolidate (& improve) the results into one holistic summary.

Overview of Map-Reduce Summarization

image

The map-reduce summarization approach consists of two main steps:

  1. Map: Split a long document into chunks and summarize each chunk independently.
  2. Reduce: Consolidate these individual summaries into a coherent final summary

This approach is useful for handling documents that exceed the context window of the summarization model, while maintaining reasonable memory usage and processing speed.

Why use a new model?

While it may seem intuitive to use the same model for both steps, the "reduce" task is fundamentally different from the initial summarization. Long-context summarization models are great at summarizing large chunks of text, but they're not necessarily the best at combining multiple summarizations - just as a marathon runner may not be the best sprinter.

The bart-large-summary-map-reduce model is specifically trained for this shorter, more focused task of merging summaries, allowing it to be both more efficient and more effective at this specific task. This keeps memory usage reasonable while maintaining quality results.

Using a Map-Reduce Model

The pszemraj/bart-large-summary-map-reduce model is trained to create one coherent summary from several consecutive summarized document chunks. Here's how to use it:

import torch
from transformers import pipeline

pipe = pipeline(
    "text2text-generation",
    model="pszemraj/bart-large-summary-map-reduce",
    device_map="auto",
)

# Input the concatenated summaries from doc chunks (preferably with \n\n between)
summaries = "Summary 1...\n\nSummary 2... \n\nSummary 3..."

# Generate consolidated summary
res = pipe(
    summaries,
    max_new_tokens=512,  # increase up to 1024 if needed
    num_beams=4,
    early_stopping=True,
    truncation=True,
)
final_summary = res[0]["generated_text"]

BART, the text2text model used for the "reduce" step, supports up to 1024 input tokens. Since the inputs themselves are summaries in this use case, it should rarely (if ever) be an issue.

Integration with textsum

Open In Colab

The map-reduce consolidation can be integrated into the 'output' of a textsum workflow easily:

from textsum.summarize import Summarizer
from transformers import pipeline

# Initialize summarizers

chunk_summarizer = Summarizer(token_batch_length=8192)  # default model supports 16384
reduce_pipe = pipeline(
    "text2text-generation",
    model="pszemraj/bart-large-summary-map-reduce",
    device_map="auto",
)

text  # load a long string from somewhere

# textsum will automatically split your text into chunks of 8192
# summarize them, and concatenate them together with "\n\n" in between

combined_summaries = chunk_summarizer.summarize_string(text)

# Consolidate
final_summary = reduce_pipe(
    combined_summaries,
    max_new_tokens=512,
    num_beams=4,
    early_stopping=True,
    truncation=True,  # bart supports up to 1024 tokens input
)[0]["generated_text"]

print(final_summary)

Integration directly into the textsum package is possible. If this is something you are interested in, please open an issue and let me know!

Further Reading

Clone this wiki locally