Delete ingest.py

ingest.py

@@ -1,158 +0,0 @@
-import os
-from tqdm import tqdm
-import pathlib
-
-from langchain_community.document_loaders import TextLoader
-from langchain.docstore.document import Document
-from langchain_community.embeddings import HuggingFaceEmbeddings
-from langchain.text_splitter import RecursiveCharacterTextSplitter
-from langchain_community.vectorstores import FAISS
-
-os.environ["RAY_memory_monitor_refresh_ms"] = "0"
-os.environ["RAY_DEDUP_LOGS"] = "0"
-import ray
-
-from common import DATASET_DIR, EMBEDDING_MODEL_NAME, MODEL_KWARGS, VECTORSTORE_FILENAME
-
-# Each document is parsed on the same CPU, to decrease paging and data copies, and up to the the number of vCPUs.
-CONCURRENCY = 32
-
-
-# @ray.remote(num_cpus=1) # Outside a container, num_cpus=1 might speed things dramatically.
-@ray.remote
-def parse_doc(document_path: str) -> Document:
-    print("Loading", document_path)
-    loader = TextLoader(document_path)
-    langchain_dataset_documents = loader.load()
-
-    # Update the metadata with the proper metadata JSON file, parsed from Arxiv.com
-    return langchain_dataset_documents
-
-
-def add_documents_to_vector_store(
-    vector_store, new_documents, text_splitter, embeddings
-):
-    split_docs = text_splitter.split_documents(new_documents)
-    # print("Embedding vectors...")
-    store = FAISS.from_documents(split_docs, embeddings)
-    if vector_store is None:
-        vector_store = store
-    else:
-        print("Updating vector store", store)
-        vector_store.merge_from(store)
-    return vector_store
-
-
-def ingest_dataset_to_vectore_store(
-    vectorstore_filename: str, dataset_directory: os.PathLike
-):
-    ray.init()
-    vector_store = None
-    text_splitter = RecursiveCharacterTextSplitter(
-        chunk_size=160,  # TODO: Finetune
-        chunk_overlap=40,  # TODO: Finetune
-        length_function=len,
-    )
-
-    dataset_documents = []
-    dataset_dir_path = pathlib.Path(dataset_directory)
-    dataset_dir_path.mkdir(exist_ok=True)
-
-    for _dirname in os.listdir(str(dataset_dir_path)):
-        if _dirname.startswith("."):
-            continue
-        catagory_path = dataset_dir_path / pathlib.Path(_dirname)
-        for filename in os.listdir(str(dataset_dir_path / catagory_path)):
-            dataset_path = dataset_dir_path / catagory_path / pathlib.Path(filename)
-            dataset_documents.append(str(dataset_path))
-    print(dataset_documents)
-    print(f"Found {len(dataset_documents)} items in dataset: ")
-    langchain_documents = []
-
-    model_name = EMBEDDING_MODEL_NAME
-    model_kwargs = MODEL_KWARGS
-    print("Creating huggingface embeddings for ", model_name)
-    embeddings = HuggingFaceEmbeddings(model_name=model_name, model_kwargs=model_kwargs)
-
-    if vector_store is None and os.path.exists(vectorstore_filename):
-        print("Loading existing vector store from", vectorstore_filename)
-        vector_store = FAISS.load_local(
-            vectorstore_filename, embeddings, allow_dangerous_deserialization=True
-        )
-
-    jobs = []
-    docs_count = len(dataset_documents)
-    failed = 0
-    print(f"Embedding {docs_count} documents with Ray...")
-    for i, document in enumerate(tqdm(dataset_documents)):
-        try:
-            # print(f"Submitting job ", i)
-            job = parse_doc.remote(document)
-            jobs.append(job)
-
-            if i > 1 and i <= docs_count and i % CONCURRENCY == 0:
-                if langchain_documents:
-                    vector_store = add_documents_to_vector_store(
-                        vector_store, langchain_documents, text_splitter, embeddings
-                    )
-                    print(f"\nSaving vector store to disk at {vectorstore_filename}...")
-                    try:
-                        os.unlink(vectorstore_filename)
-                    except:
-                        ...
-
-                    vector_store.save_local(vectorstore_filename)
-                    langchain_documents = []
-                    jobs = []
-
-            # Block jobs every CONCURRENCY iterations
-            if i > 1 and i % CONCURRENCY == 0:
-                # print(f"Collecting {len(jobs)} jobs...")
-                for _ in jobs:
-                    try:
-                        # print("waiting for ray job ", _)
-                        data = ray.get(_)
-                        langchain_documents.extend(data)
-                    except Exception as e:
-                        print("error in job: ", e)
-                        continue
-        except Exception as e:
-            print(f"\n\nERROR reading dataset {i}:", e)
-            failed = failed + 1
-            continue
-
-    # print(f"Collecting {len(jobs)} jobs...")
-    for _ in jobs:
-        try:
-            print("waiting for ray job ", _)
-            data = ray.get(_)
-            langchain_documents.extend(data)
-        except Exception as e:
-            print("error in job: ", e)
-            continue
-
-    if langchain_documents:
-        vector_store = add_documents_to_vector_store(
-            vector_store, langchain_documents, text_splitter, embeddings
-        )
-        print(f"\nSaving vector store to disk at {vectorstore_filename}...")
-        try:
-            os.unlink(vectorstore_filename)
-        except:
-            ...
-
-        vector_store.save_local(vectorstore_filename)
-
-    return vector_store
-
-
-def main():
-    vectorstore_filename = VECTORSTORE_FILENAME
-    dataset_directory = DATASET_DIR
-    ingest_dataset_to_vectore_store(
-        vectorstore_filename=vectorstore_filename, dataset_directory=dataset_directory
-    )
-
-
-if __name__ == "__main__":
-    main()