NearestNeighbourSearch.h

Go to the documentation of this file.

#pragma once
 
//
//  NearestNeighbourSearch.h
//  vector_search
//
//  Created by Chris Scott on 06/12/2022.
//
 
#include <cassert>
#include <chrono>
#include <cstdlib>
#include <fstream>
#include <iostream>
#include <map>
#include <sstream>
#include <vector>
#include "hnswlib/hnswlib.h"
 
namespace krotos
{
    struct Dataset {
        std::vector<std::string> filenames;
        std::map<std::string, std::vector<float>> vectors;
 
        // Read a vector collection in csv format.
        // Each line contains a filename and a comma-separated list of numbers.
        static Dataset read_embeddings(const std::string& filename);
    };
 
    class NearestNeighbourSearch {
        // Hnswlib - fast approximate nearest neighbor search
        // https://github.com/nmslib/hnswlib
 
    public:
 
        NearestNeighbourSearch( std::string datasetPath );
 
        ~NearestNeighbourSearch();
 
        void getNearestNeighbours(const std::string& query, const int numMatches, std::vector<std::string>& returnResults);
    
    private:
 
        std::vector<std::pair<float, hnswlib::labeltype> > knnQuery(std::string query, int k);
 
        void addDatasetItems();
        Dataset dataset;
        std::unique_ptr<hnswlib::BruteforceSearch<float>> alg;
        //hnswlib::HierarchicalNSW<float>* alg;
        hnswlib::SpaceInterface<float>* space;
 
        std::unique_ptr<hnswlib::InnerProductSpace> m_spaceInstance;
    };
 
 
}