From 36a261b0b23535111b6ccb9ff85a344cbc6b5254 Mon Sep 17 00:00:00 2001
From: FuexFollets <fuex9follets@gmail.com>
Date: Thu, 8 Feb 2024 10:17:04 -0500
Subject: [PATCH] Added nanobench segments in text prediction

---
 tests/vector_database.cpp          | 8 ++++++--
 tests/vector_database_synonyms.cpp | 5 ++++-
 2 files changed, 10 insertions(+), 3 deletions(-)

diff --git a/tests/vector_database.cpp b/tests/vector_database.cpp
index 3f8cb45..f2faa7d 100644
--- a/tests/vector_database.cpp
+++ b/tests/vector_database.cpp
@@ -3,6 +3,8 @@
 #include <string>
 #include <vector>
 
+#include <nanobench.h>
+
 #include <lexocraft/fancy_eigen_print.hpp>
 #include <lexocraft/llm/vector_database.hpp>
 
@@ -24,7 +26,7 @@ int main(int argc, char** argv) {
     if (args.at(0) == "create") {
         const std::string path_from = args.at(1);
         const std::string path_to = args.at(2);
-        const int total_trees = (args.size() > 4) ? std::stoi(args.at(3)) : 50;
+        const int total_trees = (args.size() > 3) ? std::stoi(args.at(3)) : 10;
 
         std::ifstream file {path_from};
         std::vector<lc::WordVector> words;
@@ -41,7 +43,9 @@ int main(int argc, char** argv) {
         lc::VectorDatabase database {words};
 
         std::cout << "Building Annoy index\n";
-        database.build_annoy_index(total_trees);
+        ankerl::nanobench::Bench().run("build_annoy_index", [&] {
+            ankerl::nanobench::doNotOptimizeAway(database.build_annoy_index(total_trees));
+        });
         std::cout << "Annoy index built\n";
 
         database.save_file(path_to);
diff --git a/tests/vector_database_synonyms.cpp b/tests/vector_database_synonyms.cpp
index e0328a5..a95a0df 100644
--- a/tests/vector_database_synonyms.cpp
+++ b/tests/vector_database_synonyms.cpp
@@ -66,7 +66,10 @@ int main(int argc, char** argv) {
         database = lc::VectorDatabase {words};
 
         std::cout << "Building Annoy index\n";
-        database.build_annoy_index(2);
+        ankerl::nanobench::Bench().run("build_annoy_index", [&] {
+            database.build_annoy_index(10);
+        });
+
         std::cout << "Annoy index built\n";
     }