视觉文档多模态RAG系统构建-Step by Step(附完整代码)

python3.10 -m venv venv

pip install -r requirements.txt

#%brew install tesseract poppler%pip install -q "unstructured[all-docs]"

from unstructured.partition.pdf import partition_pdf


elements = partition_pdf(filename="TAGIV.pdf", # mandatorystrategy="hi_res", # mandatory to use ``hi_res`` strategyextract_images_in_pdf=True,# mandatory to set as ``True``extract_image_block_types=["Image", "Table"],# optionalextract_image_block_to_payload=False,# optionalextract_image_block_output_dir="saved_images",# optional - only works when ``extract_image_block_to_payload=False``)

from unstructured.chunking.title import chunk_by_title # might be better for an article from typing import Any
chunks = chunk_by_title(elements)Output:# 文档中的不同类别category_counts = {}
for element in chunks:category = str(type(element))if category in category_counts:category_counts[category] += 1else:category_counts[category] = 1{"": 200, "": 3, "": 2}

from pydantic import BaseModelclass Element(BaseModel):type: strtext: Any

categorized_elements = []for element in chunks:

 if "unstructured.documents.elements.CompositeElement" in str(type(element)):

 categorized_elements.append(Element(type="text", text=str(element)))

 elif "unstructured.documents.elements.Table" in str(type(element)): 

 categorized_elements.append(Element(type="table", text=str(element)))

%pip install -q langchain langchain-chroma unstructured[all-docs] pydantic lxml langchainhub langchain-openai
from langchain_core.output_parsers import StrOutputParserfrom langchain_core.prompts import ChatPromptTemplatefrom langchain_openai import ChatOpenAI

prompt = ChatPromptTemplate.from_template(prompt_text)

model = ChatOpenAI(temperature=0, model="gpt-4o")summarize_chain = {"element": lambda x: x} | prompt | model | StrOutputParser()

texts = [i.text for i in text_elements]text_summaries = summarize_chain.batch(texts, {"max_concurrency": 5})Apply to tablestables = [i.text for i in table_elements]table_summaries = summarize_chain.batch(tables, {"max_concurrency": 5})

## 获取图像摘要import base64import osfrom langchain_core.messages import HumanMessage
def encode_image(image_path):"""获取base64字符串"""with open(image_path, "rb") as image_file:return base64.b64encode(image_file.read()).decode("utf-8")def image_summarize(img_base64, prompt):"""生成图像摘要"""chat = ChatOpenAI(model="gpt-4o", max_tokens=1024)
msg = chat.invoke([HumanMessage(content=[{"type": "text", "text": prompt},{"type": "image_url","image_url": {"url": f"data:image/jpg;base64,{img_base64}"},},])])return msg.content
def generate_img_summaries(path):"""生成图像摘要和base64编码字符串path: 通过Unstructured提取的.jpg文件列表的路径"""
# 存储base64编码的图像img_base64_list = []
# 存储图像摘要image_summaries = []
# 提示prompt = """您是一名助手，负责为检索摘要图像。\这些摘要将被嵌入并用于检索原始图像。\给出一份简洁的图像摘要，以便更好地进行检索。"""
# 应用于图像for img_file in sorted(os.listdir(path)):if img_file.endswith(".jpg"):img_path = os.path.join(path, img_file)base64_image = encode_image(img_path)img_base64_list.append(base64_image)image_summaries.append(image_summarize(base64_image, prompt))
return img_base64_list, image_summaries





fpath = "saved_images" 图像摘要img_base64_list, image_summaries = generate_img_summaries(fpath)

import uuid
from langchain.retrievers.multi_vector import MultiVectorRetriever from langchain.storage import InMemoryStorefrom langchain_chroma import Chromafrom langchain_core.documents import Documentfrom langchain_openai import OpenAIEmbeddingsdef create_multi_vector_retriever(vectorstore, text_summaries, texts, table_summaries, tables, image_summaries, images):"""创建检索器，索引摘要，但返回原始图像、表格或文本"""


# 初始化存储层store = InMemoryStore()id_key = "doc_id"
# 创建多向量检索器retriever = MultiVectorRetriever(vectorstore=vectorstore,docstore=store,id_key=id_key,search_kwargs={"k": 2}# 限制为前 2 个结果)
# 添加文档到向量存储和文档存储的辅助函数def add_documents(retriever, doc_summaries, doc_contents):doc_ids = [str(uuid.uuid4()) for _ in doc_contents]summary_docs = [Document(page_content=s, metadata={id_key: doc_ids[i]}) for i, s inenumerate(doc_summaries)]retriever.vectorstore.add_documents(summary_docs)retriever.docstore.mset(list(zip(doc_ids, doc_contents)))


# 添加文本、表格和图像# 在添加之前检查 text_summaries 不为空if text_summaries:add_documents(retriever, text_summaries, texts)# 在添加之前检查 table_summaries 不为空if table_summaries:add_documents(retriever, table_summaries, tables)# 在添加之前检查 image_summaries 不为空if image_summaries:add_documents(retriever, image_summaries, images)


return retriever

# 用于索引摘要的向量存储vectorstore = Chroma(collection_name="mm_tagiv_paper", embedding_function=OpenAIEmbeddings())


# 创建检索器retriever_multi_vector_img = create_multi_vector_retriever(vectorstore,text_summaries,texts,table_summaries,tables,image_summaries,img_base64_list,)