pythonchatbotlarge-language-modelllamachromadb

How to load vectors from stored chroma db?


I'm working on creating a RAG-based LLM. The system is working correctly, i.e., the vector embeddings are successfully created and stored in the respective directory. However, the issue i'm facing is loading back those vectors from the stored chroma db file. Following is my function that handles the creation and retrieval of vectors:

def vector_embedding():
    persist_directory = "./chroma_db"

    if os.path.exists(persist_directory):
        st.write("Loading vectors from disk...")
        st.session_state.vectors = Chroma(persist_directory=persist_directory, embedding_function=OllamaEmbeddings(model="nomic-embed-text"))
        st.write("Loaded vectors from disk.")
        return
    
    st.write("Creating new vectors...")
    
    st.session_state.embeddings = OllamaEmbeddings(model="nomic-embed-text", show_progress=True)
    
    file_path = "data2.pdf"
    st.session_state.loader = UnstructuredPDFLoader(file_path)
    
    if not os.path.exists(file_path):
        st.write(f"File does not exist: {file_path}")
        return

    st.session_state.docs = st.session_state.loader.load()
    
    if not st.session_state.docs:
        st.write("No doduments loaded.")
        return
    
    st.write(f"Loaded {len(st.session_state.docs)} documents.")
    
    st.session_state.text_splitter = RecursiveCharacterTextSplitter(chunk_size=7500, chunk_overlap=200)
    st.session_state.final_documents = st.session_state.text_splitter.split_documents(st.session_state.docs[:200])

    if not st.session_state.final_documents:
        st.write("No final documents after splitting.")
        return
    
    st.write(f"Created {len(st.session_state.final_documents)} document chunks.")
    
    for idx, doc in enumerate(st.session_state.final_documents):
        if 'id' not in doc.metadata or not doc.metadata['id']:
            doc.metadata['id'] = f"doc_{idx}"
    
    ids = [doc.metadata['id'] for doc in st.session_state.final_documents]
    
    
    st.session_state.vectors = Chroma.from_documents(
        documents=st.session_state.final_documents,
        embedding=st.session_state.embeddings,
        collection_name="local-rag",
        persist_directory=persist_directory
    )
    
    st.session_state.vectors.persist()
    st.write("Vectors saved to disak.")

I have tried loading vectors by:

    if os.path.exists(persist_directory):
        st.write("Loading vectors from disk...")
        st.session_state.vectors = Chroma(persist_directory=persist_directory, embedding_function=OllamaEmbeddings(model="nomic-embed-text"))
        st.write("Loaded vectors from disk.")
        return

However, when I run the app, the model does not have any context, and therefore doesn't answer appropriately.


Solution

  • Turns out I did not specified the collection name while loading the vectors, following solved my problem:

    st.session_state.vectors = Chroma(collection_name="local-rag",persist_directory=persist_directory, embedding_function=OllamaEmbeddings(model="nomic-embed-text"))