jaynopponep/LinuxGPT

1from langchain_community.document_loaders import PyPDFLoader, PyPDFDirectoryLoader2from langchain_chroma import Chroma3from langchain.schema import Document4from langchain_openai import OpenAIEmbeddings5from langchain.text_splitter import RecursiveCharacterTextSplitter6from dotenv import load_dotenv7import os8import shutil9 10load_dotenv()11OPEN_AI_KEY = os.getenv('OPEN_AI_KEY')12 13CHROMA_PATH = "chroma"14DATA_PATH = "data/"15TEST_PATH = "data/theory_of_computation.pdf"16 17embed = OpenAIEmbeddings(18    api_key=OPEN_AI_KEY,19    model="text-embedding-3-large"20)21 22 23def main():24    generate_data_store()25    # print(load_documents())26 27 28def generate_data_store():29    documents = load_documents()30    chunks = split_text(documents)31    save_to_chroma(chunks)32 33 34def load_documents():35    loader = PyPDFDirectoryLoader(DATA_PATH)36    docs = loader.load()37    print(docs[0].metadata)38    return docs39 40    # loader = PyPDFLoader(TEST_PATH)41    # docs = []42    # docs_lazy = loader.load()43    # for doc in docs_lazy:44    #     docs.append(doc)45    # return docs_lazy46 47 48def split_text(documents: list[Document]):49    # chunk_size = 1000,50    # chunk_overlap = 200,51    # length_function = len,52    # add_start_index = True,53    text_splitter = RecursiveCharacterTextSplitter(54        chunk_size=1100,55        chunk_overlap=100,56        length_function=len,57    )58    chunks = text_splitter.split_documents(documents)59    print(f"Split {len(documents)} documents into {len(chunks)} chunks.")60    document = chunks[10]61    print(document.page_content)62    print(document.metadata)63    return chunks64 65 66def save_to_chroma(chunks: list[Document]):67    if os.path.exists(CHROMA_PATH):  # clear out the DB first68        shutil.rmtree(CHROMA_PATH)69 70    db = Chroma(71        collection_name="linux_funds",72        embedding_function=embed,73        persist_directory=CHROMA_PATH74    )75 76    # below breaks text & metadata down to Chroma vector store77    texts = [chunk.page_content for chunk in chunks]78    metadatas = [chunk.metadata for chunk in chunks]79    db.add_texts(texts=texts, metadatas=metadatas)80    print(f"Saved {len(chunks)} chunks to CHROMA PATH {CHROMA_PATH}.")81 82 83if __name__ == "__main__":84    main()85 86