diff --git a/cmake/vendor.cmake b/cmake/vendor.cmake index f47b82bf..c48d97c4 100644 --- a/cmake/vendor.cmake +++ b/cmake/vendor.cmake @@ -271,3 +271,12 @@ CPMAddPackage( GIT_TAG v3.12.0 EXCLUDE_FROM_ALL TRUE ) + +# usearch + +CPMAddPackage( + NAME usearch + GITHUB_REPOSITORY unum-cloud/usearch + GIT_TAG v2.17.7 + EXCLUDE_FROM_ALL TRUE +) diff --git a/profiler/CMakeLists.txt b/profiler/CMakeLists.txt index cabeb9db..023c5faf 100644 --- a/profiler/CMakeLists.txt +++ b/profiler/CMakeLists.txt @@ -61,6 +61,7 @@ set(SERVER_FILES TracyImGui.cpp TracyLlm.cpp TracyLlmApi.cpp + TracyLlmEmbeddings.cpp TracyLlmTools.cpp TracyMicroArchitecture.cpp TracyMouse.cpp @@ -221,7 +222,7 @@ else() endif() find_package(Threads REQUIRED) -target_link_libraries(${PROJECT_NAME} PRIVATE TracyServer TracyImGui Threads::Threads TracyLibcurl base64 tidy-static TracyPugixml nlohmann_json::nlohmann_json) +target_link_libraries(${PROJECT_NAME} PRIVATE TracyServer TracyImGui Threads::Threads TracyLibcurl base64 tidy-static TracyPugixml nlohmann_json::nlohmann_json usearch) target_include_directories(${PROJECT_NAME} PRIVATE ${tidy_SOURCE_DIR}/include) if(NOT DEFINED GIT_REV) diff --git a/profiler/src/profiler/TracyLlmEmbeddings.cpp b/profiler/src/profiler/TracyLlmEmbeddings.cpp new file mode 100644 index 00000000..bf51a522 --- /dev/null +++ b/profiler/src/profiler/TracyLlmEmbeddings.cpp @@ -0,0 +1,34 @@ +#include "TracyLlmEmbeddings.hpp" + +namespace tracy +{ + +TracyLlmEmbeddings::TracyLlmEmbeddings( size_t length, size_t reserve ) +{ + unum::usearch::metric_punned_t metric( length ); + m_index = unum::usearch::index_dense_t::make( metric ); + if( reserve > 0 ) m_index.reserve( reserve ); +} + +void TracyLlmEmbeddings::Add( std::string str, const std::vector& embedding ) +{ + m_index.add( m_data.size(), embedding.data() ); + m_data.emplace_back( std::move( str ) ); +} + +std::vector TracyLlmEmbeddings::Search( const std::vector& embedding, size_t k ) const +{ + std::vector ret; + auto result = m_index.search( embedding.data(), k ); + ret.reserve( result.size() ); + for( size_t i=0; i +#include +#include +#include + +namespace tracy +{ + +class TracyLlmEmbeddings +{ +public: + struct Result + { + size_t idx; + float distance; + }; + + explicit TracyLlmEmbeddings( size_t length, size_t reserve = 0 ); + + void Add( std::string str, const std::vector& embedding ); + [[nodiscard]] std::vector Search( const std::vector& embedding, size_t k ) const; + [[nodiscard]] const std::string& Get( size_t idx ) const { return m_data[idx]; } + +private: + unum::usearch::index_dense_t m_index; + std::vector m_data; +}; + +} + +#endif