Spaces:
Running
Running
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 |