Search JARVIS's quotes

CMakeLists.txt

@@ -43,11 +43,6 @@ target_link_libraries(
     pthread
     dl
 )
-add_executable(
-    search_doc
-    examples/search_doc.cpp
-)
-target_link_libraries(search_doc embeddb)
 
 # Build CustomChar-core
 set(TARGET customchar-core)
@@ -71,6 +66,12 @@ target_include_directories(
 )
 target_link_libraries(${TARGET} PUBLIC ${SDL2_LIBRARIES} ${CMAKE_THREAD_LIBS_INIT} ${OpenCV_LIBS} whisper subprocess embeddb)
 
+add_executable(
+    search_doc
+    examples/search_doc.cpp
+)
+target_link_libraries(search_doc PUBLIC customchar-core)
+
 # CustomChar - cli
 add_executable(
     customchar-cli

customchar/llm/llm.cpp

@@ -20,6 +20,7 @@ LLM::LLM(const std::string& model_path, const std::string& path_session,
   lparams_.n_ctx = 2048;
   lparams_.seed = 1;
   lparams_.f16_kv = true;
+  lparams_.embedding = true;
 
   // Load model to ram
   ctx_llama_ = llama_init_from_file(model_path_.c_str(), lparams_);
@@ -349,3 +350,28 @@ std::string LLM::get_answer(const std::string& user_input) {
 
   return output_text;
 }
+
+std::vector<float> LLM::get_embedding(const std::string& text) {
+  std::vector<llama_token> embd(text.size());
+  llama_tokenize(ctx_llama_, text.c_str(), embd.data(), embd.size(), true);
+  llama_eval(ctx_llama_, embd.data(), embd.size(), n_past_, n_threads_);
+  const int n_embd = llama_n_embd(ctx_llama_);
+  const auto embeddings = llama_get_embeddings(ctx_llama_);
+  std::vector<float> result;
+  result.reserve(n_embd);
+  for (int i = 0; i < n_embd; ++i) {
+    result.push_back(embeddings[i]);
+  }
+
+  // Normalize
+  float norm = 0;
+  for (int i = 0; i < n_embd; ++i) {
+    norm += result[i] * result[i];
+  }
+  norm = sqrt(norm);
+  for (int i = 0; i < n_embd; ++i) {
+    result[i] /= norm;
+  }
+
+  return result;
+}

customchar/llm/llm.h

@@ -7,6 +7,7 @@
 #include <cassert>
 #include <cstdio>
 #include <fstream>
+#include <iostream>
 #include <regex>
 #include <string>
 #include <thread>
@@ -65,7 +66,7 @@ class LLM {
   void init_prompt();
 
  public:
-  /// @brief Constructor
+  // @brief Constructor
   /// @param model_path Path to the model
   /// @param path_session Path to the session
   LLM(const std::string& model_path, const std::string& path_session = "",
@@ -88,6 +89,32 @@ class LLM {
 
   /// @brief Get answer from LLM
   std::string get_answer(const std::string& user_input);
+
+  /// @brief Get embedding from LLM
+  std::vector<float> get_embedding(const std::string& text) {
+    std::vector<llama_token> embd(text.size());
+    llama_tokenize(ctx_llama_, text.c_str(), embd.data(), embd.size(), true);
+    llama_eval(ctx_llama_, embd.data(), embd.size(), n_past_, n_threads_);
+    const int n_embd = llama_n_embd(ctx_llama_);
+    const auto embeddings = llama_get_embeddings(ctx_llama_);
+    std::vector<float> result;
+    result.reserve(n_embd);
+    for (int i = 0; i < n_embd; ++i) {
+      result.push_back(embeddings[i]);
+    }
+
+    // Normalize
+    float norm = 0;
+    for (int i = 0; i < n_embd; ++i) {
+      norm += result[i] * result[i];
+    }
+    norm = sqrt(norm);
+    for (int i = 0; i < n_embd; ++i) {
+      result[i] /= norm;
+    }
+
+    return result;
+  }
 };
 
 }  // namespace llm

data/jarvis.txt

@@ -0,0 +1,6 @@
+Good morning. It's 7 A.M. The weather in Malibu is 72 degrees with scattered clouds. The surf conditions are fair with waist to shoulder highlines, high tide will be at 10:52 a.m. ―J.A.R.V.I.S
+We are now running on emergency backup power. ―J.A.R.V.I.S
+I don't want this winding up in the wrong hands. Maybe in mine, it could actually do some good.
+Wake up. Daddy's home.
+Welcome home, sir. Congratulations on the opening ceremonies. They were such a success, as was your senate hearing. And may I say how refreshing it is to finally see you in a video with your clothing on, sir.
+Sir, we will lose power before we penetrate that shell.

examples/search_doc.cpp

@@ -5,44 +5,88 @@
 #include "customchar/embeddb/document.h"
 #include "customchar/embeddb/embed_search.h"
 #include "customchar/embeddb/types.h"
+#include "customchar/llm/llm.h"
 
+using namespace CC;
 using namespace CC::embeddb;
 
+std::vector<std::string> read_lines(const std::string& file_path) {
+  std::vector<std::string> lines;
+
+  // Open the file
+  std::ifstream file(file_path);
+
+  // Check if the file was opened successfully
+  if (!file.is_open()) {
+    std::cerr << "Error: Unable to open the file." << std::endl;
+    return lines;
+  }
+
+  // Read line by line and add non-empty lines to the vector
+  std::string line;
+  while (std::getline(file, line)) {
+    if (!line.empty()) {
+      lines.push_back(line);
+    }
+  }
+
+  // Close the file after reading
+  file.close();
+
+  return lines;
+}
+
 int main() {
+  std::string model_path = "../models/llama-2-7b-chat.ggmlv3.q4_0.bin";
+  llm::LLM embedding_model(model_path);
+  embedding_model.eval_model();
+
   std::string connection_name = "test_collection";
   std::string path = "test_collection";
-  int dim = 10;
+  int dim = 4096;
   int max_size = 1000;
   Collection* collection = new Collection(connection_name, path, dim, max_size);
 
-  // Test Insert
-  std::vector<float> embed;
-  for (int i = 0; i < 10; i++) {
-    embed.push_back(i);
-  }
-  std::string content = "test content";
-  std::string meta = "test meta";
-  u_int32_t id = collection->insert_doc(embed, content, meta, 1, 1, 1);
-  std::cout << "Inserted document id: " << id << std::endl;
-
-  // Test Get Doc From Ids
-  std::vector<u_int32_t> ids{0, 1};
-  std::vector<Document> docs = collection->get_docs_by_ids(ids, 2);
-  std::cout << docs.size() << std::endl;
-
-  // Test Search
-  std::vector<float> query;
-  for (int i = 0; i < 10; i++) {
-    query.push_back(i);
-  }
-  std::vector<u_int32_t> doc_ids;
-  std::vector<float> distances;
-  int top_k = 2;
-  float threshold = 100;
-  collection->search(query, top_k, threshold, doc_ids, distances);
-  for (int i = 0; i < doc_ids.size(); i++) {
-    std::cout << doc_ids[i] << " " << distances[i] << std::endl;
+  // Read the document from file
+  std::string file_path = "../data/jarvis.txt";
+  std::vector<std::string> lines = read_lines(file_path);
+
+  // Insert all documents
+  for (int i = 0; i < lines.size(); i++) {
+    std::string content = lines[i];
+    std::vector<float> embed = embedding_model.get_embedding(lines[i]);
+    std::string meta = "test meta";
+    std::cout << "Inserting document " << i << std::endl;
+    std::cout << "Embedding size " << embed.size() << std::endl;
+    std::cout << "Content size " << content.size() << std::endl;
+    std::cout << "Content: " << content << std::endl;
+    u_int32_t id = collection->insert_doc(embed, content, meta, 1, 1, 1);
   }
 
+  while (true) {
+    // Test Search
+    std::string query_str;
+    std::cout << "Enter query: ";
+    std::getline(std::cin, query_str);
+    std::vector<float> query = embedding_model.get_embedding(query_str);
+
+    std::vector<u_int32_t> doc_ids;
+    std::vector<float> distances;
+    int top_k = 2;
+    float threshold = 100000000;
+    collection->search(query, top_k, threshold, doc_ids, distances);
+    for (int i = 0; i < 10; ++i) {
+      std::cout << query[i] << " ";
+    }
+    std::cout << std::endl;
+
+    std::cout << "Search result: " << std::endl;
+    for (int i = 0; i < doc_ids.size(); i++) {
+      std::cout << "Doc id: " << doc_ids[i] << std::endl;
+      std::cout << "Distance: " << distances[i] << std::endl;
+      Document doc = collection->get_doc(doc_ids[i]);
+      std::cout << "Content: " << doc.get_content() << std::endl;
+    }
+  }
   return 0;
 }