blob: f229ba7b66b8d482e1cf78a2595a884593b86403 (
plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
|
from typing import Tuple
import faiss
import numpy as np
# TODO: read from .env
EMBEDDING_DIM = 1024
# TODO: inner product distance metric?
class VectorStore:
def __init__(self):
self.index = faiss.IndexFlatL2(EMBEDDING_DIM)
# TODO: load from file
def add(self, embeddings: np.ndarray):
# TODO: save to file
self.index.add(embeddings)
def search(
self, query: np.ndarray, neighbors: int = 4
) -> Tuple[np.ndarray, np.ndarray]:
score, indices = self.index.search(query, neighbors)
return score, indices
|
