Using Ruby and LLMs: Implementing Map Reduce for Text Summarization

Edu Depetris

- Jul 28, 2024
  • Artificial Intelligence
  • Ai
  • Text Summarization
  • Open Ai
  • Ruby
We have a client who asked us to help summarize relatively small documents. Our solution is to use a Large Language Model (LLM) to achieve this.

The idea is to take documents containing a few pages and generate a short summary.

To simplify, let’s assume these documents are articles and we want to summarize each one.

After reading about the different techniques to summarize text with LLMs from LangChain, I decided to go with Map Reduce because it suits my scenario. However, if you have large documents, consider other techniques or adding an extra layer of reduction.

This technique involves taking a document, splitting it into smaller chunks of text, summarizing each chunk, and then creating the final summary from all these small summaries.

Let’s split our work into three parts:

 1. Split the article into small chunks of text.
 2. Summarize each chunk.
 3. Create the final summary from the individual summaries.

For simplicity, I’ll assume the document is already in a text format. However, you may need to handle the loading part, such as converting an article from XML, HTML, YAML, or any format into text. This can help with that.

Split the article into small chunks of text

There are many strategies to split long text into small chunks. This article provides an overview of different methods. I’ll use the recursive strategy.

We can use this ruby gem to get text chunks with the recursive strategy, but I’ll go a step further and use the Langchainrb ruby gem, which provides all the tools we need and more.

Let’s use the gem’s Loader, which will handle loading a file or path. This method accepts a "chunker" parameter where we can specify the chunking strategy. In our case, we’ll use the Recursive one.

The code will look like this:

path = "documents/article.txt"
chunker = Langchain::Chunker::RecursiveText

data_loaded = Langchain::Loader.new(path, chunker: chunker).load
data_in_chunks = data_loaded.chunks(chunk_size: 512, chunk_overlap: 80)

Summarize each chunk

Now that we have the data split into smaller chunks, let’s summarize it using an LLM.

To get better results, I’ll provide a prompt to help with my goal.

The prompt template will look like this:

def summarize_chunk_prompt(chunk_number, total_chunks, text)
  %(
    The following text is chunk #{chunk_number} of #{total_chunks} from the article.
    Please summarize the text.

    Original text:
    #{text}

    AI Summary:
  )
end

Let’s create an instance of the LLM and summarize the chunks:

llm = Langchain::LLM::OpenAI.new(api_key: ENV["OPENAI_API_KEY"])

summaries = data_in_chunks.map.with_index(1) do |chunk, index|
  text = summarize_chunk_prompt(index, data_in_chunks.size, chunk.text)
  
  llm.chat(messages: [{ role: "user", content: text }]).chat_completion
end

I’m using OpenAI as the LLM, but you can easily switch to Ollama or another one.

Create the final summary from the individual summaries

I’ll use another template for the final summary:

def final_summary_prompt(summaries:)
  %(
    The following are summarized chunks from an article.
    Please combine these summaries into one cohesive summary and provide the final summary.

    Summarized chunks:
    #{summaries}

    Final AI Summary
  )
end

Now, using the final prompt, let’s join the summaries and create the combined summary:

joined_summaries = summaries.join("\n\n")
combine_summaries = final_summary_prompt(summaries: joined_summaries)

The final step is to use the LLM to generate the final version of the summary:

llm.chat(messages: [{role: "user", content: combine_summaries}]).chat_completion


Here's the source code.

Happy Coding!