May 10, 2024 · May 10, 2024
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 = "embed")]
 pub fn embed(transformer: &str, text: &str, kwargs: default!(JsonB, "'{}'")) -> Vec<f32> {
 embed_batch(transformer, Vec::from([text]),kwargs)
        .first()
 .unwrap()
    .to_vec()
 match crate::bindings::transformers::embed(transformer, vec![text], &kwargs.0) {
 Ok(output) => output.first().unwrap().to_vec(),
 Err(e) => error!("{e}"),
 }
 }

 #[cfg(all(feature = "python", not(feature = "use_as_lib")))]
 #[pg_extern(immutable, parallel_safe, name = "embed")]
 pub fn embed_batch(transformer: &str, inputs: Vec<&str>, kwargs: default!(JsonB, "'{}'")) -> Vec<Vec<f32>> {
 pub fn embed_batch(
    transformer: &str,
    inputs: Vec<&str>,
    kwargs: default!(JsonB, "'{}'"),
 ) -> SetOfIterator<'static, Vec<f32>> {
    match crate::bindings::transformers::embed(transformer, inputs, &kwargs.0) {
        Ok(output) => output,
        Ok(output) =>SetOfIterator::new(output.into_iter()),
        Err(e) => error!("{e}"),
    }
 }
Original file line number	Diff line number	Diff line change
Expand Up		@@ -589,17 +589,21 @@ fn load_dataset(
		#[cfg(all(feature = "python", not(feature = "use_as_lib")))]
		#[pg_extern(immutable, parallel_safe, name = "embed")]
		pub fn embed(transformer: &str, text: &str, kwargs: default!(JsonB, "'{}'")) -> Vec<f32> {
		embed_batch(transformer, Vec::from([text]),kwargs)
		.first()
		.unwrap()
		.to_vec()
		match crate::bindings::transformers::embed(transformer, vec![text], &kwargs.0) {
		Ok(output) => output.first().unwrap().to_vec(),
		Err(e) => error!("{e}"),
		}
		}

		#[cfg(all(feature = "python", not(feature = "use_as_lib")))]
		#[pg_extern(immutable, parallel_safe, name = "embed")]
		pub fn embed_batch(transformer: &str, inputs: Vec<&str>, kwargs: default!(JsonB, "'{}'")) -> Vec<Vec<f32>> {
		pub fn embed_batch(
		transformer: &str,
		inputs: Vec<&str>,
		kwargs: default!(JsonB, "'{}'"),
		) -> SetOfIterator<'static, Vec<f32>> {
		match crate::bindings::transformers::embed(transformer, inputs, &kwargs.0) {
		Ok(output) => output,
		Ok(output) =>SetOfIterator::new(output.into_iter()),
		Err(e) => error!("{e}"),
		}
		}
Expand Down