34 lines
1.0 KiB
Python
34 lines
1.0 KiB
Python
|
from haystack.document_stores.faiss import FAISSDocumentStore
|
||
|
|
import os
|
||
|
|
import sys
|
||
|
|
sys.path.append('/root/home/BA_QA_HSMA/backendd')
|
||
|
|
|
||
|
|
class FAISSHandler:
|
||
|
|
def __init__(self) -> None:
|
||
|
|
self.doc_store= FAISSDocumentStore(faiss_index_factory_str="Flat",embedding_dim=5120, index="ib" )
|
||
|
|
pass
|
||
|
|
|
||
|
|
def write_data(self, data, index):
|
||
|
|
return self.doc_store.write_documents(documents=data, index=index)
|
||
|
|
|
||
|
|
def get_data(self, index):
|
||
|
|
return self.doc_store.get_all_documents(index=index)
|
||
|
|
|
||
|
|
def query_emb(self, emb, index):
|
||
|
|
return self.doc_store.query_by_embedding(query_emb=emb, index=index)
|
||
|
|
|
||
|
|
def get_all_doc_count(self,index):
|
||
|
|
return self.doc_store.get_document_count(index=index)
|
||
|
|
|
||
|
|
def delete_all_docs(self, index):
|
||
|
|
return self.doc_store.delete_all_documents(index=index)
|
||
|
|
|
||
|
|
def remove_docs(self):
|
||
|
|
pass
|
||
|
|
def update_embeddings(self, retriever):
|
||
|
|
return self.doc_store.update_embeddings(retriever=retriever, index="ib")
|
||
|
|
|
||
|
|
|
||
|
|
if __name__ == "__main__":
|
||
|
|
pass
|