Improved examples [skip ci]

Author: ankane

examples/cohere/src/main.rs

@@ -22,7 +22,7 @@ fn main() -> Result<(), Box<dyn Error>> {
         "The cat is purring",
         "The bear is growling",
     ];
-    let embeddings = fetch_embeddings(&input, "search_document")?;
+    let embeddings = embed(&input, "search_document")?;
     for (content, embedding) in input.iter().zip(embeddings) {
         let embedding = Bit::from_bytes(&embedding);
         client.execute(
@@ -32,7 +32,7 @@ fn main() -> Result<(), Box<dyn Error>> {
     }
 
     let query = "forest";
-    let query_embedding = fetch_embeddings(&[query], "search_query")?;
+    let query_embedding = embed(&[query], "search_query")?;
     for row in client.query(
         "SELECT content FROM documents ORDER BY embedding <~> $1 LIMIT 5",
         &[&Bit::from_bytes(&query_embedding[0])],
@@ -44,7 +44,7 @@ fn main() -> Result<(), Box<dyn Error>> {
     Ok(())
 }
 
-fn fetch_embeddings(texts: &[&str], input_type: &str) -> Result<Vec<Vec<u8>>, Box<dyn Error>> {
+fn embed(texts: &[&str], input_type: &str) -> Result<Vec<Vec<u8>>, Box<dyn Error>> {
     let api_key = std::env::var("CO_API_KEY").or(Err("Set CO_API_KEY"))?;
 
     let response: Value = ureq::post("https://api.cohere.com/v1/embed")

examples/openai/src/main.rs

@@ -22,26 +22,28 @@ fn main() -> Result<(), Box<dyn Error>> {
         "The cat is purring",
         "The bear is growling",
     ];
-    let embeddings = fetch_embeddings(&input)?;
-
+    let embeddings = embed(&input)?;
     for (content, embedding) in input.iter().zip(embeddings) {
-        let embedding = Vector::from(embedding);
         client.execute(
             "INSERT INTO documents (content, embedding) VALUES ($1, $2)",
-            &[&content, &embedding],
+            &[&content, &Vector::from(embedding)],
         )?;
     }
 
-    let document_id = 2;
-    for row in client.query("SELECT content FROM documents WHERE id != $1 ORDER BY embedding <=> (SELECT embedding FROM documents WHERE id = $1) LIMIT 5", &[&document_id])? {
+    let query = "forest";
+    let query_embedding = embed(&[query])?.drain(..).next().unwrap();
+    for row in client.query(
+        "SELECT content FROM documents ORDER BY embedding <=> $1 LIMIT 5",
+        &[&Vector::from(query_embedding)],
+    )? {
         let content: &str = row.get(0);
         println!("{}", content);
     }
 
     Ok(())
 }
 
-fn fetch_embeddings(input: &[&str]) -> Result<Vec<Vec<f32>>, Box<dyn Error>> {
+fn embed(input: &[&str]) -> Result<Vec<Vec<f32>>, Box<dyn Error>> {
     let api_key = std::env::var("OPENAI_API_KEY").or(Err("Set OPENAI_API_KEY"))?;
 
     let response: Value = ureq::post("https://api.openai.com/v1/embeddings")