refactor: Updated embed(...) and rerank() signatures (#201)

Signed-off-by: Anush008 <[email protected]>

Author: Anush008

src/image_embedding/impl.rs

@@ -157,13 +157,16 @@ impl ImageEmbedding {
         Ok(output)
     }
 
-    /// Method to generate image embeddings for a Vec of image path
-    // Generic type to accept String, &str, OsString, &OsStr
+    /// Method to generate image embeddings for a collection of image paths.
+    ///
+    /// Accepts anything that can be referenced as a slice of elements implementing
+    /// [`AsRef<Path>`], such as `Vec<String>`, `Vec<PathBuf>`, `&[&str]`, or `&[&Path]`.
     pub fn embed<S: AsRef<Path> + Send + Sync>(
         &mut self,
-        images: Vec<S>,
+        images: impl AsRef<[S]>,
         batch_size: Option<usize>,
     ) -> anyhow::Result<Vec<Embedding>> {
+        let images = images.as_ref();
         // Determine the batch size, default if not specified
         let batch_size = batch_size.unwrap_or(DEFAULT_BATCH_SIZE);
 

src/reranking/impl.rs

@@ -122,13 +122,16 @@ impl TextRerank {
     }
 
     /// Rerank documents using the reranker model and returns the results sorted by score in descending order.
+    ///
+    /// Accepts a query and a collection of documents implementing [`AsRef<str>`].
     pub fn rerank<S: AsRef<str> + Send + Sync>(
         &mut self,
         query: S,
-        documents: Vec<S>,
+        documents: impl AsRef<[S]>,
         return_documents: bool,
         batch_size: Option<usize>,
     ) -> Result<Vec<RerankResult>> {
+        let documents = documents.as_ref();
         let batch_size = batch_size.unwrap_or(DEFAULT_BATCH_SIZE);
         let q = query.as_ref();
 

src/sparse_text_embedding/impl.rs

@@ -104,13 +104,16 @@ impl SparseTextEmbedding {
             .expect("Model not found in supported models list. This is a bug - please report it.")
     }
 
-    /// Method to generate sentence embeddings for a Vec of texts
-    // Generic type to accept String, &str, OsString, &OsStr
+    /// Method to generate sentence embeddings for a collection of texts.
+    ///
+    /// Accepts anything that can be referenced as a slice of elements implementing
+    /// [`AsRef<str>`], such as `Vec<String>`, `Vec<&str>`, `&[String]`, or `&[&str]`.
     pub fn embed<S: AsRef<str> + Send + Sync>(
         &mut self,
-        texts: Vec<S>,
+        texts: impl AsRef<[S]>,
         batch_size: Option<usize>,
     ) -> Result<Vec<SparseEmbedding>> {
+        let texts = texts.as_ref();
         // Determine the batch size, default if not specified
         let batch_size = batch_size.unwrap_or(DEFAULT_BATCH_SIZE);
 

src/text_embedding/impl.rs

@@ -265,9 +265,10 @@ impl TextEmbedding {
     /// embeddings with your custom output type.
     pub fn transform<S: AsRef<str> + Send + Sync>(
         &mut self,
-        texts: Vec<S>,
+        texts: impl AsRef<[S]>,
         batch_size: Option<usize>,
     ) -> Result<EmbeddingOutput> {
+        let texts = texts.as_ref();
         // Determine the batch size according to the quantization method used.
         // Default if not specified
         let batch_size = match self.quantization {
@@ -360,10 +361,10 @@ impl TextEmbedding {
         Ok(EmbeddingOutput::new(batches))
     }
 
-    /// Method to generate sentence embeddings for a Vec of texts.
+    /// Method to generate sentence embeddings for a collection of texts.
     ///
-    /// Accepts a [`Vec`] consisting of elements of either [`String`], &[`str`],
-    /// [`std::ffi::OsString`], &[`std::ffi::OsStr`].
+    /// Accepts anything that can be referenced as a slice of elements implementing
+    /// [`AsRef<str>`], such as `Vec<String>`, `Vec<&str>`, `&[String]`, or `&[&str]`.
     ///
     /// The output is a [`Vec`] of [`Embedding`]s.
     ///
@@ -373,10 +374,10 @@ impl TextEmbedding {
     /// the default output precedence and array transformer for the [`TextEmbedding`] model.
     pub fn embed<S: AsRef<str> + Send + Sync>(
         &mut self,
-        texts: Vec<S>,
+        texts: impl AsRef<[S]>,
         batch_size: Option<usize>,
     ) -> Result<Vec<Embedding>> {
-        let batches = self.transform(texts, batch_size)?;
+        let batches = self.transform(texts.as_ref(), batch_size)?;
         if let Some(output_key) = &self.output_key {
             batches.export_with_transformer(output::transformer_with_precedence(
                 output_key,