Spaces:
Runtime error
Runtime error
| """ | |
| Python Backend API to chat with private data | |
| 08/15/2023 | |
| D.M. Theekshana Samaradiwakara | |
| """ | |
| import os | |
| import time | |
| import glob | |
| from multiprocessing import Pool | |
| from tqdm import tqdm | |
| from dotenv import load_dotenv | |
| from chromaDb import save_files | |
| from langchain.document_loaders import ( | |
| CSVLoader, | |
| EverNoteLoader, | |
| PyMuPDFLoader, | |
| TextLoader, | |
| UnstructuredEmailLoader, | |
| UnstructuredEPubLoader, | |
| UnstructuredHTMLLoader, | |
| UnstructuredMarkdownLoader, | |
| UnstructuredODTLoader, | |
| UnstructuredPowerPointLoader, | |
| UnstructuredWordDocumentLoader, | |
| ) | |
| from langchain.document_loaders import DirectoryLoader | |
| text_loader_kwargs={'autodetect_encoding': True} | |
| from langchain.text_splitter import RecursiveCharacterTextSplitter | |
| from langchain.embeddings import HuggingFaceEmbeddings | |
| from langchain.docstore.document import Document | |
| from chroma import load_store | |
| load_dotenv() | |
| chunk_size = os.environ.get('EMBEDDING_CHUNK_SIZE') | |
| chunk_overlap = os.environ.get('EMBEDDING_CHUNK_OVERLAP') | |
| embeddings_model_name = os.environ.get("EMBEDDINGS_MODEL_NAME") | |
| # Map file extensions to document loaders and their arguments | |
| LOADER_MAPPING = { | |
| ".csv": (CSVLoader, {}), | |
| # ".docx": (Docx2txtLoader, {}), | |
| ".doc": (UnstructuredWordDocumentLoader, {}), | |
| ".docx": (UnstructuredWordDocumentLoader, {}), | |
| ".enex": (EverNoteLoader, {}), | |
| ".eml": (UnstructuredEmailLoader, {}), | |
| ".epub": (UnstructuredEPubLoader, {}), | |
| ".html": (UnstructuredHTMLLoader, {}), | |
| ".md": (UnstructuredMarkdownLoader, {}), | |
| ".odt": (UnstructuredODTLoader, {}), | |
| ".pdf": (PyMuPDFLoader, {}), | |
| ".ppt": (UnstructuredPowerPointLoader, {}), | |
| ".pptx": (UnstructuredPowerPointLoader, {}), | |
| ".txt": (TextLoader, {"encoding": "utf8"}), | |
| # Add more mappings for other file extensions and loaders as needed | |
| } | |
| class DataPipeline: | |
| def __init__(self): | |
| self.dataset_name = None | |
| self.vectorstore = None | |
| def load_documents_in_folder(self, folder): | |
| print(f"loading documents...") | |
| loader = DirectoryLoader(folder, glob="**/[!.]*", loader_cls=TextLoader, loader_kwargs=text_loader_kwargs) | |
| pages = loader.load() | |
| return pages | |
| def load_single_document(self, doc): | |
| ext = "." + doc.name.rsplit(".", 1)[-1] | |
| if ext in LOADER_MAPPING: | |
| loader_class, loader_args = LOADER_MAPPING[ext] | |
| loader = loader_class(doc, **loader_args) | |
| return loader.load() | |
| raise ValueError(f"Unsupported file extension '{ext}'") | |
| def load_documents(self, uploaded_files): | |
| with Pool(processes=os.cpu_count()) as pool: | |
| results = [] | |
| with tqdm(total=len(uploaded_files), desc='Loading new documents', ncols=80) as pbar: | |
| for i, docs in enumerate(pool.imap_unordered(self.load_single_document, uploaded_files)): | |
| results.extend(docs) | |
| pbar.update() | |
| return results | |
| def load_streamlit_documents(self, uploaded_files, year): | |
| documents = [] | |
| for uploaded_file in uploaded_files: | |
| print(print("\n\n uploaded_file \n\n",uploaded_file,"\n")) | |
| source = uploaded_file.name | |
| print(print("\n\n source \n\n",source,"\n")) | |
| content = uploaded_file.read().decode('latin-1') | |
| print(print("\n\n content \n\n",content[:10],"\n")) | |
| doc = Document( | |
| page_content=content, | |
| metadata={ | |
| "source": source, | |
| 'year': year | |
| } | |
| ) | |
| print(print("doc")) | |
| print(print("\n doc \n\n",doc,"\n\n\n\n")) | |
| documents.append(doc) | |
| return documents | |
| def process_documents(self, documents): | |
| print(f"Creating embeddings. May take some minutes...") | |
| text_splitter = RecursiveCharacterTextSplitter( | |
| chunk_size=chunk_size, | |
| chunk_overlap=chunk_overlap, | |
| separators=["\n\n", "\n", "(?<=\. )", " ", ""] | |
| ) | |
| texts = text_splitter.split_documents(documents) | |
| return texts | |
| def persist_documents(self, persist_directory, document_splits): | |
| save_files(persist_directory, document_splits) | |
| def add_metadata(self, documents, metadata, value): | |
| for doc in documents: | |
| doc.metadata[metadata]=value | |
| return documents |