128 lines
4.5 KiB
Python
128 lines
4.5 KiB
Python
import os
|
||
import glob
|
||
from typing import List
|
||
from dotenv import load_dotenv
|
||
|
||
from langchain.document_loaders import (
|
||
CSVLoader,
|
||
EverNoteLoader,
|
||
PDFMinerLoader,
|
||
TextLoader,
|
||
UnstructuredEmailLoader,
|
||
UnstructuredEPubLoader,
|
||
UnstructuredHTMLLoader,
|
||
UnstructuredMarkdownLoader,
|
||
UnstructuredODTLoader,
|
||
UnstructuredPowerPointLoader,
|
||
UnstructuredWordDocumentLoader,
|
||
)
|
||
|
||
from langchain.text_splitter import RecursiveCharacterTextSplitter
|
||
from langchain.vectorstores import Chroma
|
||
from langchain.embeddings import HuggingFaceEmbeddings
|
||
from langchain.docstore.document import Document
|
||
from constants import CHROMA_SETTINGS
|
||
|
||
|
||
load_dotenv()
|
||
|
||
|
||
# Load environment variables
|
||
persist_directory = os.environ.get('PERSIST_DIRECTORY')
|
||
source_directory = os.environ.get('SOURCE_DIRECTORY', 'source_documents')
|
||
embeddings_model_name = os.environ.get('EMBEDDINGS_MODEL_NAME')
|
||
chunk_size = 500
|
||
chunk_overlap = 50
|
||
|
||
# Map file extensions to document loaders and their arguments
|
||
LOADER_MAPPING = {
|
||
".csv": (CSVLoader, {}),
|
||
# ".docx": (Docx2txtLoader, {}),
|
||
".docx": (UnstructuredWordDocumentLoader, {}),
|
||
".enex": (EverNoteLoader, {}),
|
||
".eml": (UnstructuredEmailLoader, {}),
|
||
".epub": (UnstructuredEPubLoader, {}),
|
||
".html": (UnstructuredHTMLLoader, {}),
|
||
".md": (UnstructuredMarkdownLoader, {}),
|
||
".odt": (UnstructuredODTLoader, {}),
|
||
".pdf": (PDFMinerLoader, {}),
|
||
".pptx": (UnstructuredPowerPointLoader, {}),
|
||
".txt": (TextLoader, {"encoding": "utf8"}),
|
||
# Add more mappings for other file extensions and loaders as needed
|
||
}
|
||
|
||
|
||
load_dotenv()
|
||
|
||
def load_single_document(file_path: str) -> Document:
|
||
ext = "." + file_path.rsplit(".", 1)[-1]
|
||
if ext in LOADER_MAPPING:
|
||
loader_class, loader_args = LOADER_MAPPING[ext]
|
||
loader = loader_class(file_path, **loader_args)
|
||
return loader.load()[0]
|
||
|
||
raise ValueError(f"Unsupported file extension '{ext}'")
|
||
|
||
|
||
def load_documents(source_dir: str, ignored_files: List[str] = []) -> List[Document]:
|
||
"""
|
||
Loads all documents from the source documents directory, ignoring specified files
|
||
"""
|
||
all_files = []
|
||
for ext in LOADER_MAPPING:
|
||
all_files.extend(
|
||
glob.glob(os.path.join(source_dir, f"**/*{ext}"), recursive=True)
|
||
)
|
||
filtered_files = [file_path for file_path in all_files if file_path not in ignored_files]
|
||
return [load_single_document(file_path) for file_path in filtered_files]
|
||
|
||
def process_documents(ignored_files: List[str] = []) -> List[Document]:
|
||
"""
|
||
Load documents and split in chunks
|
||
"""
|
||
print(f"Loading documents from {source_directory}")
|
||
documents = load_documents(source_directory, ignored_files)
|
||
if not documents:
|
||
print("No new documents to load")
|
||
exit(0)
|
||
print(f"Loaded {len(documents)} new documents from {source_directory}")
|
||
text_splitter = RecursiveCharacterTextSplitter(chunk_size=chunk_size, chunk_overlap=chunk_overlap)
|
||
texts = text_splitter.split_documents(documents)
|
||
print(f"Split into {len(texts)} chunks of text (max. 500 tokens each)")
|
||
return texts
|
||
|
||
def does_vectorstore_exist(persist_directory: str) -> bool:
|
||
"""
|
||
Checks if vectorstore exists
|
||
"""
|
||
if os.path.exists(os.path.join(persist_directory, 'index')):
|
||
if os.path.exists(os.path.join(persist_directory, 'chroma-collections.parquet')) and os.path.exists(os.path.join(persist_directory, 'chroma-embeddings.parquet')):
|
||
list_index_files = glob.glob(os.path.join(persist_directory, 'index/*.bin'))
|
||
list_index_files += glob.glob(os.path.join(persist_directory, 'index/*.pkl'))
|
||
if len(list_index_files) == 4:
|
||
return True
|
||
return False
|
||
|
||
def main():
|
||
# Create embeddings
|
||
embeddings = HuggingFaceEmbeddings(model_name=embeddings_model_name)
|
||
|
||
if does_vectorstore_exist(persist_directory):
|
||
# Update and store locally vectorstore
|
||
print(f"Appending to existing vectorstore at {persist_directory}")
|
||
db = Chroma(persist_directory=persist_directory, embedding_function=embeddings, client_settings=CHROMA_SETTINGS)
|
||
collection = db.get()
|
||
texts = process_documents([metadata['source'] for metadata in collection['metadatas']])
|
||
db.add_documents(texts)
|
||
else:
|
||
# Create and store locally vectorstore
|
||
print("Creating new vectorstore")
|
||
texts = process_documents()
|
||
db = Chroma.from_documents(texts, embeddings, persist_directory=persist_directory, client_settings=CHROMA_SETTINGS)
|
||
db.persist()
|
||
db = None
|
||
|
||
|
||
if __name__ == "__main__":
|
||
main()
|