| from langchain_mongodb import MongoDBAtlasVectorSearch | |
| from langchain.text_splitter import RecursiveCharacterTextSplitter | |
| from langchain_community.document_loaders import PyPDFLoader, TextLoader | |
| from embed_with_db import embeddings, config, client | |
| from tqdm import tqdm | |
| class VectorDataBase(): | |
| def __init__(self, file_path, db_collection, file_type='pdf', page_start=0): | |
| self.file_path = file_path | |
| self.file_type= file_type | |
| self.text_splitter = RecursiveCharacterTextSplitter(chunk_size=512, chunk_overlap=32) | |
| self.db_collection = client[config['DB_NAME']][db_collection] | |
| self.start_page = int(page_start) | |
| def load_docs_split(self): | |
| if str(self.file_type).lower() == 'pdf': | |
| loader = PyPDFLoader(self.file_path) | |
| elif str(self.file_type).lower() == 'text': | |
| loader = TextLoader(self.file_path) | |
| else: | |
| loader = None | |
| if loader: | |
| docs = loader.load() | |
| return self.text_splitter.split_documents(docs) | |
| else: | |
| return self.text_splitter.create_documents([self.file_path]) | |
| def docs_embeddings(self): | |
| texts = self.load_docs_split() | |
| if texts: | |
| docsearch = MongoDBAtlasVectorSearch.from_documents( | |
| texts, | |
| embeddings, | |
| collection=self.db_collection, | |
| index_name=config['VECTOR_SEARCH_INDEX']) | |
| print('done!') | |
| return docsearch | |
| else: | |
| print('documents is not embedded') | |
| return 'Some issues' | |
| def add_collection_database(self,doc): | |
| self.db_collection.insert_one( | |
| { | |
| 'text': doc.page_content, | |
| 'embedding': embeddings.embed_query(doc.page_content), | |
| 'source': doc.metadata.get('source', 'Unknown'), | |
| 'page': doc.metadata.get('page', 0) | |
| } | |
| ) | |
| def embedding_with_loop(self): | |
| docs = self.load_docs_split() | |
| if docs: | |
| for doc in tqdm(docs[self.start_page:]): | |
| self.add_collection_database(doc) | |
| print('Done') | |
| else: | |
| raise Exception('Some issue with it') | |