From fc081d9da6b4848ba41f602efdad00772617830a Mon Sep 17 00:00:00 2001 From: Roman Rizzi Date: Wed, 3 Jul 2024 18:10:31 -0300 Subject: [PATCH] FIX: Restore ability to fold summaries, which was accidentally removed (#700) --- lib/summarization/strategies/fold_content.rb | 107 +++++++++++++++++- .../strategies/fold_content_spec.rb | 32 ++++++ 2 files changed, 133 insertions(+), 6 deletions(-) diff --git a/lib/summarization/strategies/fold_content.rb b/lib/summarization/strategies/fold_content.rb index 38defd3a..36f76276 100644 --- a/lib/summarization/strategies/fold_content.rb +++ b/lib/summarization/strategies/fold_content.rb @@ -21,27 +21,122 @@ module DiscourseAi llm = DiscourseAi::Completions::Llm.proxy(completion_model.model_name) - summary_content = - content[:contents].map { |c| { ids: [c[:id]], summary: format_content_item(c) } } + initial_chunks = + rebalance_chunks( + llm.tokenizer, + content[:contents].map { |c| { ids: [c[:id]], summary: format_content_item(c) } }, + ) - { - summary: - summarize_single(llm, summary_content.first[:summary], user, opts, &on_partial_blk), - } + # Special case where we can do all the summarization in one pass. + if initial_chunks.length == 1 + { + summary: + summarize_single(llm, initial_chunks.first[:summary], user, opts, &on_partial_blk), + chunks: [], + } + else + summarize_chunks(llm, initial_chunks, user, opts, &on_partial_blk) + end end private + def summarize_chunks(llm, chunks, user, opts, &on_partial_blk) + # Safely assume we always have more than one chunk. + summarized_chunks = summarize_in_chunks(llm, chunks, user, opts) + total_summaries_size = + llm.tokenizer.size(summarized_chunks.map { |s| s[:summary].to_s }.join) + + if total_summaries_size < completion_model.available_tokens + # Chunks are small enough, we can concatenate them. + { + summary: + concatenate_summaries( + llm, + summarized_chunks.map { |s| s[:summary] }, + user, + &on_partial_blk + ), + chunks: summarized_chunks, + } + else + # We have summarized chunks but we can't concatenate them yet. Split them into smaller summaries and summarize again. + rebalanced_chunks = rebalance_chunks(llm.tokenizer, summarized_chunks) + + summarize_chunks(llm, rebalanced_chunks, user, opts, &on_partial_blk) + end + end + def format_content_item(item) "(#{item[:id]} #{item[:poster]} said: #{item[:text]} " end + def rebalance_chunks(tokenizer, chunks) + section = { ids: [], summary: "" } + + chunks = + chunks.reduce([]) do |sections, chunk| + if tokenizer.can_expand_tokens?( + section[:summary], + chunk[:summary], + completion_model.available_tokens, + ) + section[:summary] += chunk[:summary] + section[:ids] = section[:ids].concat(chunk[:ids]) + else + sections << section + section = chunk + end + + sections + end + + chunks << section if section[:summary].present? + + chunks + end + def summarize_single(llm, text, user, opts, &on_partial_blk) prompt = summarization_prompt(text, opts) llm.generate(prompt, user: user, feature_name: "summarize", &on_partial_blk) end + def summarize_in_chunks(llm, chunks, user, opts) + chunks.map do |chunk| + prompt = summarization_prompt(chunk[:summary], opts) + + chunk[:summary] = llm.generate( + prompt, + user: user, + max_tokens: 300, + feature_name: "summarize", + ) + chunk + end + end + + def concatenate_summaries(llm, summaries, user, &on_partial_blk) + prompt = DiscourseAi::Completions::Prompt.new(<<~TEXT.strip) + You are a summarization bot that effectively concatenates disjoint summaries, creating a cohesive narrative. + The narrative you create is in the form of one or multiple paragraphs. + Your reply MUST BE a single concatenated summary using the summaries I'll provide to you. + I'm NOT interested in anything other than the concatenated summary, don't include additional text or comments. + You understand and generate Discourse forum Markdown. + You format the response, including links, using Markdown. + TEXT + + prompt.push(type: :user, content: <<~TEXT.strip) + THESE are the summaries, each one separated by a newline, all of them inside XML tags: + + + #{summaries.join("\n")} + + TEXT + + llm.generate(prompt, user: user, &on_partial_blk) + end + def summarization_prompt(input, opts) insts = +<<~TEXT You are an advanced summarization bot that generates concise, coherent summaries of provided text. diff --git a/spec/lib/modules/summarization/strategies/fold_content_spec.rb b/spec/lib/modules/summarization/strategies/fold_content_spec.rb index df7f1298..0333dd45 100644 --- a/spec/lib/modules/summarization/strategies/fold_content_spec.rb +++ b/spec/lib/modules/summarization/strategies/fold_content_spec.rb @@ -32,5 +32,37 @@ RSpec.describe DiscourseAi::Summarization::Strategies::FoldContent do expect(result[:summary]).to eq(single_summary) end end + + context "when the content to summarize doesn't fit in a single call" do + it "summarizes each chunk and then concatenates them" do + content[:contents] << { poster: "asd2", id: 2, text: summarize_text } + + result = + DiscourseAi::Completions::Llm.with_prepared_responses( + [single_summary, single_summary, concatenated_summary], + ) { |spy| strategy.summarize(content, user).tap { expect(spy.completions).to eq(3) } } + + expect(result[:summary]).to eq(concatenated_summary) + end + + it "keeps splitting into chunks until the content fits into a single call to create a cohesive narrative" do + content[:contents] << { poster: "asd2", id: 2, text: summarize_text } + max_length_response = "(1 asd said: This is a text " + chunk_of_chunks = "I'm smol" + + result = + DiscourseAi::Completions::Llm.with_prepared_responses( + [ + max_length_response, + max_length_response, + chunk_of_chunks, + chunk_of_chunks, + concatenated_summary, + ], + ) { |spy| strategy.summarize(content, user).tap { expect(spy.completions).to eq(5) } } + + expect(result[:summary]).to eq(concatenated_summary) + end + end end end