Nov 23, 2023 · Nov 25, 2023 · Nov 25, 2023 · Nov 25, 2023 · Nov 25, 2023
diff --git a/pgml-extension/src/api.rs b/pgml-extension/src/api.rs
    }
 }

 #[cfg(all(feature = "python", not(feature = "use_as_lib")))]
 #[pg_extern(immutable, parallel_safe, name = "embed2")]
 pub fn embed_batch2<'a>(
    transformer: &str,
    inputs: Vec<&'a str>,
    kwargs: default!(JsonB, "'{}'"),
 ) -> TableIterator<'a, (name!(text, String), name!(embedding, Vec<f32>))> {
    let rows = match crate::bindings::transformers::embed(transformer, inputs.clone(), &kwargs.0) {
        Ok(rows) => rows,
        Err(e) => {
            error!("{e}");
        }
    };
    TableIterator::new(
        inputs.into_iter().zip(rows.into_iter()).map(|(text, embedding)| {
            (text.to_string(), embedding)
        }),
    )
 }

 /// Clears the GPU cache.
 ///
 /// # Arguments
Original file line number	Diff line number	Diff line change
Expand Up		@@ -558,6 +558,26 @@ pub fn embed_batch(
		}
		}

		#[cfg(all(feature = "python", not(feature = "use_as_lib")))]
		#[pg_extern(immutable, parallel_safe, name = "embed2")]
		pub fn embed_batch2<'a>(
Copy link Contributor montanalowNov 25, 2023• edited Loading Choose a reason for hiding this comment The reason will be displayed to describe this comment to others.Learn more. My rough thoughts, without running the code on some examples. I think we should name this SQL`embed_3` and Rust`embed_batch_3` with the goal of establishing this as the 3.0`embed` API, as well as a pattern for releasing 3.0 APIs early as we developing them in an alpha state (with potentially breaking changes, where we completely drop them in 3.1 in favor of the newly established default behavior). Your example convinces me that batch APIs should return a table, but I think that table's rows should be`JSONB` with {id, embedding} keys (at least), unless there is a significant performance implication on that front. My thinking is that embedding models are getting more complicated and now some take JSON rather than TEXT for`inputs` including a`prompt`. It would be nice to have an optional`id` in the input JSON, and if it's not present, then just return the entire input JSON as the`id`, which acts just like your TEXT as the key. Final thought is that`kwargs` is JSONB currently, which works well with the underlying Python dependencies, but I'd like to structure it as much as possible for final 3.0. We should find a way to flag this obviously as an alpha API, that will be broken and eventually dropped when a final version is available.
		transformer: &str,
		inputs: Vec<&'a str>,
		kwargs: default!(JsonB, "'{}'"),
		) -> TableIterator<'a, (name!(text, String), name!(embedding, Vec<f32>))> {
		let rows = match crate::bindings::transformers::embed(transformer, inputs.clone(), &kwargs.0) {
		Ok(rows) => rows,
		Err(e) => {
		error!("{e}");
		}
		};
		TableIterator::new(
		inputs.into_iter().zip(rows.into_iter()).map(\|(text, embedding)\| {
		(text.to_string(), embedding)
		}),
		)
		}

		/// Clears the GPU cache.
		///
		/// # Arguments
Expand Down