From de7ed724263fa6fc9f922da96acae33d44b275f7 Mon Sep 17 00:00:00 2001
From: SilasMarvin <19626586+SilasMarvin@users.noreply.github.com>
Date: Mon, 13 May 2024 12:06:47 -0700
Subject: [PATCH] Added batch embeddings in the SDK

---
 pgml-sdks/pgml/src/queries.rs | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/pgml-sdks/pgml/src/queries.rs b/pgml-sdks/pgml/src/queries.rs
index 1ea7001bf..3fcdb9c31 100644
--- a/pgml-sdks/pgml/src/queries.rs
+++ b/pgml-sdks/pgml/src/queries.rs
@@ -217,9 +217,9 @@ ON CONFLICT (chunk_id) DO UPDATE SET ts = EXCLUDED.ts;
 pub const GENERATE_EMBEDDINGS_FOR_CHUNK_IDS: &str = r#"
 INSERT INTO %s (chunk_id, embedding) 
 SELECT 
-  id, 
+  unnest(array_agg(id)), 
   pgml.embed(
-    text => chunk, 
+    inputs => array_agg(chunk), 
     transformer => $1, 
     kwargs => $2 
   )