DEV: Backfill embeddings concurrently.

We are adding a new method for generating and storing embeddings in bulk, which relies on `Concurrent::Promises::Future`. Generating an embedding consists of three steps:

Prepare text
HTTP call to retrieve the vector
Save to DB.
Each one is independently executed on whatever thread the pool gives us.

We are bringing a custom thread pool instead of the global executor since we want control over how many threads we spawn to limit concurrency. We also avoid firing thousands of HTTP requests when working with large batches.

Author: romanrizzi

app/jobs/scheduled/embeddings_backfill.rb

@@ -75,9 +75,9 @@ def execute(args)
       # First, we'll try to backfill embeddings for posts that have none
       posts
         .where("#{table_name}.post_id IS NULL")
-        .find_each do |t|
-          vector_rep.generate_representation_from(t)
-          rebaked += 1
+        .find_in_batches do |batch|
+          vector_rep.gen_bulk_reprensentations(batch)
+          rebaked += batch.size
         end
 
       return if rebaked >= limit
@@ -90,24 +90,28 @@ def execute(args)
           OR
           #{table_name}.strategy_version < #{strategy.version}
         SQL
-        .find_each do |t|
-          vector_rep.generate_representation_from(t)
-          rebaked += 1
+        .find_in_batches do |batch|
+          vector_rep.gen_bulk_reprensentations(batch)
+          rebaked += batch.size
         end
 
       return if rebaked >= limit
 
       # Finally, we'll try to backfill embeddings for posts that have outdated
       # embeddings due to edits. Here we only do 10% of the limit
-      posts
-        .where("#{table_name}.updated_at < ?", 7.days.ago)
-        .order("random()")
-        .limit((limit - rebaked) / 10)
-        .pluck(:id)
-        .each do |id|
-          vector_rep.generate_representation_from(Post.find_by(id: id))
-          rebaked += 1
-        end
+      posts_batch_size = 1000
+
+      outdated_post_ids =
+        posts
+          .where("#{table_name}.updated_at < ?", 7.days.ago)
+          .order("random()")
+          .limit((limit - rebaked) / 10)
+          .pluck(:id)
+
+      outdated_post_ids.each_slice(posts_batch_size) do |batch|
+        vector_rep.gen_bulk_reprensentations(Post.where(id: batch).order("topics.bumped_at DESC"))
+        rebaked += batch.length
+      end
 
       rebaked
     end
@@ -120,14 +124,13 @@ def populate_topic_embeddings(vector_rep, topics, force: false)
       topics = topics.where("#{vector_rep.topic_table_name}.topic_id IS NULL") if !force
 
       ids = topics.pluck("topics.id")
+      batch_size = 1000
 
-      ids.each do |id|
-        topic = Topic.find_by(id: id)
-        if topic
-          vector_rep.generate_representation_from(topic)
-          done += 1
-        end
+      ids.each_slice(batch_size) do |batch|
+        vector_rep.gen_bulk_reprensentations(Topic.where(id: batch).order("topics.bumped_at DESC"))
+        done += batch.length
       end
+
       done
     end
   end

lib/embeddings/vector_representations/base.rb

@@ -50,8 +50,37 @@ def vector_from(text, asymetric: false)
           raise NotImplementedError
         end
 
+        def gen_bulk_reprensentations(relation)
+          raw_inputs = relation.map { |record| { target: record, text: prepare_text(record) } }
+
+          pool_size = 10
+          pool =
+            Concurrent::CachedThreadPool.new(min_threads: 0, max_threads: pool_size, idletime: 30)
+
+          embedding_gen = inference_client
+          db = RailsMultisite::ConnectionManagement.current_db
+          promised_embeddings =
+            raw_inputs.map do |raw|
+              Concurrent::Promises
+                .fulfilled_future(raw, pool)
+                .then_on(pool) do |w_prepared_text|
+                  w_prepared_text.merge(
+                    embedding: embedding_gen.perform!(w_prepared_text[:text]),
+                    digest: OpenSSL::Digest::SHA1.hexdigest(w_prepared_text[:text]),
+                  )
+                end
+                .then_on(pool) do |w_embedding|
+                  RailsMultisite::ConnectionManagement.with_connection(db) do
+                    save_to_db(w_embedding[:target], w_embedding[:embedding], w_embedding[:digest])
+                  end
+                end
+            end
+
+          Concurrent::Promises.zip(*promised_embeddings).value!
+        end
+
         def generate_representation_from(target, persist: true)
-          text = @strategy.prepare_text_from(target, tokenizer, max_sequence_length - 2)
+          text = prepare_text(target)
           return if text.blank?
 
           target_column =
@@ -429,6 +458,10 @@ def save_to_db(target, vector, digest)
         def inference_client
           raise NotImplementedError
         end
+
+        def prepare_text(record)
+          @strategy.prepare_text_from(record, tokenizer, max_sequence_length - 2)
+        end
       end
     end
   end

lib/embeddings/vector_representations/multilingual_e5_large.rb

@@ -34,7 +34,7 @@ def vector_from(text, asymetric: false)
           needs_truncation = client.class.name.include?("HuggingFaceTextEmbeddings")
           if needs_truncation
             text = tokenizer.truncate(text, max_sequence_length - 2)
-          else
+          elsif !text.starts_with?("query:")
             text = "query: #{text}"
           end
 
@@ -79,6 +79,14 @@ def inference_client
             raise "No inference endpoint configured"
           end
         end
+
+        def prepare_text(record)
+          if inference_client.class.name.include?("DiscourseClassifier")
+            return "query: #{super(record)}"
+          end
+
+          super(record)
+        end
       end
     end
   end

spec/lib/modules/embeddings/vector_representations/vector_rep_shared_examples.rb

@@ -1,14 +1,15 @@
 # frozen_string_literal: true
 
 RSpec.shared_examples "generates and store embedding using with vector representation" do
-  before { @expected_embedding = [0.0038493] * vector_rep.dimensions }
+  let(:expected_embedding_1) { [0.0038493] * vector_rep.dimensions }
+  let(:expected_embedding_2) { [0.0037684] * vector_rep.dimensions }
 
   describe "#vector_from" do
     it "creates a vector from a given string" do
       text = "This is a piece of text"
-      stub_vector_mapping(text, @expected_embedding)
+      stub_vector_mapping(text, expected_embedding_1)
 
-      expect(vector_rep.vector_from(text)).to eq(@expected_embedding)
+      expect(vector_rep.vector_from(text)).to eq(expected_embedding_1)
     end
   end
 
@@ -24,11 +25,11 @@
           vector_rep.tokenizer,
           vector_rep.max_sequence_length - 2,
         )
-      stub_vector_mapping(text, @expected_embedding)
+      stub_vector_mapping(text, expected_embedding_1)
 
       vector_rep.generate_representation_from(topic)
 
-      expect(vector_rep.topic_id_from_representation(@expected_embedding)).to eq(topic.id)
+      expect(vector_rep.topic_id_from_representation(expected_embedding_1)).to eq(topic.id)
     end
 
     it "creates a vector from a post and stores it in the database" do
@@ -38,11 +39,45 @@
           vector_rep.tokenizer,
           vector_rep.max_sequence_length - 2,
         )
-      stub_vector_mapping(text, @expected_embedding)
+      stub_vector_mapping(text, expected_embedding_1)
 
       vector_rep.generate_representation_from(post)
 
-      expect(vector_rep.post_id_from_representation(@expected_embedding)).to eq(post.id)
+      expect(vector_rep.post_id_from_representation(expected_embedding_1)).to eq(post.id)
+    end
+  end
+
+  describe "#gen_bulk_reprensentations" do
+    fab!(:topic) { Fabricate(:topic) }
+    fab!(:post) { Fabricate(:post, post_number: 1, topic: topic) }
+    fab!(:post2) { Fabricate(:post, post_number: 2, topic: topic) }
+
+    fab!(:topic_2) { Fabricate(:topic) }
+    fab!(:post_2_1) { Fabricate(:post, post_number: 1, topic: topic_2) }
+    fab!(:post_2_2) { Fabricate(:post, post_number: 2, topic: topic_2) }
+
+    it "creates a vector for each object in the relation" do
+      text =
+        truncation.prepare_text_from(
+          topic,
+          vector_rep.tokenizer,
+          vector_rep.max_sequence_length - 2,
+        )
+
+      text2 =
+        truncation.prepare_text_from(
+          topic_2,
+          vector_rep.tokenizer,
+          vector_rep.max_sequence_length - 2,
+        )
+
+      stub_vector_mapping(text, expected_embedding_1)
+      stub_vector_mapping(text2, expected_embedding_2)
+
+      vector_rep.gen_bulk_reprensentations(Topic.where(id: [topic.id, topic_2.id]))
+
+      expect(vector_rep.topic_id_from_representation(expected_embedding_1)).to eq(topic.id)
+      expect(vector_rep.topic_id_from_representation(expected_embedding_1)).to eq(topic.id)
     end
   end
 
@@ -58,7 +93,7 @@
           vector_rep.tokenizer,
           vector_rep.max_sequence_length - 2,
         )
-      stub_vector_mapping(text, @expected_embedding)
+      stub_vector_mapping(text, expected_embedding_1)
       vector_rep.generate_representation_from(topic)
 
       expect(