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