OxVec/oxvec.hpp at main · ToyMath/OxVec · GitHub

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
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
#ifndef OXVEC_HPP
#define OXVEC_HPP

#include <iostream>
#include <string>
#include <vector>
#include <unordered_map>
#include <memory>
#include <cmath>
#include <queue>
#include <algorithm>

using Vector = std::vector<double>;
using VectorMap = std::unordered_map<std::string, std::shared_ptr<class VectorEntry>>;

class VectorEntry;
class SimilarityMetric;
class EuclideanDistance;
class CosineSimilarity;
class DotProductSimilarity;
class VectorIndex;
class VectorCollection;
class OxVec;

void printVector(const Vector& vec);

class VectorEntry {
public:
    VectorEntry(const std::string& id, const Vector& vector);

    const std::string& getId() const;
    const Vector& getVector() const;

private:
    std::string id_;
    Vector vector_;
};

class SimilarityMetric {
public:
    virtual ~SimilarityMetric() = default;
    virtual double computeSimilarity(const Vector& vector1, const Vector& vector2) const = 0;
};

class EuclideanDistance : public SimilarityMetric {
public:
    double computeSimilarity(const Vector& vector1, const Vector& vector2) const override;
};

class CosineSimilarity : public SimilarityMetric {
public:
    double computeSimilarity(const Vector& vector1, const Vector& vector2) const override;
};

class DotProductSimilarity : public SimilarityMetric {
public:
    double computeSimilarity(const Vector& vector1, const Vector& vector2) const override;
};

class VectorIndex {
public:
    void buildIndex(const VectorMap& vectors);

    std::vector<std::shared_ptr<VectorEntry>> search(
        const Vector& queryVector,
        const SimilarityMetric& metric,
        int topK,
        bool isDistanceMetric = false) const;

private:
    VectorMap vectors_;
};

class VectorCollection {
public:
    explicit VectorCollection(const std::string& name);

    void addVector(const std::string& id, const Vector& vector);
    void deleteVector(const std::string& id);
    void updateVector(const std::string& id, const Vector& vector);
    std::shared_ptr<VectorEntry> getVector(const std::string& id) const;

    std::vector<std::shared_ptr<VectorEntry>> query(
        const Vector& queryVector,
        const SimilarityMetric& metric,
        int topK) const;

    const VectorMap& getVectors() const;

private:
    std::string name_;
    VectorMap vectors_;
    std::shared_ptr<VectorIndex> index_;
};

class OxVec {
public:
    void createCollection(const std::string& name);
    void deleteCollection(const std::string& name);
    std::shared_ptr<VectorCollection> getCollection(const std::string& name) const;

private:
    std::unordered_map<std::string, std::shared_ptr<VectorCollection>> collections_;
};

#endif