Spaces:
Running
Running
File size: 1,131 Bytes
c42a70a | 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 | from sentence_transformers import SentenceTransformer
import faiss
import numpy as np
class VectorStore:
def __init__(self):
self.index = None
self.chunks = []
self.model = None # lazy load
def load_model(self):
if self.model is None:
print("Loading model...")
self.model = SentenceTransformer("all-MiniLM-L6-v2")
def create_index(self, chunks):
self.load_model()
self.chunks = chunks
embeddings = self.model.encode(chunks)
if len(embeddings.shape) == 1:
embeddings = np.array([embeddings])
else:
embeddings = np.array(embeddings)
dim = embeddings.shape[1]
self.index = faiss.IndexFlatL2(dim)
self.index.add(embeddings)
def retrieve(self, query, k=3):
self.load_model()
query_embedding = self.model.encode([query])
if len(query_embedding.shape) == 1:
query_embedding = np.array([query_embedding])
distances, indices = self.index.search(query_embedding, k)
return [self.chunks[i] for i in indices[0]]
|