File size: 1,496 Bytes
931cdad
 
 
 
 
bfdcc17
931cdad
 
012a510
 
 
931cdad
012a510
 
 
931cdad
 
 
 
 
 
012a510
 
 
 
931cdad
 
 
 
012a510
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
import uuid
from langchain_chroma import Chroma
from langchain_core.stores import InMemoryStore
from langchain_core.documents import Document
from langchain_classic.retrievers.multi_vector import MultiVectorRetriever
from langchain_huggingface import HuggingFaceEmbeddings

def initialize_retriever():
    vectorstore = Chroma(
        collection_name="multi_modal_rag",
        embedding_function=HuggingFaceEmbeddings(model_name="sentence-transformers/all-MiniLM-L6-v2")
    )
    store = InMemoryStore()
    id_key = "doc_id"
    retriever = MultiVectorRetriever(vectorstore=vectorstore, docstore=store, id_key=id_key)
    return retriever, id_key

def add_assets_to_vectorstore(retriever, id_key, raw_elements, summaries):
    if not raw_elements:
        return
    doc_ids = [str(uuid.uuid4()) for _ in raw_elements]
    summary_documents = [
        Document(page_content=summary, metadata={id_key: doc_ids[i]})
        for i, summary in enumerate(summaries)
    ]
    retriever.vectorstore.add_documents(summary_documents)
    retriever.docstore.mset(list(zip(doc_ids, raw_elements)))

def run_database_pipeline(text_chunks, text_sums, tables, table_sums, images, image_sums):
    retriever, id_key = initialize_retriever()
    add_assets_to_vectorstore(retriever, id_key, text_chunks, text_sums)
    add_assets_to_vectorstore(retriever, id_key, tables, table_sums)
    add_assets_to_vectorstore(retriever, id_key, images, image_sums)
    return retriever