PaddleX/paddlex/inference/pipelines_new/pp_chatocrv3_doc/pipeline.py

# copyright (c) 2024 PaddlePaddle Authors. All Rights Reserve.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#    http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from ..base import BasePipeline

from typing import Any, Dict, Optional

# import numpy as np
# import cv2
from .result import VisualInfoResult
import re

########## [TODO]后续需要更新路径
from ...components.transforms import ReadImage

import json

from ....utils import logging

from ...utils.pp_option import PaddlePredictorOption

from ..layout_parsing.result import LayoutParsingResult

import numpy as np


class PP_ChatOCRv3_doc_Pipeline(BasePipeline):
    """PP-ChatOCRv3-doc Pipeline"""

    entities = "PP-ChatOCRv3-doc"

    def __init__(
        self,
        config: Dict,
        device: str = None,
        pp_option: PaddlePredictorOption = None,
        use_hpip: bool = False,
        hpi_params: Optional[Dict[str, Any]] = None,
    ) -> None:
        """Initializes the pp-chatocrv3-doc pipeline.

        Args:
            config (Dict): Configuration dictionary containing various settings.
            device (str, optional): Device to run the predictions on. Defaults to None.
            pp_option (PaddlePredictorOption, optional): PaddlePredictor options. Defaults to None.
            use_hpip (bool, optional): Whether to use high-performance inference (hpip) for prediction. Defaults to False.
            hpi_params (Optional[Dict[str, Any]], optional): HPIP parameters. Defaults to None.
        """

        super().__init__(
            device=device, pp_option=pp_option, use_hpip=use_hpip, hpi_params=hpi_params
        )

        self.inintial_predictor(config)

        self.img_reader = ReadImage(format="BGR")

        self.table_structure_len_max = 500

    def inintial_predictor(self, config: dict) -> None:
        """
        Initializes the predictor with the given configuration.

        Args:
            config (dict): The configuration dictionary containing the necessary
                                parameters for initializing the predictor.
        Returns:
            None
        """
        layout_parsing_config = config["SubPipelines"]["LayoutParser"]
        self.layout_parsing_pipeline = self.create_pipeline(layout_parsing_config)

        from .. import create_chat_bot

        chat_bot_config = config["SubModules"]["LLM_Chat"]
        self.chat_bot = create_chat_bot(chat_bot_config)

        from .. import create_retriever

        retriever_config = config["SubModules"]["LLM_Retriever"]
        self.retriever = create_retriever(retriever_config)

        from .. import create_prompt_engeering

        text_pe_config = config["SubModules"]["PromptEngneering"]["KIE_CommonText"]
        self.text_pe = create_prompt_engeering(text_pe_config)

        table_pe_config = config["SubModules"]["PromptEngneering"]["KIE_Table"]
        self.table_pe = create_prompt_engeering(table_pe_config)

        return

    def decode_visual_result(
        self, layout_parsing_result: LayoutParsingResult
    ) -> VisualInfoResult:
        """
        Decodes the visual result from the layout parsing result.

        Args:
            layout_parsing_result (LayoutParsingResult): The result of layout parsing.

        Returns:
            VisualInfoResult: The decoded visual information.
        """
        text_paragraphs_ocr_res = layout_parsing_result["text_paragraphs_ocr_res"]
        seal_res_list = layout_parsing_result["seal_res_list"]
        normal_text_dict = {}

        for seal_res in seal_res_list:
            for text in seal_res["rec_text"]:
                layout_type = "印章"
                if layout_type not in normal_text_dict:
                    normal_text_dict[layout_type] = f"{text}"
                else:
                    normal_text_dict[layout_type] += f"\n {text}"

        for text in text_paragraphs_ocr_res["rec_text"]:
            layout_type = "words in text block"
            if layout_type not in normal_text_dict:
                normal_text_dict[layout_type] = text
            else:
                normal_text_dict[layout_type] += f"\n {text}"

        table_res_list = layout_parsing_result["table_res_list"]
        table_text_list = []
        table_html_list = []
        for table_res in table_res_list:
            table_html_list.append(table_res["pred_html"])
            single_table_text = " ".join(table_res["table_ocr_pred"]["rec_text"])
            table_text_list.append(single_table_text)

        visual_info = {}
        visual_info["normal_text_dict"] = normal_text_dict
        visual_info["table_text_list"] = table_text_list
        visual_info["table_html_list"] = table_html_list
        return VisualInfoResult(visual_info)

    # Function to perform visual prediction on input images
    def visual_predict(
        self,
        input: str | list[str] | np.ndarray | list[np.ndarray],
        use_doc_orientation_classify: bool = False,  # Whether to use document orientation classification
        use_doc_unwarping: bool = False,  # Whether to use document unwarping
        use_common_ocr: bool = True,  # Whether to use common OCR
        use_seal_recognition: bool = True,  # Whether to use seal recognition
        use_table_recognition: bool = True,  # Whether to use table recognition
        **kwargs,
    ) -> dict:
        """
        This function takes an input image or a list of images and performs various visual
        prediction tasks such as document orientation classification, document unwarping,
        common OCR, seal recognition, and table recognition based on the provided flags.

        Args:
            input (str | list[str] | np.ndarray | list[np.ndarray]): Input image path, list of image paths,
                                                                        numpy array of an image, or list of numpy arrays.
            use_doc_orientation_classify (bool): Flag to use document orientation classification.
            use_doc_unwarping (bool): Flag to use document unwarping.
            use_common_ocr (bool): Flag to use common OCR.
            use_seal_recognition (bool): Flag to use seal recognition.
            use_table_recognition (bool): Flag to use table recognition.
            **kwargs: Additional keyword arguments.

        Returns:
            dict: A dictionary containing the layout parsing result and visual information.
        """

        if not isinstance(input, list):
            input_list = [input]
        else:
            input_list = input

        img_id = 1
        for input in input_list:
            if isinstance(input, str):
                image_array = next(self.img_reader(input))[0]["img"]
            else:
                image_array = input

            assert len(image_array.shape) == 3

            layout_parsing_result = next(
                self.layout_parsing_pipeline.predict(
                    image_array,
                    use_doc_orientation_classify=use_doc_orientation_classify,
                    use_doc_unwarping=use_doc_unwarping,
                    use_common_ocr=use_common_ocr,
                    use_seal_recognition=use_seal_recognition,
                    use_table_recognition=use_table_recognition,
                )
            )

            visual_info = self.decode_visual_result(layout_parsing_result)

            visual_predict_res = {
                "layout_parsing_result": layout_parsing_result,
                "visual_info": visual_info,
            }
            yield visual_predict_res

    def save_visual_info_list(
        self, visual_info: VisualInfoResult, save_path: str
    ) -> None:
        """
        Save the visual info list to the specified file path.

        Args:
            visual_info (VisualInfoResult): The visual info result, which can be a single object or a list of objects.
            save_path (str): The file path to save the visual info list.

        Returns:
            None
        """
        if not isinstance(visual_info, list):
            visual_info_list = [visual_info]
        else:
            visual_info_list = visual_info

        with open(save_path, "w") as fout:
            fout.write(json.dumps(visual_info_list, ensure_ascii=False) + "\n")
        return

    def load_visual_info_list(self, data_path: str) -> list[VisualInfoResult]:
        """
        Loads visual info list from a JSON file.

        Args:
            data_path (str): The path to the JSON file containing visual info.

        Returns:
            list[VisualInfoResult]: A list of VisualInfoResult objects parsed from the JSON file.
        """
        with open(data_path, "r") as fin:
            data = fin.readline()
            visual_info_list = json.loads(data)
        return visual_info_list

    def merge_visual_info_list(
        self, visual_info_list: list[VisualInfoResult]
    ) -> tuple[list, list, list]:
        """
        Merge visual info lists.

        Args:
            visual_info_list (list[VisualInfoResult]): A list of visual info results.

        Returns:
            tuple[list, list, list]: A tuple containing three lists, one for normal text dicts,
                                               one for table text lists, and one for table HTML lists.
        """
        all_normal_text_list = []
        all_table_text_list = []
        all_table_html_list = []
        for single_visual_info in visual_info_list:
            normal_text_dict = single_visual_info["normal_text_dict"]
            table_text_list = single_visual_info["table_text_list"]
            table_html_list = single_visual_info["table_html_list"]
            all_normal_text_list.append(normal_text_dict)
            all_table_text_list.extend(table_text_list)
            all_table_html_list.extend(table_html_list)
        return all_normal_text_list, all_table_text_list, all_table_html_list

    def build_vector(
        self,
        visual_info: VisualInfoResult,
        min_characters: int = 3500,
        llm_request_interval: float = 1.0,
    ) -> dict:
        """
        Build a vector representation from visual information.

        Args:
            visual_info (VisualInfoResult): The visual information input, can be a single instance or a list of instances.
            min_characters (int): The minimum number of characters required for text processing, defaults to 3500.
            llm_request_interval (float): The interval between LLM requests, defaults to 1.0.

        Returns:
            dict: A dictionary containing the vector info and a flag indicating if the text is too short.
        """

        if not isinstance(visual_info, list):
            visual_info_list = [visual_info]
        else:
            visual_info_list = visual_info

        all_visual_info = self.merge_visual_info_list(visual_info_list)
        all_normal_text_list, all_table_text_list, all_table_html_list = all_visual_info

        vector_info = {}

        all_items = []
        for i, normal_text_dict in enumerate(all_normal_text_list):
            for type, text in normal_text_dict.items():
                all_items += [f"{type}：{text}\n"]

        for table_html, table_text in zip(all_table_html_list, all_table_text_list):
            if len(table_html) > min_characters - self.table_structure_len_max:
                all_items += [f"table：{table_text}\n"]

        all_text_str = "".join(all_items)

        if len(all_text_str) > min_characters:
            vector_info["flag_too_short_text"] = False
            vector_info["vector"] = self.retriever.generate_vector_database(all_items)
        else:
            vector_info["flag_too_short_text"] = True
            vector_info["vector"] = all_items
        return vector_info

    def format_key(self, key_list: str | list[str]) -> list[str]:
        """
        Formats the key list.

        Args:
            key_list (str|list[str]): A string or a list of strings representing the keys.

        Returns:
            list[str]: A list of formatted keys.
        """
        if key_list == "":
            return []

        if isinstance(key_list, list):
            return key_list

        if isinstance(key_list, str):
            key_list = re.sub(r"[\t\n\r\f\v]", "", key_list)
            key_list = key_list.replace("，", ",").split(",")
            return key_list

        return []

    def fix_llm_result_format(self, llm_result: str) -> dict:
        """
        Fix the format of the LLM result.

        Args:
            llm_result (str): The result from the LLM (Large Language Model).

        Returns:
            dict: A fixed format dictionary from the LLM result.
        """
        if not llm_result:
            return {}

        if "json" in llm_result or "```" in llm_result:
            llm_result = (
                llm_result.replace("```", "").replace("json", "").replace("/n", "")
            )
            llm_result = llm_result.replace("[", "").replace("]", "")

        try:
            llm_result = json.loads(llm_result)
            llm_result_final = {}
            for key in llm_result:
                value = llm_result[key]
                if isinstance(value, list):
                    if len(value) > 0:
                        llm_result_final[key] = value[0]
                else:
                    llm_result_final[key] = value
            return llm_result_final

        except:
            results = (
                llm_result.replace("\n", "")
                .replace("    ", "")
                .replace("{", "")
                .replace("}", "")
            )
            if not results.endswith('"'):
                results = results + '"'
            pattern = r'"(.*?)": "([^"]*)"'
            matches = re.findall(pattern, str(results))
            if len(matches) > 0:
                llm_result = {k: v for k, v in matches}
                return llm_result
            else:
                return {}

    def generate_and_merge_chat_results(
        self, prompt: str, key_list: list, final_results: dict, failed_results: dict
    ) -> None:
        """
        Generate and merge chat results into the final results dictionary.

        Args:
            prompt (str): The input prompt for the chat bot.
            key_list (list): A list of keys to track which results to merge.
            final_results (dict): The dictionary to store the final merged results.
            failed_results (dict): A dictionary of failed results to avoid merging.

        Returns:
            None
        """

        llm_result = self.chat_bot.generate_chat_results(prompt)
        if llm_result is None:
            logging.warning(
                "chat bot error: \n [prompt:]\n %s\n [result:] %s\n"
                % (prompt, self.chat_bot.ERROR_MASSAGE)
            )
            return

        llm_result = self.fix_llm_result_format(llm_result)
        for key, value in llm_result.items():
            if value not in failed_results and key in key_list:
                key_list.remove(key)
                final_results[key] = value
        return

    def chat(
        self,
        key_list: str | list[str],
        visual_info: VisualInfoResult,
        use_vector_retrieval: bool = True,
        vector_info: dict = None,
        min_characters: int = 3500,
        text_task_description: str = None,
        text_output_format: str = None,
        text_rules_str: str = None,
        text_few_shot_demo_text_content: str = None,
        text_few_shot_demo_key_value_list: str = None,
        table_task_description: str = None,
        table_output_format: str = None,
        table_rules_str: str = None,
        table_few_shot_demo_text_content: str = None,
        table_few_shot_demo_key_value_list: str = None,
    ) -> dict:
        """
        Generates chat results based on the provided key list and visual information.

        Args:
            key_list (str | list[str]): A single key or a list of keys to extract information.
            visual_info (VisualInfoResult): The visual information result.
            use_vector_retrieval (bool): Whether to use vector retrieval.
            vector_info (dict): The vector information for retrieval.
            min_characters (int): The minimum number of characters required.
            text_task_description (str): The description of the text task.
            text_output_format (str): The output format for text results.
            text_rules_str (str): The rules for generating text results.
            text_few_shot_demo_text_content (str): The text content for few-shot demos.
            text_few_shot_demo_key_value_list (str): The key-value list for few-shot demos.
            table_task_description (str): The description of the table task.
            table_output_format (str): The output format for table results.
            table_rules_str (str): The rules for generating table results.
            table_few_shot_demo_text_content (str): The text content for table few-shot demos.
            table_few_shot_demo_key_value_list (str): The key-value list for table few-shot demos.

        Returns:
            dict: A dictionary containing the chat results.
        """

        key_list = self.format_key(key_list)
        if len(key_list) == 0:
            return {"error": "输入的key_list无效！"}

        if not isinstance(visual_info, list):
            visual_info_list = [visual_info]
        else:
            visual_info_list = visual_info

        all_visual_info = self.merge_visual_info_list(visual_info_list)
        all_normal_text_list, all_table_text_list, all_table_html_list = all_visual_info

        final_results = {}
        failed_results = ["大模型调用失败", "未知", "未找到关键信息", "None", ""]

        for table_html, table_text in zip(all_table_html_list, all_table_text_list):
            if len(table_html) <= min_characters - self.table_structure_len_max:
                for table_info in [table_html, table_text]:
                    if len(key_list) > 0:
                        prompt = self.table_pe.generate_prompt(
                            table_info,
                            key_list,
                            task_description=table_task_description,
                            output_format=table_output_format,
                            rules_str=table_rules_str,
                            few_shot_demo_text_content=table_few_shot_demo_text_content,
                            few_shot_demo_key_value_list=table_few_shot_demo_key_value_list,
                        )

                        self.generate_and_merge_chat_results(
                            prompt, key_list, final_results, failed_results
                        )

        if len(key_list) > 0:
            if use_vector_retrieval and vector_info is not None:
                question_key_list = [f"抽取关键信息:{key}" for key in key_list]
                vector = vector_info["vector"]
                if not vector_info["flag_too_short_text"]:
                    related_text = self.retriever.similarity_retrieval(
                        question_key_list, vector
                    )
                    # print(question_key_list, related_text)
                else:
                    if len(vector) > 0:
                        related_text = "".join(vector)
                    else:
                        related_text = ""
            else:
                all_items = []
                for i, normal_text_dict in enumerate(all_normal_text_list):
                    for type, text in normal_text_dict.items():
                        all_items += [f"{type}：{text}\n"]
                related_text = "".join(all_items)
                if len(related_text) > min_characters:
                    logging.warning(
                        "The input text content is too long, the large language model may truncate it."
                    )

            if len(related_text) > 0:
                prompt = self.text_pe.generate_prompt(
                    related_text,
                    key_list,
                    task_description=text_task_description,
                    output_format=text_output_format,
                    rules_str=text_rules_str,
                    few_shot_demo_text_content=text_few_shot_demo_text_content,
                    few_shot_demo_key_value_list=text_few_shot_demo_key_value_list,
                )
                # print(prompt)
                self.generate_and_merge_chat_results(
                    prompt, key_list, final_results, failed_results
                )

        return {"chat_res": final_results}

    def predict(self, *args, **kwargs) -> None:
        logging.error(
            "PP-ChatOCRv3-doc Pipeline do not support to call `predict()` directly! Please invoke `visual_predict`, `build_vector`, `chat` sequentially to obtain the result."
        )
        return