NearestNeighbourSearch.cpp

Go to the documentation of this file.

//
//  NearestNeighbourSearch.cpp
//  vector_search
//
//  Created by Chris Scott on 06/12/2022.
//
 
#include "NearestNeighbourSearch.h"
 
namespace krotos
{
 
    NearestNeighbourSearch::NearestNeighbourSearch(std::string datasetPath )
    {
        dataset = Dataset::read_embeddings(datasetPath);
 
        if (dataset.filenames.size() == 0)
        {
            assert(false); // Empty dataset
            return;
        }
 
        assert(dataset.filenames.size() >= 0);
        const auto dim = dataset.vectors[dataset.filenames[0]].size();
        const std::size_t n = 2 * dataset.filenames.size();
 
        //mSink = std::make_unique<SinkFloat>(sinkChannelCount);
 
        m_spaceInstance = std::make_unique <hnswlib::InnerProductSpace>(dim);
        // The line below is for comapatability with hnswlib. space is a ptr to hnswlib::SpaceInterface,
        // an abstract class which can't be instantiated as is - it is used as an interface class
        // Assigning m_spaceInstance raw pointer here should be safe because m_spaceInstance
        // will persist until owning class is destroyed
        // If anyone can see a better way to do this outside of re-factoring hnswlib, please do !
        space = m_spaceInstance.get();  
        
        alg = std::make_unique<hnswlib::BruteforceSearch<float>>(space, n);
        //alg = new hnswlib::HierarchicalNSW<float>(space, n); // alternative
 
        addDatasetItems();
    }
 
    NearestNeighbourSearch::~NearestNeighbourSearch() {}
 
    std::vector<std::pair<float, hnswlib::labeltype>> NearestNeighbourSearch::knnQuery(std::string query, int k) {
        if (dataset.vectors.count(query) > 0)
            return alg->searchKnnCloserFirst(dataset.vectors[query].data(), k);
        else
            return std::vector<std::pair<float, hnswlib::labeltype> >();
    }
 
    void NearestNeighbourSearch::addDatasetItems() {
        // Insert each vector into the index.
        std::size_t index = 0;
        for (auto word : dataset.filenames) {
            alg->addPoint(dataset.vectors[word].data(), index);
            ++index;
        }
    }
 
    // Read a vector collection in csv format.
    // Each line contains a filename and a comma-separated list of numbers.
 
    Dataset Dataset::read_embeddings(const std::string& filename) {
        std::ifstream file(filename);
 
        if (!file.is_open()) {
            assert(false);          // File not found - should never actually get here as this method should not get called is file doesn't exist
            return Dataset();       // Return empty dataset
        }
 
        std::vector<std::string> filenames;
        std::map<std::string, std::vector<float>> vectors;
        while (!file.eof()) {
            std::string full_line;
            std::getline(file, full_line);
            std::istringstream line(full_line);
 
            std::string key;
            getline(line, key, ',');
 
            std::string field;
            std::vector<float> row;
            while (getline(line, field, ',')) {
                auto val = (float)atof(field.c_str());
                row.push_back(val);
            }
 
            if (row.size() != 0) {
                filenames.push_back(key);
                vectors[key] = row;
            }
        }
 
        Dataset res;
        res.filenames = filenames;
        res.vectors = vectors;
        return res;
    }
 
    void NearestNeighbourSearch::getNearestNeighbours(const std::string& query, const int numMatches, std::vector<std::string>& returnResults)
    {
        // Perform a search
        const auto results = knnQuery(query, numMatches);   // find k matches (including query)
 
        // Results will be empty if the query file wasn't in the dataset
        if (results.size() > 0)
        {
            // Return the k closest matches (including query).
            //assert(results[0].second == 0); // first match should be query
            for (std::size_t i = 0; i < results.size(); ++i) {
                // std::pair<distance score, index>
                returnResults.push_back( dataset.filenames[results[i].second] );
            }
        }
    }
    
}