Hnsw.cpp

#include "Hnsw.hpp"
#include <algorithm>
#include <cmath>
#include <numeric>
#include <random>

Hnsw::Hnsw(int max_connections, int vec_size) {
  this->max_connections = max_connections;
  this->vec_size = vec_size;
  this->highest_layer = 0;
  this->current_id = 0;
}

int Hnsw::_assign_layer() {
  std::random_device dev;
  std::mt19937 rng(dev());
  std::uniform_real_distribution<float> dist(0.0, 1.0);

  int layer = 0;
  while (dist(rng) > 0.5)
    layer++;
  return layer;
}

int Hnsw::insert(std::vector<float> new_embedding, int document_id,
                 std::string document_title) {
  if (new_embedding.size() != this->vec_size)
    throw "Error: Inconsistent vector size";
  int max_layer = this->_assign_layer();
  std::vector<int> vector_layers(max_layer + 1);
  std::iota(vector_layers.begin(), vector_layers.end(), 0);

  std::vector<std::vector<int>> candidates_by_layer;
  int vector_id = this->current_id;
  Embedding embedding =
      Embedding(vector_id, new_embedding.size(), new_embedding);
  if (this->current_id > 0)
    candidates_by_layer = this->_search(embedding, 50);

  this->highest_layer = std::max(max_layer, this->highest_layer);

  this->vec2embeds.push_back(embedding);

  Document d = Document(document_id, document_title);
  this->vec2docs.push_back(d);

  this->vec2neighbors.push_back(std::vector<std::vector<int>>());
  for (int layer : vector_layers)
    this->vec2neighbors[vector_id].push_back(std::vector<int>());

  if (candidates_by_layer.size() != 0)
    this->_make_connections(vector_id, candidates_by_layer, vector_layers);

  this->current_id++;
  return vector_id;
}

float Hnsw::_cosine_similarity(Embedding a, Embedding b) {
  if (a.vec_size != b.vec_size)
    throw "Error: Embedding size mismatch";
  float dot = 0.0;
  float a_mag = 0.0;
  float b_mag = 0.0;
  for (int i = 0; i < a.vec_size; i++) {
    dot += a.components[i] * b.components[i];
    a_mag += a.components[i] * a.components[i];
    b_mag += b.components[i] * b.components[i];
  }
  float denom = (sqrt(a_mag) * sqrt(b_mag));
  if (denom == 0.0)
    throw "Error: Division by zero";
  return dot / (sqrt(a_mag) * sqrt(b_mag));
}
std::vector<std::vector<int>> Hnsw::_search(Embedding query,
                                            int num_candidates) {
  std::vector<int> current_candidates;
  if (this->vec2neighbors.size() == 0)
    return std::vector<std::vector<int>>();

  for (int vector_id = 0; vector_id < this->current_id;
       vector_id++) // collect vectors in the highest layer (sparse layer)
    if (this->vec2neighbors[vector_id].size() == this->highest_layer + 1)
      current_candidates.push_back(vector_id);

  std::vector<std::vector<int>> layer_candidates;

  for (int layer_id = this->highest_layer; layer_id >= 0; layer_id--) {
    std::vector<int> all_neighbors;

    for (int current_candidate : current_candidates) {
      if (this->vec2neighbors[current_candidate].size() - 1 >= layer_id) {
        for (int i = 0;
             i < this->vec2neighbors[current_candidate][layer_id].size(); i++)
          all_neighbors.push_back(
              this->vec2neighbors[current_candidate][layer_id][i]);
      }
    }
    std::vector<int> unique_candidates;
    unique_candidates.insert(unique_candidates.end(), all_neighbors.begin(),
                             all_neighbors.end());
    unique_candidates.insert(unique_candidates.end(),
                             current_candidates.begin(),
                             current_candidates.end());
    std::sort(unique_candidates.begin(), unique_candidates.end());
    auto temp_vec =
        std::unique(unique_candidates.begin(), unique_candidates.end());
    unique_candidates.erase(temp_vec, unique_candidates.end());

    if (unique_candidates.size() == 0) {
      current_candidates.clear();
      continue;
    }

    std::vector<std::pair<float, int>> scored_candidates;
    for (int candidate : unique_candidates) {
      float similarity =
          this->_cosine_similarity(query, this->vec2embeds[candidate]);
      scored_candidates.push_back({similarity, candidate});
    }

    std::sort(scored_candidates.begin(), scored_candidates.end(),
              [](const std::pair<float, int> &a,
                 const std::pair<float, int> &b) { return a.first > b.first; });

    std::vector<int> top_candidates;
    for (int i = 0; i < std::min(num_candidates, int(scored_candidates.size()));
         i++)
      top_candidates.push_back(scored_candidates[i].second);

    layer_candidates.resize(this->highest_layer + 1);
    for (int i = 0; i < top_candidates.size(); i++)
      layer_candidates[layer_id].push_back(top_candidates[i]);

    int next_layer_size =
        std::min(int(num_candidates / 2),
                 int(top_candidates
                         .size())); // this is a simpler variant where it just
                                    // halves -- explore ef (exploration factor)
    current_candidates.clear();
    for (int i = 0; i < next_layer_size; i++)
      current_candidates.push_back(top_candidates[i]);
  }
  return layer_candidates;
}

void Hnsw::_make_connections(int new_vector_id,
                             std::vector<std::vector<int>> candidates_by_layer,
                             std::vector<int> new_vector_layers) {
  for (int layer : new_vector_layers) {
    if (layer >= candidates_by_layer.size())
      continue;

    std::vector<int> candidates =
        candidates_by_layer[layer]; // CANDIDATES SHOULD BE SORTED FROM
                                    // _SEARCH()

    int connections_to_make =
        std::min(this->max_connections, int(candidates.size()));
    std::vector<int> best_candidates;
    for (int i = 0; i < connections_to_make; i++)
      best_candidates.push_back(candidates[i]);

    for (int candidate_id : best_candidates) {
      this->vec2neighbors[new_vector_id][layer].push_back(candidate_id);

      // the worst neighbour search can be O(degree^2) in the worst case, this
      // is an HNSW quirk but can be mitigated by caching
      if (layer < this->vec2neighbors[candidate_id].size()) {
        std::vector<int> &candidate_neighbors =
            this->vec2neighbors[candidate_id][layer];
        if (candidate_neighbors.size() < this->max_connections) {
          candidate_neighbors.push_back(new_vector_id);
        } else {
          Embedding candidate_embedding = this->vec2embeds[candidate_id];
          Embedding new_embedding = this->vec2embeds[new_vector_id];
          float new_similarity =
              this->_cosine_similarity(candidate_embedding, new_embedding);

          float worst_similarity = 1.1;
          int worst_neighbor;

          for (int neighbor_id : candidate_neighbors) {
            Embedding neighbor_embedding = this->vec2embeds[neighbor_id];
            float similarity = this->_cosine_similarity(candidate_embedding,
                                                        neighbor_embedding);
            if (similarity < worst_similarity) {
              worst_similarity = similarity;
              worst_neighbor = neighbor_id;
            }
          }
          if (new_similarity > worst_similarity) {
            candidate_neighbors.erase(std::remove(candidate_neighbors.begin(),
                                                  candidate_neighbors.end(),
                                                  worst_neighbor),
                                      candidate_neighbors.end());
            candidate_neighbors.push_back(new_vector_id);
          }
        }
      }
    }
  }
}

std::vector<Hnsw::query_result> Hnsw::query(Embedding query_embedding, int k) {
  std::vector<Hnsw::query_result> scored_results;
  if (this->vec2neighbors.size() == 0)
    return scored_results;

  std::vector<std::vector<int>> layer_candidates =
      this->_search(query_embedding, k * 2);

  if (layer_candidates.size() == 0)
    return scored_results;

  std::vector<int> candidates = layer_candidates[0];

  for (int candidate_id : candidates) {
    float similarity = this->_cosine_similarity(query_embedding,
                                                this->vec2embeds[candidate_id]);
    scored_results.push_back(
        {this->vec2docs[candidate_id], similarity, candidate_id});
  }
  std::sort(scored_results.begin(), scored_results.end(),
            [](const query_result &a, const query_result &b) {
              return a.similarity > b.similarity;
            });

  if (scored_results.size() > (size_t)k)
    scored_results.erase(scored_results.begin() + k, scored_results.end());
  return scored_results;
}