Mayank Aamseek
220 lines
8.5 KiB
Python
220 lines
8.5 KiB
Python
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()
|