From 4807f73128de86e76be1369a690516fdb1d11f45 Mon Sep 17 00:00:00 2001
From: Santiago Medina <santiagm08@gmail.com>
Date: Tue, 31 Oct 2023 20:56:53 -0700
Subject: [PATCH] docs

---
 examples/pdf/Cargo.toml  |  3 +--
 examples/pdf/src/main.rs |  9 ++-----
 orca/src/llm/mod.rs      | 52 +++++++++++++++++++++++++++++-----------
 3 files changed, 41 insertions(+), 23 deletions(-)

diff --git a/examples/pdf/Cargo.toml b/examples/pdf/Cargo.toml
index 534960d..0e09ce1 100644
--- a/examples/pdf/Cargo.toml
+++ b/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"
\ No newline at end of file
diff --git a/examples/pdf/src/main.rs b/examples/pdf/src/main.rs
index 0a66fc9..e3bbe3b 100644
--- a/examples/pdf/src/main.rs
+++ b/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?;
diff --git a/orca/src/llm/mod.rs b/orca/src/llm/mod.rs
index 68868bb..c3b0437 100644
--- a/orca/src/llm/mod.rs
+++ b/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>;
 }