This is my favorite and most used tool

[clearsitedesigns]

Hi all I have to fork off the main repo. I use this tool a lot in my research, and building some complexity. Since I am ingesting a ton of stuff I wrote a custom_ingest.py that gives me a bit more insight into what is going on and controls the token ingesting a bit more. Including, telling me what files are going, and giving a count of how much time is left. If there is interest I could merge this. It's not perfect. I usually throw a hundred files or so at a time, but if they are large PDF's that can be problematic.

For example,

• I get a count of the total amount ingested, documents that fail because of UTF (or other issues)
• Full read of all files, before ingestion, so if anything fails.
• calculate the total anticipated ingestion amount,
• and % remaining based on Token count split

I used to sit here and wonder how many hours something was going to take now I can see how long it will take to ingest these 195 documents, at least a better idea of how far done we are. I did have to update matlib + charset and a few other libs to get this to work.

1000
2024-02-05 16:14:26,396 - WARNING - text_splitter.py:176 - Created a chunk of size 1713, which is longer than the specified 1000
2024-02-05 16:14:26,970 - INFO - SentenceTransformer.py:66 - Load pre trained SentenceTransformer: hkunlp/instructor-large
load INSTRUCTOR_Transformer
max_seq_length 512
2024-02-05 16:14:34,784 - INFO - custom_ingest.py:151 - Ingested batch 1/195, 0.51% complete
2024-02-05 16:14:40,216 - INFO - custom_ingest.py:151 - Ingested batch 2/195, 1.03% complete
2024-02-05 16:14:45,999 - INFO - custom_ingest.py:151 - Ingested batch 3/195, 1.54% complete
2024-02-05 16:14:52,932 - INFO - custom_ingest.py:151 - Ingested batch 4/195, 2.05% complete
2024-02-05 16:14:58,689 - INFO - custom_ingest.py:151 - Ingested batch 5/195, 2.56% complete
2024-02-05 16:15:04,510 - INFO - custom_ingest.py:151 - Ingested batch 6/195, 3.08% complete
2024-02-05 16:15:10,455 - INFO - custom_ingest.py:151 - Ingested batch 7/195, 3.59% complete
2024-02-05 16:15:16,187 - INFO - custom_ingest.py:151 - Ingested batch 8/195, 4.10% complete
2024-02-05 16:15:21,968 - INFO - custom_ingest.py:151 - Ingested batch 9/195, 4.62% complete
2024-02-05 16:15:28,333 - INFO - custom_ingest.py:151 - Ingested batch 10/195, 5.13% complete
2024-02-05 16:15:33,562 - INFO - custom_ingest.py:151 - Ingested batch 11/195, 5.64% complete
2024-02-05 16:17:34,903 - INFO - custom_ingest.py:151 - Ingested batch 31/195, 15.90% complete

import logging
import os
import csv
from math import ceil
import click
import sys
import torch
from concurrent.futures import ProcessPoolExecutor, ThreadPoolExecutor, as_completed
from langchain.docstore.document import Document
from langchain.embeddings import HuggingFaceInstructEmbeddings
from langchain.text_splitter import CharacterTextSplitter
from langchain.vectorstores import Chroma

from constants import (
    CHROMA_SETTINGS,
    DOCUMENT_MAP,
    EMBEDDING_MODEL_NAME,
    INGEST_THREADS,
    PERSIST_DIRECTORY,
    SOURCE_DIRECTORY,
)

def process_batch(batch, file_path):
    logging.info(f"Processing batch from file: {file_path}")
    return [Document(page_content=' '.join(f'{k}: {"" if v is None else v}' for k, v in row.items()), 
                     metadata={"source": file_path}) for row in batch]

def load_csv_document(file_path: str, batch_size=10000) -> list[Document]:
    logging.info(f"Loading CSV document: {file_path}")
    documents = []
    batch = []
    csv.field_size_limit(sys.maxsize)
    with open(file_path, mode='r', encoding='utf-8') as csvfile:
        reader = csv.DictReader(csvfile)
        for row in reader:
            batch.append(row)
            if len(batch) >= batch_size:
                documents.extend(process_batch(batch, file_path))
                batch.clear()  # Clear the batch
                logging.info(f"Processed a batch of {batch_size} rows from {file_path}")

        if batch:
            documents.extend(process_batch(batch, file_path))
            logging.info(f"Processed remaining batch from {file_path}")

    return documents

def load_single_document(file_path: str) -> list[Document]:
    logging.info(f"Loading single document: {file_path}")
    file_extension = os.path.splitext(file_path)[1].lower()
    if file_extension == '.csv':
        return load_csv_document(file_path)
    else:
        loader_class = DOCUMENT_MAP.get(file_extension)
        if loader_class:
            loader = loader_class(file_path)
            return loader.load()
        else:
            logging.warning(f"Document type '{file_extension}' is undefined")
            return []

def load_document_batch(filepaths):
    logging.info("Loading document batch")
    with ThreadPoolExecutor(max(1, len(filepaths))) as exe:
        futures = [exe.submit(load_single_document, filepath) for filepath in filepaths]
        data_list = []
        for future in futures:
            result = future.result()
            if isinstance(result, list):
                data_list.extend(result)
            else:
                data_list.append(result)
        return (data_list, filepaths)

def load_documents(source_dir: str) -> list[Document]:
    logging.info(f"Scanning directory for documents: {source_dir}")
    paths = []
    for root, _, files in os.walk(source_dir):
        for file_name in files:
            file_extension = os.path.splitext(file_name)[1].lower()
            source_file_path = os.path.join(root, file_name)
            if file_extension in DOCUMENT_MAP.keys() or file_extension == '.csv':
                paths.append(source_file_path)

    if not paths:
        logging.info("No documents found to process.")
        return []

    total_files = len(paths)
    n_workers = min(INGEST_THREADS, max(len(paths), 1))
    chunksize = max(1, round(len(paths) / n_workers))
    docs = []

    with ProcessPoolExecutor(n_workers) as executor:
        futures = {}
        for i in range(0, total_files, chunksize):
            filepaths = paths[i : i + chunksize]
            future = executor.submit(load_document_batch, filepaths)
            futures[future] = filepaths
            logging.info(f"Processing batch: {filepaths}")

        files_processed = 0
        for future in as_completed(futures):
            contents, _ = future.result()
            docs.extend(contents)
            files_processed += len(futures[future])
            logging.info(f"Processed {files_processed}/{total_files} files.")

    return docs

@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):
    logging.basicConfig(
        format="%(asctime)s - %(levelname)s - %(filename)s:%(lineno)s - %(message)s", level=logging.INFO
    )
    logging.info(f"Script started with device_type: {device_type}")
    documents = load_documents(SOURCE_DIRECTORY)

    text_splitter = CharacterTextSplitter(chunk_size=1000, chunk_overlap=0)
    docs = text_splitter.split_documents(documents)
    total_chunks = len(docs)

    batch_size = 100  # Define a suitable batch size
    total_batches = ceil(total_chunks / batch_size)
    embeddings = HuggingFaceInstructEmbeddings(
        model_name=EMBEDDING_MODEL_NAME,
        model_kwargs={"device": device_type},
    )

    for i in range(total_batches):
        batch_start = i * batch_size
        batch_end = min(batch_start + batch_size, total_chunks)
        batch_docs = docs[batch_start:batch_end]

        db = Chroma.from_documents(
            batch_docs,
            embeddings,
            persist_directory=PERSIST_DIRECTORY,
            client_settings=CHROMA_SETTINGS,
        )

        percentage_complete = ((i + 1) / total_batches) * 100
        logging.info(f"Ingested batch {i + 1}/{total_batches}, {percentage_complete:.2f}% complete")

    logging.info("Chroma ingestion process completed")
    logging.info("Script completed")

if __name__ == "__main__":
    main()`

[PromtEngineer]

@clearsitedesigns this looks really useful. We can integrate it to the main code.