File size: 768 Bytes
e87a50a
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
#pragma once
#include<vector>
#include<cstdint>
#include "IndexPQ.h"
#include "IndexFlat.h"
#include <cstddef>
class IndexIVFPQ{
private:
    int d;
    int m;//bitquant 
    int nbucket; //no of centroid
    int ntotal; //no of vector index
    bool trained=false;
    size_t nprobe;//how many voronoi i should look at
    IndexFlatL2 router;
    IndexPQ pq;
    std::vector<float>coarse_centroids;
    std::vector<std::vector<uint8_t>>codes;
    std::vector<std::vector<int64_t>>ids;

public: 
    IndexIVFPQ(int d, int nbucket, int m);
    void train(int n, const float *x, bool subsampling, int seed);
    void add(int n, const float *x, const uint64_t* xids);
    void search(int n, const float *query, int k, int nprobe, float* distances, int64_t* labels);
};