testraft.cu

#include <cstdint>
#include <raft/core/device_mdarray.hpp>
#include <raft/core/device_resources.hpp>
#include <raft/neighbors/cagra.cuh>
#include <raft/random/make_blobs.cuh>
#include <iostream>
#include <fstream>
#include <string>

#include <rmm/mr/device/device_memory_resource.hpp>
#include <rmm/mr/device/pool_memory_resource.hpp>

#include "common.cuh"
#include "rapidcsv.h"

void cagra_build_search_simple(raft::device_resources const &dev_resources,
                               raft::device_matrix_view<const float, int64_t> dataset,
                               raft::device_matrix_view<const float, int64_t> queries, int64_t dim, int64_t tk)
{
    using namespace raft::neighbors;

    int64_t topk = tk;
    int64_t n_queries = queries.extent(0);

    std::cout << "build started ..." << std::endl;

    std::cout << n_queries << std::endl;

//   for(int i = 0 ; i < n_queries; i++) {
//     for (int j = 0 ; j < queries.extent(1); j++) {
//         std::cout << queries(i, j) << ",";
//     }
//     std::cout << std::endl;
//   }

    // create output arrays
    auto neighbors = raft::make_device_matrix<uint32_t>(dev_resources, n_queries, topk);
    auto distances = raft::make_device_matrix<float>(dev_resources, n_queries, topk);

    std::cout << "Allocated neighbors and distances ..." << std::endl;

    // use default index parameters
    cagra::index_params index_params;

    index_params.graph_degree = 200;
    index_params.intermediate_graph_degree = 200;

    std::cout << "Building CAGRA index (search graph)" << std::endl;
    auto index = cagra::build<float, uint32_t>(dev_resources, index_params, dataset);

    std::cout << "Build complete. " << std::endl;

    std::cout << "CAGRA index has " << index.size() << " vectors" << std::endl;
    std::cout << "CAGRA graph has degree " << index.graph_degree() << ", graph size ["
              << index.graph().extent(0) << ", " << index.graph().extent(1) << "]" << std::endl;

    std::cout << "Starting search ..." << std::endl;

    // use default search parameters
    cagra::search_params search_params;
    // search K nearest neighbors
    cagra::search<float, uint32_t>(
            dev_resources, search_params, index, queries, neighbors.view(), distances.view());

    // The call to ivf_flat::search is asynchronous. Before accessing the data, sync by calling
    // raft::resource::sync_stream(dev_resources);

    print_results(dev_resources, neighbors.view(), distances.view());
}

int main()
{
    raft::device_resources dev_resources;

    // Set pool memory resource with 1 GiB initial pool size. All allocations use the same pool.
    rmm::mr::pool_memory_resource<rmm::mr::device_memory_resource> pool_mr(
            rmm::mr::get_current_device_resource(), 1 * 1024 * 1024 * 1024ull);
    rmm::mr::set_current_device_resource(&pool_mr);

    // Alternatively, one could define a pool allocator for temporary arrays (used within RAFT
    // algorithms). In that case only the internal arrays would use the pool, any other allocation
    // uses the default RMM memory resource. Here is how to change the workspace memory resource to
    // a pool with 2 GiB upper limit.
    // raft::resource::set_workspace_to_pool_resource(dev_resources, 20 * 1024 * 1024 * 1024ull);

    //  Create input arrays.
    //  int64_t n_samples = 10;
    //  int64_t n_dim = 10;
    //  int64_t n_queries = 10;
    //  auto dataset = raft::make_device_matrix<float, int64_t>(dev_resources, n_samples, n_dim);
    //  auto queries = raft::make_device_matrix<float, int64_t>(dev_resources, n_queries, n_dim);
    //  generate_dataset(dev_resources, dataset.view(), queries.view());


    // Preparing query vector

    std::fstream queryfile;
    queryfile.open("/home/narangvivek/searchscale/repos/lucene-raft/examples/raft-cagra-example/src/query.txt",std::ios::in);
    std::vector<std::string> v;

    if (queryfile.is_open()){
        std::string tp;
        getline(queryfile, tp);
        tp = tp.substr(1,tp.length()-2);

        std::stringstream ss(tp);

        while (ss.good()) {
            std::string substr;
            getline(ss, substr, ',');
            v.push_back(substr);
        }

        queryfile.close();
    }

    std::cout << "Input vector size: " << v.size() << std::endl;

    // for(int i = 0 ; i < v.size() ; i ++) {
    //    std::cout << v[i] << ", ";
    // }
    std::cout << std::endl << "Now allocating query vector into device matrix" << std::endl;

    int64_t n_dim = v.size();
    int64_t n_queries = 1;

    auto queries = raft::make_device_matrix<float, int64_t>(dev_resources, n_queries, n_dim);

    for(int j = 0; j < n_dim; ++j) {
        queries(0, j) = std::stof(v[j]);
    }

    std::cout << "Printing queries dm" << std::endl;
    std::cout << "e0 : " << queries.extent(0) << " e1 : " << queries.extent(1) << std::endl;
    for(int i = 0 ; i < queries.extent(0); i++) {
        for (int j = 0 ; j < queries.extent(1); j++) {
            std::cout << queries(i, j) << ",";
        }
        std::cout << std::endl;
    }
    std::cout << "END Printing queries dm" << std::endl;

    // Preparing data vector
    std::cout << "dataset csv processing starting ..." << std::endl;
    rapidcsv::Document doc("/home/narangvivek/searchscale/repos/lucene-raft/examples/raft-cagra-example/src/vector_database_wikipedia_articles_embedded.csv", rapidcsv::LabelParams(),
                           rapidcsv::SeparatorParams(',' /* pSeparator */,
                                                     false /* pTrim */,
                                                     rapidcsv::sPlatformHasCR /* pHasCR */,
                                                     true /* pQuotedLinebreaks */,
                                                     false /* pAutoQuote */));


    int64_t n_dataset = doc.GetRowCount();
    auto dataset = raft::make_device_matrix<float, int64_t>(dev_resources, n_dataset, n_dim);

    for(size_t i = 0; i < n_dataset ; i ++) {

        // Get row
        std::vector row = doc.GetRow<std::string>(i);

        // Since the row has a quote and a bracket in the beginning and the end, remove those
        std::string tpx;
        tpx = row[5].substr(2,row[5].length()-3); // Content vector column

        // Split by a comma
        std::vector<std::string> vx;
        std::stringstream ss(tpx);
        while (ss.good()) {
            std::string substr;
            getline(ss, substr, ',');
            vx.push_back(substr);
        }

        // Populate dataset with converted float values for the row
        for(size_t j = 0; j < n_dim; ++j) {
            dataset(i, j) = std::stof(vx[j]);
        }

    }
    std::cout << "dataset csv processing END" << std::endl;


    // Simple build and search example.
    cagra_build_search_simple(dev_resources,
                              raft::make_const_mdspan(dataset.view()),
                              raft::make_const_mdspan(queries.view()), n_dim, 20);
}

The dataset and query both have 1536 dimensions. Dataset is here:
https://cdn.openai.com/API/examples/data/vector_database_wikipedia_articles_embedded.zip

Query file: https://termbin.com/i9of