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

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

    entities = "PP-ChatOCRv3-doc"
    def __init__(self,
        config,
        device=None,
        pp_option=None, 
        use_hpip: bool = False,
        hpi_params: Optional[Dict[str, Any]] = 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")
        
    def inintial_predictor(self, config):
        # layout_parsing_config = config['SubPipelines']["LayoutParser"]
        # self.layout_parsing_pipeline = self.create_pipeline(layout_parsing_config)

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

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

        text_pe_config = config['SubModules']['PromptEngneering']['KIE_CommonText']
        self.text_pe = self.create_prompt_engeering(text_pe_config)
        
        table_pe_config = config['SubModules']['PromptEngneering']['KIE_Table']
        self.table_pe = self.create_prompt_engeering(table_pe_config)

        return 

    def decode_visual_result(self, layout_parsing_result):
        text_paragraphs_ocr_res = layout_parsing_result['text_paragraphs_ocr_res']
        seal_res_list = layout_parsing_result['seal_res_list']
        normal_text_dict = {}
        layout_type = "text"
        for text in text_paragraphs_ocr_res['rec_text']:
            if layout_type not in normal_text_dict:
                normal_text_dict[layout_type] = text
            else:
                normal_text_dict[layout_type] += f"\n {text}"
        
        layout_type = "seal"
        for seal_res in seal_res_list:
            for text in seal_res['rec_text']:
                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)

    def visual_predict(self, input,
        use_doc_orientation_classify=True,
        use_doc_unwarping=True,
        use_common_ocr=True,
        use_seal_recognition=True,
        use_table_recognition=True,
        **kwargs):

        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, save_path):
        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):
        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):
        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, 
        min_characters=3500,
        llm_request_interval=1.0):

        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

        all_normal_text_str = "".join(["\n".join(e.values()) for e in all_normal_text_list])
        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}"]

        if len(all_normal_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):
        """format key"""
        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):
        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, key_list,
        final_results, failed_results):

        llm_result = self.chat_bot.generate_chat_results(prompt)
        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, visual_info, 
        key_list, 
        vector_info,
        text_task_description=None,
        text_output_format=None,
        text_rules_str=None,
        text_few_shot_demo_text_content=None,
        text_few_shot_demo_key_value_list=None,        
        table_task_description=None,
        table_output_format=None,
        table_rules_str=None,
        table_few_shot_demo_text_content=None,
        table_few_shot_demo_key_value_list=None):

        key_list = self.format_key(key_list)
        if len(key_list) == 0:
            return {"chat_res": "输入的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 all_table_info in [all_table_html_list, all_table_text_list]:
            for table_info in all_table_info:
                if len(key_list) == 0:
                    continue

                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:
            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)
            else:
                related_text = " ".join(vector)
            
            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)
            
            self.generate_and_merge_chat_results(prompt, 
                key_list, final_results, failed_results)

        return final_results

    def predict(self, *args, **kwargs):
        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