mirror of
https://github.com/zylon-ai/private-gpt.git
synced 2025-12-22 10:45:42 +01:00
122 lines
4 KiB
Python
Executable file
122 lines
4 KiB
Python
Executable file
#!/usr/bin/env python3
|
||
import os
|
||
import glob
|
||
from typing import List
|
||
from dotenv import load_dotenv
|
||
from multiprocessing import Pool
|
||
|
||
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
|
||
|
||
|
||
class MyElmLoader(UnstructuredEmailLoader):
|
||
"""Wrapper to fallback to text/plain when default does not work"""
|
||
|
||
def load(self) -> List[Document]:
|
||
"""Wrapper adding fallback for elm without html"""
|
||
try:
|
||
try:
|
||
doc = UnstructuredEmailLoader.load(self)
|
||
except ValueError as e:
|
||
if 'text/html content not found in email' in str(e):
|
||
# Try plain text
|
||
self.unstructured_kwargs["content_source"]="text/plain"
|
||
doc = UnstructuredEmailLoader.load(self)
|
||
else:
|
||
raise
|
||
except Exception as e:
|
||
# Add file_path to exception message
|
||
raise type(e)(f"{self.file_path}: {e}") from e
|
||
|
||
return doc
|
||
|
||
|
||
# Map file extensions to document loaders and their arguments
|
||
LOADER_MAPPING = {
|
||
".csv": (CSVLoader, {}),
|
||
# ".docx": (Docx2txtLoader, {}),
|
||
".doc": (UnstructuredWordDocumentLoader, {}),
|
||
".docx": (UnstructuredWordDocumentLoader, {}),
|
||
".enex": (EverNoteLoader, {}),
|
||
".eml": (MyElmLoader, {}),
|
||
".epub": (UnstructuredEPubLoader, {}),
|
||
".html": (UnstructuredHTMLLoader, {}),
|
||
".md": (UnstructuredMarkdownLoader, {}),
|
||
".odt": (UnstructuredODTLoader, {}),
|
||
".pdf": (PDFMinerLoader, {}),
|
||
".ppt": (UnstructuredPowerPointLoader, {}),
|
||
".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) -> List[Document]:
|
||
# Loads all documents from source documents directory
|
||
all_files = []
|
||
for ext in LOADER_MAPPING:
|
||
all_files.extend(
|
||
glob.glob(os.path.join(source_dir, f"**/*{ext}"), recursive=True)
|
||
)
|
||
with Pool(processes=os.cpu_count()) as pool:
|
||
documents = pool.map(load_single_document, all_files)
|
||
return documents
|
||
|
||
|
||
def main():
|
||
# 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')
|
||
|
||
# Load documents and split in chunks
|
||
print(f"Loading documents from {source_directory}")
|
||
chunk_size = 500
|
||
chunk_overlap = 50
|
||
documents = load_documents(source_directory)
|
||
text_splitter = RecursiveCharacterTextSplitter(chunk_size=chunk_size, chunk_overlap=chunk_overlap)
|
||
texts = text_splitter.split_documents(documents)
|
||
print(f"Loaded {len(documents)} documents from {source_directory}")
|
||
print(f"Split into {len(texts)} chunks of text (max. {chunk_size} characters each)")
|
||
|
||
# Create embeddings
|
||
embeddings = HuggingFaceEmbeddings(model_name=embeddings_model_name)
|
||
|
||
# Create and store locally vectorstore
|
||
db = Chroma.from_documents(texts, embeddings, persist_directory=persist_directory, client_settings=CHROMA_SETTINGS)
|
||
db.persist()
|
||
db = None
|
||
|
||
|
||
if __name__ == "__main__":
|
||
main()
|