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

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,            # 指定输出图像的目录路径
)
1.
2.
3.
4.
5.
6.
7.
8.
9.
10.
11.
12.
13.
14.
15.
16.
17.
18.
19.

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

# 引入所需的库
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
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.
35.
36.
37.
38.
39.
40.
41.
42.
43.
44.
45.
46.
47.
48.
49.
50.

# 创建图像数据和图像摘要的列表
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)
1.
2.
3.
4.
5.
6.
7.
8.
9.
10.
11.
12.
13.
14.
15.
16.

image_summary_list
1.

# 定义一个基本模型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)))
1.
2.
3.
4.
5.
6.
7.
8.
9.
10.
11.
12.
13.
14.
15.
16.
17.
18.
19.
20.
21.

# 导入所需模块
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})
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.

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和内容存储在内存存储中
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.
35.
36.
37.
38.
39.
40.
41.
42.
43.
44.
45.
46.
47.

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)]
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.
35.
36.
37.
38.
39.
40.
41.
42.
43.
44.
45.
46.
47.
48.
49.
50.
51.
52.
53.
54.
55.
56.
57.
58.
59.
60.
61.
62.
63.
64.
65.
66.
67.
68.
69.
70.
71.
72.
73.
74.
75.
76.
77.

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)
1.
2.
3.
4.
5.
6.
7.
8.
9.
10.
11.
12.
13.
14.
15.
16.
17.
18.
19.
20.