Spaces:
Runtime error
Runtime error
File size: 801 Bytes
7f7b773 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 |
from langchain.vectorstores import FAISS
class FAISSVectorStore:
def __init__(self, embedding_model):
self.embedding_model = embedding_model
self.db = None
def initialize_from_documents(self, docs):
self.db = FAISS.from_documents(docs, self.embedding_model.model)
def initialize_from_file(self, path):
self.db = FAISS.load_local(path, self.embedding_model.model)
def save(self, path):
self.db.save_local(path)
def add_documents(self, documents):
return self.db.add_documents(documents)
def query(self, query: str, k: int = 4):
# TODO adjust fetch_k parameter. It is now set to match the defaults k=4, fetch_k=20 in the original code.
return self.db.similarity_search_with_score(query, k=k, fetch_k=5*k)
|