-
Notifications
You must be signed in to change notification settings - Fork 1
/
ingest.py
59 lines (48 loc) · 2.54 KB
/
ingest.py
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
from langchain.embeddings import HuggingFaceEmbeddings, CacheBackedEmbeddings
from langchain.vectorstores import FAISS
from langchain.retrievers import BM25Retriever,EnsembleRetriever
from langchain.storage import LocalFileStore
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain.document_loaders import PyPDFLoader,DirectoryLoader
DATA_PATH = 'data/'
DB_FAISS_PATH = 'vectorstoredb/db_faiss'
def load_doc():
loader = DirectoryLoader(DATA_PATH, glob="*.pdf", loader_cls=PyPDFLoader)
documents = loader.load()
text_splitter = RecursiveCharacterTextSplitter(chunk_size=1500, chunk_overlap=200)
esops_documents = text_splitter.transform_documents(documents)
return esops_documents
def vectorize_doc(cache_embedder):
doc = load_doc()
vectorstore = FAISS.from_documents(doc, cache_embedder)
vectorstore.save_local(DB_FAISS_PATH)
bm25_retriever = BM25Retriever.from_documents(doc)
bm25_retriever.k = 5
return vectorstore, bm25_retriever
def load_db():
store = LocalFileStore("./cache")
embed_model = 'BAAI/bge-small-en-v1.5'
embedding_model = HuggingFaceEmbeddings(model_name=embed_model)
embedder = CacheBackedEmbeddings.from_bytes_store(embedding_model, store, namespace=embed_model)
vectorstore, bm25_retriever = vectorize_doc(embedder)
faiss_retriever = vectorstore.as_retriever(search_kwargs = {"k": 5})
ensemble_retriever = EnsembleRetriever(retrievers=[bm25_retriever,faiss_retriever], weights = [0.5,0.5])
return ensemble_retriever
# def create_vector_db():
# loader = DirectoryLoader(DATA_PATH, glob="*.pdf", loader_cls=PyPDFLoader)
# documents = loader.load()
# text_splitter = RecursiveCharacterTextSplitter(chunk_size=500, chunk_overlap=200)
# esops_documents = text_splitter.transform_documents(documents)
# store = LocalFileStore("./cache")
# embed_model = 'BAAI/bge-small-en-v1.5'
# embedding_model = HuggingFaceEmbeddings(model_name=embed_model)
# embedder = CacheBackedEmbeddings.from_bytes_store(embedding_model, store, namespace=embed_model)
# vectorstore = FAISS.from_documents(esops_documents, embedder)
# vectorstore.save_local(DB_FAISS_PATH)
# bm25_retriever = BM25Retriever.from_documents(esops_documents)
# bm25_retriever.k = 5
# faiss_retriever = vectorstore.as_retriever(search_kwargs = {"k": 5})
# ensemble_retriever = EnsembleRetriever(retrievers=[bm25_retriever,faiss_retriever], weights = [0.5,0.5])
# return ensemble_retriever
if __name__ == "__main__":
load_db()