import logging import os from concurrent.futures import ProcessPoolExecutor, ThreadPoolExecutor, as_completed import click import torch from langchain.docstore.document import Document from langchain.embeddings import HuggingFaceInstructEmbeddings from langchain.text_splitter import Language, RecursiveCharacterTextSplitter from langchain.vectorstores import Chroma import uuid from constants import ( CHROMA_SETTINGS, DOCUMENT_MAP, EMBEDDING_MODEL_NAME, INGEST_THREADS, PERSIST_DIRECTORY, SOURCE_DIRECTORY, ) import boto3 from langchain.embeddings.bedrock import BedrockEmbeddings def file_log(logentry): file1 = open("file_ingest.log","a") file1.write(logentry + "\n") file1.close() print(logentry + "\n") def load_single_document(file_path: str) -> Document: # Loads a single document from a file path try: file_extension = os.path.splitext(file_path)[1] loader_class = DOCUMENT_MAP.get(file_extension) if loader_class: file_log(file_path + ' loaded.') loader = loader_class(file_path) else: file_log(file_path + ' document type is undefined.') raise ValueError("Document type is undefined") return loader.load()[0] except Exception as ex: file_log('%s loading error: \n%s' % (file_path, ex)) return None def load_document_batch(filepaths): logging.info("Loading document batch") # create a thread pool with ThreadPoolExecutor(len(filepaths)) as exe: # load files futures = [exe.submit(load_single_document, name) for name in filepaths] # collect data if futures is None: file_log(name + ' failed to submit') return None else: data_list = [future.result() for future in futures] # return data and file paths return (data_list, filepaths) def load_documents(source_dir: str) -> list[Document]: # Loads all documents from the source documents directory, including nested folders paths = [] for root, _, files in os.walk(source_dir): for file_name in files: print('Importing: ' + file_name) file_extension = os.path.splitext(file_name)[1] source_file_path = os.path.join(root, file_name) if file_extension in DOCUMENT_MAP.keys(): paths.append(source_file_path) # Have at least one worker and at most INGEST_THREADS workers n_workers = min(INGEST_THREADS, max(len(paths), 1)) chunksize = round(len(paths) / n_workers) docs = [] with ProcessPoolExecutor(n_workers) as executor: futures = [] # split the load operations into chunks for i in range(0, len(paths), chunksize): # select a chunk of filenames filepaths = paths[i : (i + chunksize)] # submit the task try: future = executor.submit(load_document_batch, filepaths) except Exception as ex: file_log('executor task failed: %s' % (ex)) future = None if future is not None: futures.append(future) # process all results for future in as_completed(futures): # open the file and load the data try: contents, _ = future.result() docs.extend(contents) except Exception as ex: file_log('Exception: %s' % (ex)) return docs def split_documents(documents: list[Document]) -> tuple[list[Document], list[Document]]: # Splits documents for correct Text Splitter text_docs, python_docs = [], [] for doc in documents: if doc is not None: file_extension = os.path.splitext(doc.metadata["source"])[1] if file_extension == ".py": python_docs.append(doc) else: text_docs.append(doc) return text_docs, python_docs def process_in_batches(texts, batch_size): for i in range(0, len(texts), batch_size): yield texts[i:i+batch_size] @click.command() @click.option( "--device_type", default="cuda" if torch.cuda.is_available() else "cpu", type=click.Choice( [ "cpu", "cuda", "ipu", "xpu", "mkldnn", "opengl", "opencl", "ideep", "hip", "ve", "fpga", "ort", "xla", "lazy", "vulkan", "mps", "meta", "hpu", "mtia", ], ), help="Device to run on. (Default is cuda)", ) def main(device_type): # Load documents and split in chunks logging.info(f"Loading documents from {SOURCE_DIRECTORY}") for filename in os.listdir(SOURCE_DIRECTORY): with open(os.path.join(SOURCE_DIRECTORY, filename), 'r') as infile: text = infile.read() text_splitted = [i for i in text.split('Start of Page No. = ')] for i, txt in enumerate(text_splitted): if len(txt) > 2: page_path = os.path.join(SOURCE_DIRECTORY, f'{filename[:-4]}_page{i}.txt') with open(page_path, 'w') as f: f.write(txt) documents = [load_single_document(page_path)] os.remove(page_path) text_documents, python_documents = split_documents(documents) text_splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=200) python_splitter = RecursiveCharacterTextSplitter.from_language( language=Language.PYTHON, chunk_size=880, chunk_overlap=200 ) texts = text_splitter.split_documents(text_documents) texts.extend(python_splitter.split_documents(python_documents)) logging.info(f"Loaded {len(documents)} documents from {SOURCE_DIRECTORY}") logging.info(f"Split into {len(texts)} chunks of text") # Create embeddings # embeddings = HuggingFaceInstructEmbeddings( # model_name=EMBEDDING_MODEL_NAME, # model_kwargs={"device": device_type}, # ) bedrock_runtime = boto3.client( service_name="bedrock-runtime", region_name="us-east-1", aws_access_key_id="ASIAZTMXAXNXD3TOUOAJ", aws_secret_access_key="pk7k69CqXZPB/bf2hdFsW+47D5WYkoWXFdxQ633X", aws_session_token="IQoJb3JpZ2luX2VjELj//////////wEaCXVzLWVhc3QtMiJIMEYCIQCqcyLSwLeN9RX6tz+TgB4VMabHYwPcA3bTzx7xZ5b75QIhAMiOKbDlmPZ+eHkeskzelK4M9gWtS7GqlzJq5qmWbbrOKo4DCKH//////////wEQABoMNjYwMTMxMDY4NzgyIgzPuv0vReC65Mo/bQcq4gLXtII+VTaxES+JrjHHWpWywdmpsJneN6bcB3U37z7J8BFA5aaUqkATKmwJI0brmk5ZJKL8SpgBEC2TNdA/V9nzbMlf2HPunhEv6OPLzSWp5iJqaNL945MP764CbkYvfN9QWd6durUv1WgGZRNcbMzXg2UFsxKcRql795vtOmL207+R7uIouWl73So7NaCkEgaj4FdEJ9lbnfvWFeNcBlHbjwUx8e9EJjwm8D60OkTdS4w7Q3EacoEKLO94/kp2RtsaggAUV33OcvO/32VwYJzhRJYuveQUZnIzfkmybGtrXkkWLqMO9pls1bkTmIjaeMwcL8Uo7oowR9sTFCT87rY711yIYBGVOjkN9mfavPH4FSCNeeI6ta5aXqa7iVDU7rPerpFtle2i1VGfTW5bKoaWxjO13IdqI4yI9Pgibl+FuVeRB7md2tDuS6SMAuX4qpWnMVDLKCItlAxXlZdmVhuDFf8wr+7rrgY6pQH02EIrvCkYYXbJd8u7t78hOam3lSNPz+nghQHA5ppl7TnyRBL5/9Rfqp5Ib8y54HdC4hbzz/7w6lvfS0QbHE3Q6r2GC6XWk9L4FAxD8pyZkIYbtwf/WueX1g0r+0x7uLCxKxbWsYum/bigyvNxsDyFIcQcUOm+OJsXGPH3Z7z+qWKni5N9RRA4SEEpdSTRkAWipDtgVNII6kd/eWvyE74d+LYUrmA=" ) embeddings = BedrockEmbeddings( client=bedrock_runtime, model_id="amazon.titan-embed-text-v1", ) """ db = Chroma.from_documents( texts, embeddings, persist_directory=PERSIST_DIRECTORY, client_settings=CHROMA_SETTINGS, ) """ # for batch_texts in process_in_batches(texts, 20000): # https://github.com/PromtEngineer/localGPT/issues/489y print(filename) db = Chroma.from_documents( texts, embeddings, persist_directory=PERSIST_DIRECTORY, client_settings=CHROMA_SETTINGS, collection_metadata={"hnsw:space": "cosine"}, # ids = [str(filename)] ) if __name__ == "__main__": logging.basicConfig( format="%(asctime)s - %(levelname)s - %(filename)s:%(lineno)s - %(message)s", level=logging.INFO ) main()