doczyai-pipelines/archive/streamlit/ingest.py

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",
                )
                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()