从RAG到多模态搜索看OpenAI的技术演进

!pip install langchain unstructured[all-docs] pydantic lxml openai chromadb tiktoken -q -U
!apt-get install poppler-utils tesseract-ocr

from typing import Any

from pydantic import BaseModel
from unstructured.partition.pdf import partition_pdf

# 设置存放图像的路径
images_path = "./images"

# 调用 partition_pdf 函数处理PDF文件，提取其中的元素
raw_pdf_elements = partition_pdf(
    filename="weekly-market-recap.pdf",           # 指定要处理的PDF文件名
    extract_images_in_pdf=True,                   # 设置为True以从PDF中提取图像
    infer_table_structure=True,                   # 设置为True以推断PDF中的表格结构
    chunking_strategy="by_title",                 # 设置文档切分策略为按标题切分
    max_characters=4000,                          # 设置每个块的最大字符数为4000
    new_after_n_chars=3800,                       # 在达到3800字符后开始新的块
    combine_text_under_n_chars=2000,              # 将少于2000字符的文本组合在一起
    image_output_dir_path=images_path,            # 指定输出图像的目录路径
)

from IPython.display import Image
Image('images/figure-1-1.jpg')

# 引入所需的库
import base64
import os

# 从langchain包导入ChatOpenAI类和HumanMessage模块
from langchain.chat_models import ChatOpenAI
from langchain.schema.messages import HumanMessage

# 定义图像摘要类
class ImageSummarizer:

    # 初始化函数，设置图像路径
    def __init__(self, image_path) -> None:
        self.image_path = image_path
        self.prompt = """你的任务是将图像内容生成摘要信息，以便检索。
        这些摘要将被嵌入并用于检索原始图像。请给出一个简洁的图像摘要，以便于检索优化。
"""

    # 定义函数将图像转换为base64编码
    def base64_encode_image(self):
        with open(self.image_path, "rb") as image_file:
            return base64.b64encode(image_file.read()).decode("utf-8")

    # 定义摘要函数
    def summarize(self, prompt = None):
        # 获取图像的base64编码数据
        base64_image_data = self.base64_encode_image()
        # 创建ChatOpenAI对象，使用gpt-4-vision预览模型，最大token数为1000
        chat = ChatOpenAI(model="gpt-4-vision-preview", max_tokens=1000)

        # 调用chat对象的invoke方法，发送包含文本和图像的消息
        response = chat.invoke(
            [
                #人类提示语的输入
                HumanMessage(
                    content=[
                        {
                            "type": "text",
                            "text": prompt if prompt else self.prompt
                        },
                        {
                            "type": "image_url",
                            "image_url": {"url": f"data:image/jpeg;base64,{base64_image_data}"},
                        },
                    ]
                )
            ]
        )
        # 返回base64编码的图像数据和响应内容
        return base64_image_data, response.content

# 创建图像数据和图像摘要的列表
image_data_list = []
image_summary_list = []

# 遍历指定路径下的所有文件
for img_file in sorted(os.listdir(images_path)):
    # 检查文件扩展名是否为.jpg
    if img_file.endswith(".jpg"):
        # 创建ImageSummarizer对象
        summarizer = ImageSummarizer(os.path.join(images_path, img_file))
        # 调用summarize方法生成摘要
        data, summary = summarizer.summarize()
        # 将base64编码的图像数据添加到列表
        image_data_list.append(data)
        # 将图像摘要添加到列表
        image_summary_list.append(summary)

image_summary_list

# 定义一个基本模型Element，用于存储文档元素的类型和文本内容
class Element(BaseModel):
    type: str
    text: Any

# 创建表格元素和文本元素的空列表
#表格信息
table_elements = []
#文本信息
text_elements = []

# 遍历从PDF中分割出的所有元素
for element in raw_pdf_elements:
    # 判断元素是否为表格：非结构化-文档-元素-表格
    if "unstructured.documents.elements.Table" in str(type(element)):
        # 如果是表格，将其添加到表格元素列表
        table_elements.append(Element(type="table", text=str(element)))
    # 判断元素是否为复合元素
    elif "unstructured.documents.elements.CompositeElement" in str(type(element)):
        # 如果是复合元素，将其添加到文本元素列表
        text_elements.append(Element(type="text", text=str(element)))

# 导入所需模块
from langchain.chat_models import ChatOpenAI
from langchain.prompts import ChatPromptTemplate
from langchain.schema.output_parser import StrOutputParser

# 定义用于生成总结的提示文本
prompt_text = """
  您负责简洁地总结表格或文本块，中文输出。:
  {element}
"""
# 从模板创建提示
prompt = ChatPromptTemplate.from_template(prompt_text)

# 创建总结链，结合提示和GPT-3.5模型
summarize_chain = {"element": lambda x: x} | prompt | ChatOpenAI(temperature=0, model="gpt-3.5-turbo") | StrOutputParser()

# 从表格元素中提取文本
tables = [i.text for i in table_elements]
# 使用总结链批量处理表格文本，设置最大并发数为5
table_summaries = summarize_chain.batch(tables, {"max_concurrency": 5})

# 从文本元素中提取文本
texts = [i.text for i in text_elements]
# 使用总结链批量处理文本，设置最大并发数为5
text_summaries = summarize_chain.batch(texts, {"max_concurrency": 5})

import uuid  # 导入uuid库，用于生成唯一标识符

# 导入所需的langchain模块
from langchain.embeddings import OpenAIEmbeddings
from langchain.retrievers.multi_vector import MultiVectorRetriever
from langchain.schema.document import Document
from langchain.storage import InMemoryStore
from langchain.vectorstores import Chroma

id_key = "doc_id"  # 设置文档的唯一标识键

# 初始化多向量检索器，初始为空，文字，图片，表格
retriever = MultiVectorRetriever(
    vectorstore=Chroma(collection_name="summaries", embedding_function=OpenAIEmbeddings()),  # 使用OpenAI的嵌入方法
    docstore=InMemoryStore(),  # 使用内存存储文档
    id_key=id_key,  # 设置文档的唯一标识键
)

# 为文本添加唯一标识
doc_ids = [str(uuid.uuid4()) for _ in texts]  # 为每个文本生成一个唯一的UUID
# 创建文本的文档对象并添加到检索器
summary_texts = [
    Document(page_content=s, metadata={id_key: doc_ids[i]})  # 将文本封装为文档对象
    for i, s in enumerate(text_summaries)
]
retriever.vectorstore.add_documents(summary_texts)  # 将文本文档添加到向量存储中
retriever.docstore.mset(list(zip(doc_ids, texts)))  # 将文本文档的ID和内容存储在内存存储中

# 为表格添加唯一标识
table_ids = [str(uuid.uuid4()) for _ in tables]  # 为每个表格生成一个唯一的UUID
# 创建表格的文档对象并添加到检索器
summary_tables = [
    Document(page_content=s, metadata={id_key: table_ids[i]})  # 将表格封装为文档对象
    for i, s in enumerate(table_summaries)
]
retriever.vectorstore.add_documents(summary_tables)  # 将表格文档添加到向量存储中
retriever.docstore.mset(list(zip(table_ids, tables)))  # 将表格文档的ID和内容存储在内存存储中

# 为图像添加唯一标识
doc_ids = [str(uuid.uuid4()) for _ in image_data_list]  # 为每个图像生成一个唯一的UUID
# 创建图像的文档对象并添加到检索器
summary_images = [
    Document(page_content=s, metadata={id_key: doc_ids[i]})  # 将图像封装为文档对象
    for i, s in enumerate(image_summary_list)
]
retriever.vectorstore.add_documents(summary_images)  # 将图像文档添加到向量存储中
retriever.docstore.mset(list(zip(doc_ids, image_data_list)))  # 将图像文档的ID和内容存储在内存存储中

from PIL import Image
from IPython.display import HTML, display
import io
import re

# 显示基于base64编码的图片
def plt_img_base64(img_base64):
    display(HTML(f''))

# 检查base64数据是否为图片
def is_image_data(b64data):
    """
    通过查看数据开头来检查base64数据是否为图片
    """
    image_signatures = {
        b"\xFF\xD8\xFF": "jpg",
        b"\x89\x50\x4E\x47\x0D\x0A\x1A\x0A": "png",
        b"\x47\x49\x46\x38": "gif",
        b"\x52\x49\x46\x46": "webp",
    }
    try:
        header = base64.b64decode(b64data)[:8]  # 解码并获取前8个字节
        for sig, format in image_signatures.items():
            if header.startswith(sig):
                return True
        return False
    except Exception:
        return False

# 分离base64编码的图片和文本
def split_image_text_types(docs):
    """
    分离base64编码的图片和文本
    """
    b64_images = []
    texts = []
    for doc in docs:
        # 检查文档是否为Document类型并提取page_content
        if isinstance(doc, Document):
            doc = doc.page_content
        #是图片
        if is_image_data(doc):
            b64_images.append(doc)
        else:
            texts.append(doc)
    return {"images": b64_images, "texts": texts}

# 根据数据生成提示信息
def img_prompt_func(data_dict):
    messages = []

    # 如果存在图片，则添加到消息中
    #图片的信息单独拿出来， 需要和提示语一起传给大模型的。
    if data_dict["context"]["images"]:
        for image in data_dict["context"]["images"]:
            image_message = {
                "type": "image_url",
                "image_url": {"url": f"data:image/jpeg;base64,{image}"},
            }
            messages.append(image_message)

    # 添加文本到消息
    formatted_texts = "\n".join(data_dict["context"]["texts"])
    text_message = {
        "type": "text",
        "text": (
            "You are financial analyst.\n"
            "You will be given a mixed of text, tables, and image(s) usually of charts or graphs.\n"
            "Use this information to answer the user question in the finance. \n"
            f"Question: {data_dict['question']}\n\n"
            "Text and / or tables:\n"
            f"{formatted_texts}"
        ),
    }
    messages.append(text_message)
    print(messages)
    return [HumanMessage(content=messages)]

from langchain.schema.runnable import RunnableLambda, RunnablePassthrough
# 创建ChatOpenAI模型实例
model = ChatOpenAI(temperature=0, model="gpt-4-vision-preview", max_tokens=1024)

# 构建RAG（Retrieval Augmented Generation）管道
chain = (
    {
        "context": retriever | RunnableLambda(split_image_text_types), # 使用检索器获取内容并分离图片和文本
        "question": RunnablePassthrough(), # 传递问题
    }
    | RunnableLambda(img_prompt_func) # 根据图片和文本生成提示信息
    | model # 使用ChatOpenAI模型生成回答
    | StrOutputParser() # 解析模型的字符串输出
)

# 定义查询问题
query = "Which year had the highest holiday sales growth?"

# 调用chain的invoke方法执行查询
chain.invoke(query)