pip install evalscope[rag] # 提供RAG评估依赖

modelscope download --dataset modelscope/plants-image

pip install "unstructured[all-docs]"
apt install poppler-utils tesseract-ocr -y

from unstructured.partition.pdf import partition_pdf
import os
import json

file_path = "./pdf/plants.pdf"  # PDF文件路径
image_path = "./pdf/images"  # 图片存储路径

def extract_pdf_elements(file_path, output_image_path):
    """
    从PDF文件中提取图片、表格和文本块。
    output_image_path: 文件路径，用于存储图片（.jpg）
    file_path: 文件名
    """
    os.makedirs(output_image_path, exist_ok=True)
    return partition_pdf(
        filename=file_path,
        strategy="hi_res",
        extract_images_in_pdf=True,
        infer_table_structure=False,
        chunking_strategy="by_title",
        max_characters=1000,
        new_after_n_chars=1000,
        combine_text_under_n_chars=1000,
        hi_res_model_name="yolox",
        extract_image_block_output_dir=output_image_path,
    )

import os
import numpy as np
from langchain_chroma import Chroma
from evalscope.backend.rag_eval import VisionModel

# 创建 Chroma 实例
vectorstore = Chroma(
    collection_name="mm_rag_clip",
    # 使用ModelScope上的CLIP模型计算向量
    embedding_function=VisionModel.load(model_name="AI-ModelScope/chinese-clip-vit-large-patch14-336px")  
)

# 获取所有.jpg扩展名的图片URI
image_uris = sorted(
    [
        # 获取.jpg扩展名的文件
        os.path.join(image_path, image_name)
        for image_name in os.listdir(image_path)  
        if image_name.endswith(".jpg")  
    ]
)

# 添加图片到Chroma向量存储，也可通过add_texts添加文本向量，本文不需要
vectorstore.add_images(uris=image_uris)

# 创建检索器实例
retriever = vectorstore.as_retriever(
        search_type="similarity",  # 使用相似性搜索类型
        search_kwargs={'k': 2}  # 搜索结果返回前2个最相似项（top_k = 2）
)

docs = retriever.invoke("玉米")

import os
import json
import base64
import matplotlib.pyplot as plt
from io import BytesIO
from PIL import Image
from langchain_core.messages import HumanMessage
from langchain_core.output_parsers import StrOutputParser
from langchain_core.runnables import RunnableLambda, RunnablePassthrough

class RAGPipeline:
    def __init__(self, retriever, model):
        self.retriever = retriever
        self.model = model
        self.question = None
        self.context = None
        self.response = None

    def split_image_text_types(self, docs):
        """
        将文档内容分为图片和文本两类。
        docs: 文档列表
        """
        images = []
        text = []
        for doc in docs:
            doc = doc.page_content  # 提取文档内容
            if self.is_base64(doc):
                images.append(doc)
            else:
                text.append(doc)
        self.context = {"images": images, "texts": text}
        return self.context

    def prompt_func(self, data_dict):
        """
        生成提示信息，将上下文文本合并为一个字符串。
        data_dict: 包含上下文信息的数据字典
        """
        # 将上下文文本合并为一个字符串
        formatted_texts = "\n".join(data_dict["context"]["texts"])
        messages = [{"type": "text", "text": (
            f"""你是一位植物图鉴科普者，专门回答有关各种植物的问题。请根据用户输入的问题和检索到的图像提供详细且准确的回答，详细描述检索到的图片细节，结合这些细节，请确保你的回答足够详细，以方便用户理解。以下是用户的问题和相关的上下文信息：
用户问题: {self.question}
检索到的上下文信息:
{formatted_texts}
"""
        )}]

        # 如果有图片，将其添加到消息中
        for image_base64 in data_dict["context"]["images"]:
            messages.append({
                "type": "image_url",
                "image_url": {"url": f"data:image/jpeg;base64,{image_base64}"}
            })

        return [HumanMessage(content=messages)]


    def is_base64(self, s):
        """
        检查字符串是否为base64编码。
        s: 输入字符串
        """
        try:
            return base64.b64encode(base64.b64decode(s)) == s.encode()
        except Exception:
            return False

    def invoke(self, question):
        """
        调用链式操作执行检索和生成回答。
        question: 用户问题
        """
        self.question = question
        chain = (
            {
                "context": self.retriever | RunnableLambda(self.split_image_text_types),
                "question": RunnablePassthrough(),
            }
            | RunnableLambda(self.prompt_func)
            | self.model
            | StrOutputParser()
        )

        self.response = chain.invoke(question)
        return self.response

    def show_result(self):
        """
        显示最终结果，包括问题、上下文和生成的回答。
        """
        print(f"Question: {self.question}")
        print(f"Context: {self.context}")
        print(f"Response: {self.response}")
        for image_base64 in self.context.get("images", []):
            self.plt_img_base64(image_base64)

    def save_case(self, file_path):
        """
        保存案例数据到指定文件路径。
        file_path: 保存路径
        """
        case_data = {
            "user_input": self.question,
            "retrieved_contexts": self.context['texts'] + self.context['images'],
            "response": self.response
        }

        if os.path.exists(file_path):
            # 读取现有的案例数据
            with open(file_path, 'r', encoding='utf-8') as f:
                cases = json.load(f)
        else:
            cases = []

        # 添加新的案例数据
        cases.append(case_data)

        # 写回所有案例数据
        with open(file_path, 'w', encoding='utf-8') as f:
            json.dump(cases, f, ensure_ascii=False, indent=4)

from evalscope.backend.rag_eval import LLM

model = LLM.load(
            model_name="qwen-vl-plus",
            api_base="https://dashscope.aliyuncs.com/compatible-mode/v1",
            api_key=os.getenv("DASHSCOPE_API_KEY"),
        )

pipeline = RAGPipeline(retriever, model)
for prompt in ["玉米", "西兰花", "洋芋", "水稻", "汽车"]:
    response = pipeline.invoke(f"{prompt}是什么？")
    pipeline.show_result()
    pipeline.save_case("./pdf/output.json")

[
    {
        "user_input": "玉米是什么？",
        "retrieved_contexts": [
          "/9j/4AAQSkZJRgABAQAAAQABAAD/2wBDAAgGB...", # base64 图像
          "/9j/4AAQSkZJRgABAQAAAQABAAD/2wBDAAgGB...", # base64 图像
        ],
        "response": "玉米是一种禾本科作物，在全球范围内广泛种植作为粮食、饲料或工业原料使用。\n\n在第一张图片中，我们可以看到一个正在生长中的玉米植株。它的叶子是绿色的，并带有明显的叶脉结构。茎部呈现出鲜艳的黄色，这是由于它含有丰富的类胡萝卜素色素所致。在这根茎上，我们还可以观察到一些红色的毛状物，这实际上是玉米花序的一部分，也就是将来会发育成玉米棒的地方。\n\n第二张图片则展示了两个不同阶段的玉米植株。左边是一颗已经开花并开始形成穗轴（即玉米棒）的成熟期玉米；右边则是另一棵刚刚进入抽穗生长期的小苗，可以看到其叶片较小而颜色较浅，与成熟的玉米相比显得更为稚嫩。\n\n总的来说，无论是从形态还是生命周期的角度来看，这两幅图片都为我们提供了关于玉米这种重要农作物的重要视觉参考点。"
    },
  ...
]

multi_modal_task_cfg = {
    "eval_backend": "RAGEval",
    "eval_config": {
        "tool": "RAGAS",
        "eval": {
            "testset_file": "./pdf/output.json",
            "critic_llm": {
                "model_name": "qwen-vl-max",
                "api_base": "https://dashscope.aliyuncs.com/compatible-mode/v1",
                "api_key": os.getenv("DASHSCOPE_API_KEY"),
            },
            "embeddings": {
                "model_name_or_path": "AI-ModelScope/bge-large-zh",
            },
            "metrics": [
                "MultiModalFaithfulness",
                "MultiModalRelevance",
                "AnswerCorrectness"
            ],
            "language": "chinese"
        },
    },
}

from evalscope.run import run_task

run_task(task_cfg=multi_modal_task_cfg)