docs

examples/pdf/Cargo.toml

@@ -11,5 +11,4 @@ anyhow = "1.0.75"
 tokio = { version = "1.12.0", features = ["full"] }
 clap = "4.4.7"
 serde_json = "1.0.108"
-env_logger = "0.10.0"
-rayon = "1.8.0"
+env_logger = "0.10.0"

examples/pdf/src/main.rs

@@ -16,7 +16,6 @@ use orca::record::pdf;
 use orca::record::pdf::Pdf;
 use orca::record::Spin;
 use orca::{prompt, prompts};
-use rayon::prelude::*;
 use serde_json::json;
 
 #[derive(Parser, Debug)]
@@ -54,12 +53,8 @@ async fn main() -> Result<()> {
 
     let qdrant = Qdrant::new("localhost", 6334);
     if qdrant.create_collection(&collection, 384).await.is_ok() {
-        let mut embeddings = Vec::new();
-        for record in &pdf_records {
-            let embedding = bert.generate_embedding(prompt!(record)).await?;
-            embeddings.push(embedding.to_vec()?);
-        }
-        qdrant.insert_many(&collection, embeddings.clone(), pdf_records).await?;
+        let embeddings = bert.generate_embeddings(prompts!(&pdf_records)).await?;
+        qdrant.insert_many(&collection, embeddings.to_vec2()?, pdf_records).await?;
     }
 
     let query_embedding = bert.generate_embedding(prompt!(args.prompt)).await?;

orca/src/llm/mod.rs

@@ -52,28 +52,52 @@ pub trait LLM: Sync + Send {
 /// Embedding trait is used to generate an embedding from an Online Service.
 #[async_trait::async_trait]
 pub trait Embedding {
-    /// Generate an embedding from an Online Service.
+    /// Generate an embedding.
     /// # Arguments
-    /// * `input` - A Record trait object.
+    /// * `input` - Boxed prompt trait object.
+    ///
+    /// # Returns
+    /// * `EmbeddingResponse` - An embedding response.
     ///
     /// # Examples
     /// This example uses the OpenAI chat models.
     /// ```
-    /// use orca::prompt;
-    /// use orca::llm::Embedding;
-    /// use orca::record::Record;
-    /// use orca::llm::openai::OpenAI;
-    ///
-    /// #[tokio::main]
-    /// async fn main() {
-    ///    let client = OpenAI::new();
-    ///    let input = prompt!("Hello, world");
-    ///    let response = client.generate_embedding(input).await.unwrap();
-    ///    assert!(response.to_vec().expect("embedding is empty").len() > 0);
-    /// }
+    /// # use orca::prompt;
+    /// # use orca::llm::Embedding;
+    /// # use orca::llm::openai::OpenAI;
+    /// # #[tokio::main]
+    /// # async fn main() {
+    /// let client = OpenAI::new();
+    /// let input = prompt!("Hello, world");
+    /// let response = client.generate_embedding(input).await.unwrap();
+    /// assert!(response.to_vec().expect("embedding is empty").len() > 0);
+    /// # }
     /// ```
     async fn generate_embedding(&self, prompt: Box<dyn Prompt>) -> Result<EmbeddingResponse>;
 
+    /// Generate an embedding by batch
+    /// # Arguments
+    /// * `prompts` - A vector of boxed prompt trait objects.
+    ///
+    /// # Returns
+    /// * `EmbeddingResponse` - An embedding response.
+    ///
+    /// # Example
+    /// This example uses the Bert model.
+    /// ```
+    /// # use orca::prompts;
+    /// # use orca::llm::Embedding;
+    /// # use orca::llm::bert::Bert;
+    /// # use orca::prompt::Prompt;
+    /// # #[tokio::main]
+    /// # async fn main() {
+    /// let bert = Bert::new().build_model_and_tokenizer().await.unwrap();
+    /// let response = bert.generate_embeddings(prompts!("Hello World", "Goodbye World")).await;
+    /// let response = response.unwrap();
+    /// let vec = response.to_vec2().unwrap();
+    /// assert_eq!(vec.len(), 2);
+    /// # }
+    /// ````
     async fn generate_embeddings(&self, prompts: Vec<Box<dyn Prompt>>) -> Result<EmbeddingResponse>;
 }