pipeline_v4.py 30 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723724725726727728729730731732733734735736737738739740
  1. # copyright (c) 2024 PaddlePaddle Authors. All Rights Reserve.
  2. #
  3. # Licensed under the Apache License, Version 2.0 (the "License");
  4. # you may not use this file except in compliance with the License.
  5. # You may obtain a copy of the License at
  6. #
  7. # http://www.apache.org/licenses/LICENSE-2.0
  8. #
  9. # Unless required by applicable law or agreed to in writing, software
  10. # distributed under the License is distributed on an "AS IS" BASIS,
  11. # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  12. # See the License for the specific language governing permissions and
  13. # limitations under the License.
  14. from typing import Any, Dict, Optional
  15. import re
  16. import cv2
  17. import json
  18. import base64
  19. import numpy as np
  20. import copy
  21. from .pipeline_base import PP_ChatOCR_Pipeline
  22. from ...common.reader import ReadImage
  23. from ...common.batch_sampler import ImageBatchSampler
  24. from ....utils import logging
  25. from ...utils.pp_option import PaddlePredictorOption
  26. from ..layout_parsing.result import LayoutParsingResult
  27. class PP_ChatOCRv4_Pipeline(PP_ChatOCR_Pipeline):
  28. """PP-ChatOCRv4 Pipeline"""
  29. entities = ["PP-ChatOCRv4-doc"]
  30. def __init__(
  31. self,
  32. config: Dict,
  33. device: str = None,
  34. pp_option: PaddlePredictorOption = None,
  35. use_hpip: bool = False,
  36. ) -> None:
  37. """Initializes the pp-chatocrv3-doc pipeline.
  38. Args:
  39. config (Dict): Configuration dictionary containing various settings.
  40. device (str, optional): Device to run the predictions on. Defaults to None.
  41. pp_option (PaddlePredictorOption, optional): PaddlePredictor options. Defaults to None.
  42. use_hpip (bool, optional): Whether to use high-performance inference (hpip) for prediction. Defaults to False.
  43. use_layout_parsing (bool, optional): Whether to use layout parsing. Defaults to True.
  44. """
  45. super().__init__(device=device, pp_option=pp_option, use_hpip=use_hpip)
  46. self.pipeline_name = config["pipeline_name"]
  47. self.inintial_predictor(config)
  48. self.batch_sampler = ImageBatchSampler(batch_size=1)
  49. self.img_reader = ReadImage(format="BGR")
  50. self.table_structure_len_max = 500
  51. def inintial_predictor(self, config: dict) -> None:
  52. """
  53. Initializes the predictor with the given configuration.
  54. Args:
  55. config (dict): The configuration dictionary containing the necessary
  56. parameters for initializing the predictor.
  57. Returns:
  58. None
  59. """
  60. self.use_layout_parser = config.get("use_layout_parser", True)
  61. if self.use_layout_parser:
  62. layout_parsing_config = config.get("SubPipelines", {}).get(
  63. "LayoutParser",
  64. {"pipeline_config_error": "config error for layout_parsing_pipeline!"},
  65. )
  66. self.layout_parsing_pipeline = self.create_pipeline(layout_parsing_config)
  67. from .. import create_chat_bot
  68. chat_bot_config = config.get("SubModules", {}).get(
  69. "LLM_Chat",
  70. {"chat_bot_config_error": "config error for llm chat bot!"},
  71. )
  72. self.chat_bot = create_chat_bot(chat_bot_config)
  73. from .. import create_retriever
  74. retriever_config = config.get("SubModules", {}).get(
  75. "LLM_Retriever",
  76. {"retriever_config_error": "config error for llm retriever!"},
  77. )
  78. self.retriever = create_retriever(retriever_config)
  79. from .. import create_prompt_engeering
  80. text_pe_config = (
  81. config.get("SubModules", {})
  82. .get("PromptEngneering", {})
  83. .get(
  84. "KIE_CommonText",
  85. {"pe_config_error": "config error for text_pe!"},
  86. )
  87. )
  88. self.text_pe = create_prompt_engeering(text_pe_config)
  89. table_pe_config = (
  90. config.get("SubModules", {})
  91. .get("PromptEngneering", {})
  92. .get(
  93. "KIE_Table",
  94. {"pe_config_error": "config error for table_pe!"},
  95. )
  96. )
  97. self.table_pe = create_prompt_engeering(table_pe_config)
  98. self.use_mllm_predict = config.get("use_mllm_predict", True)
  99. if self.use_mllm_predict:
  100. mllm_chat_bot_config = config.get("SubModules", {}).get(
  101. "MLLM_Chat",
  102. {"mllm_chat_bot_config": "config error for mllm chat bot!"},
  103. )
  104. self.mllm_chat_bot = create_chat_bot(mllm_chat_bot_config)
  105. ensemble_pe_config = (
  106. config.get("SubModules", {})
  107. .get("PromptEngneering", {})
  108. .get(
  109. "Ensemble",
  110. {"pe_config_error": "config error for ensemble_pe!"},
  111. )
  112. )
  113. self.ensemble_pe = create_prompt_engeering(ensemble_pe_config)
  114. return
  115. def decode_visual_result(self, layout_parsing_result: LayoutParsingResult) -> dict:
  116. """
  117. Decodes the visual result from the layout parsing result.
  118. Args:
  119. layout_parsing_result (LayoutParsingResult): The result of layout parsing.
  120. Returns:
  121. dict: The decoded visual information.
  122. """
  123. text_paragraphs_ocr_res = layout_parsing_result["text_paragraphs_ocr_res"]
  124. seal_res_list = layout_parsing_result["seal_res_list"]
  125. normal_text_dict = {}
  126. for seal_res in seal_res_list:
  127. for text in seal_res["rec_texts"]:
  128. layout_type = "印章"
  129. if layout_type not in normal_text_dict:
  130. normal_text_dict[layout_type] = f"{text}"
  131. else:
  132. normal_text_dict[layout_type] += f"\n {text}"
  133. for text in text_paragraphs_ocr_res["rec_texts"]:
  134. layout_type = "words in text block"
  135. if layout_type not in normal_text_dict:
  136. normal_text_dict[layout_type] = text
  137. else:
  138. normal_text_dict[layout_type] += f"\n {text}"
  139. table_res_list = layout_parsing_result["table_res_list"]
  140. table_text_list = []
  141. table_html_list = []
  142. table_nei_text_list = []
  143. for table_res in table_res_list:
  144. table_html_list.append(table_res["pred_html"])
  145. single_table_text = " ".join(table_res["table_ocr_pred"]["rec_texts"])
  146. table_text_list.append(single_table_text)
  147. table_nei_text_list.append(table_res["neighbor_texts"])
  148. visual_info = {}
  149. visual_info["normal_text_dict"] = normal_text_dict
  150. visual_info["table_text_list"] = table_text_list
  151. visual_info["table_html_list"] = table_html_list
  152. visual_info["table_nei_text_list"] = table_nei_text_list
  153. return visual_info
  154. # Function to perform visual prediction on input images
  155. def visual_predict(
  156. self,
  157. input: str | list[str] | np.ndarray | list[np.ndarray],
  158. use_doc_orientation_classify: Optional[bool] = None,
  159. use_doc_unwarping: Optional[bool] = None,
  160. use_general_ocr: Optional[bool] = None,
  161. use_seal_recognition: Optional[bool] = None,
  162. use_table_recognition: Optional[bool] = None,
  163. text_det_limit_side_len: Optional[int] = None,
  164. text_det_limit_type: Optional[str] = None,
  165. text_det_thresh: Optional[float] = None,
  166. text_det_box_thresh: Optional[float] = None,
  167. text_det_unclip_ratio: Optional[float] = None,
  168. text_rec_score_thresh: Optional[float] = None,
  169. seal_det_limit_side_len: Optional[int] = None,
  170. seal_det_limit_type: Optional[str] = None,
  171. seal_det_thresh: Optional[float] = None,
  172. seal_det_box_thresh: Optional[float] = None,
  173. seal_det_unclip_ratio: Optional[float] = None,
  174. seal_rec_score_thresh: Optional[float] = None,
  175. **kwargs,
  176. ) -> dict:
  177. """
  178. This function takes an input image or a list of images and performs various visual
  179. prediction tasks such as document orientation classification, document unwarping,
  180. general OCR, seal recognition, and table recognition based on the provided flags.
  181. Args:
  182. input (str | list[str] | np.ndarray | list[np.ndarray]): Input image path, list of image paths,
  183. numpy array of an image, or list of numpy arrays.
  184. use_doc_orientation_classify (bool): Flag to use document orientation classification.
  185. use_doc_unwarping (bool): Flag to use document unwarping.
  186. use_general_ocr (bool): Flag to use general OCR.
  187. use_seal_recognition (bool): Flag to use seal recognition.
  188. use_table_recognition (bool): Flag to use table recognition.
  189. **kwargs: Additional keyword arguments.
  190. Returns:
  191. dict: A dictionary containing the layout parsing result and visual information.
  192. """
  193. if self.use_layout_parser == False:
  194. logging.error("The models for layout parser are not initialized.")
  195. yield {"error": "The models for layout parser are not initialized."}
  196. for layout_parsing_result in self.layout_parsing_pipeline.predict(
  197. input,
  198. use_doc_orientation_classify=use_doc_orientation_classify,
  199. use_doc_unwarping=use_doc_unwarping,
  200. use_general_ocr=use_general_ocr,
  201. use_seal_recognition=use_seal_recognition,
  202. use_table_recognition=use_table_recognition,
  203. text_det_limit_side_len=text_det_limit_side_len,
  204. text_det_limit_type=text_det_limit_type,
  205. text_det_thresh=text_det_thresh,
  206. text_det_box_thresh=text_det_box_thresh,
  207. text_det_unclip_ratio=text_det_unclip_ratio,
  208. text_rec_score_thresh=text_rec_score_thresh,
  209. seal_det_box_thresh=seal_det_box_thresh,
  210. seal_det_limit_side_len=seal_det_limit_side_len,
  211. seal_det_limit_type=seal_det_limit_type,
  212. seal_det_thresh=seal_det_thresh,
  213. seal_det_unclip_ratio=seal_det_unclip_ratio,
  214. seal_rec_score_thresh=seal_rec_score_thresh,
  215. ):
  216. visual_info = self.decode_visual_result(layout_parsing_result)
  217. visual_predict_res = {
  218. "layout_parsing_result": layout_parsing_result,
  219. "visual_info": visual_info,
  220. }
  221. yield visual_predict_res
  222. def save_visual_info_list(self, visual_info: dict, save_path: str) -> None:
  223. """
  224. Save the visual info list to the specified file path.
  225. Args:
  226. visual_info (dict): The visual info result, which can be a single object or a list of objects.
  227. save_path (str): The file path to save the visual info list.
  228. Returns:
  229. None
  230. """
  231. if not isinstance(visual_info, list):
  232. visual_info_list = [visual_info]
  233. else:
  234. visual_info_list = visual_info
  235. with open(save_path, "w") as fout:
  236. fout.write(json.dumps(visual_info_list, ensure_ascii=False) + "\n")
  237. return
  238. def load_visual_info_list(self, data_path: str) -> list[dict]:
  239. """
  240. Loads visual info list from a JSON file.
  241. Args:
  242. data_path (str): The path to the JSON file containing visual info.
  243. Returns:
  244. list[dict]: A list of dict objects parsed from the JSON file.
  245. """
  246. with open(data_path, "r") as fin:
  247. data = fin.readline()
  248. visual_info_list = json.loads(data)
  249. return visual_info_list
  250. def merge_visual_info_list(
  251. self, visual_info_list: list[dict]
  252. ) -> tuple[list, list, list, list]:
  253. """
  254. Merge visual info lists.
  255. Args:
  256. visual_info_list (list[dict]): A list of visual info results.
  257. Returns:
  258. tuple[list, list, list, list]: A tuple containing four lists, one for normal text dicts,
  259. one for table text lists, one for table HTML lists.
  260. one for table neighbor texts.
  261. """
  262. all_normal_text_list = []
  263. all_table_text_list = []
  264. all_table_html_list = []
  265. all_table_nei_text_list = []
  266. for single_visual_info in visual_info_list:
  267. normal_text_dict = single_visual_info["normal_text_dict"]
  268. for key in normal_text_dict:
  269. normal_text_dict[key] = normal_text_dict[key].replace("\n", "")
  270. table_text_list = single_visual_info["table_text_list"]
  271. table_html_list = single_visual_info["table_html_list"]
  272. table_nei_text_list = single_visual_info["table_nei_text_list"]
  273. all_normal_text_list.append(normal_text_dict)
  274. all_table_text_list.extend(table_text_list)
  275. all_table_html_list.extend(table_html_list)
  276. all_table_nei_text_list.extend(table_nei_text_list)
  277. return (
  278. all_normal_text_list,
  279. all_table_text_list,
  280. all_table_html_list,
  281. all_table_nei_text_list,
  282. )
  283. def build_vector(
  284. self,
  285. visual_info: dict,
  286. min_characters: int = 3500,
  287. llm_request_interval: float = 1.0,
  288. flag_save_bytes_vector: bool = False,
  289. ) -> dict:
  290. """
  291. Build a vector representation from visual information.
  292. Args:
  293. visual_info (dict): The visual information input, can be a single instance or a list of instances.
  294. min_characters (int): The minimum number of characters required for text processing, defaults to 3500.
  295. llm_request_interval (float): The interval between LLM requests, defaults to 1.0.
  296. flag_save_bytes_vector (bool): Whether to save the vector as bytes, defaults to False.
  297. Returns:
  298. dict: A dictionary containing the vector info and a flag indicating if the text is too short.
  299. """
  300. if not isinstance(visual_info, list):
  301. visual_info_list = [visual_info]
  302. else:
  303. visual_info_list = visual_info
  304. all_visual_info = self.merge_visual_info_list(visual_info_list)
  305. (
  306. all_normal_text_list,
  307. all_table_text_list,
  308. all_table_html_list,
  309. all_table_nei_text_list,
  310. ) = all_visual_info
  311. vector_info = {}
  312. all_items = []
  313. for i, normal_text_dict in enumerate(all_normal_text_list):
  314. for type, text in normal_text_dict.items():
  315. all_items += [f"{type}:{text}\n"]
  316. for table_html, table_text, table_nei_text in zip(
  317. all_table_html_list, all_table_text_list, all_table_nei_text_list
  318. ):
  319. if len(table_html) > min_characters - self.table_structure_len_max:
  320. all_items += [f"table:{table_text}\t{table_nei_text}"]
  321. all_text_str = "".join(all_items)
  322. vector_info["flag_save_bytes_vector"] = False
  323. if len(all_text_str) > min_characters:
  324. vector_info["flag_too_short_text"] = False
  325. vector_info["vector"] = self.retriever.generate_vector_database(all_items)
  326. if flag_save_bytes_vector:
  327. vector_info["vector"] = self.retriever.encode_vector_store_to_bytes(
  328. vector_info["vector"]
  329. )
  330. vector_info["flag_save_bytes_vector"] = True
  331. else:
  332. vector_info["flag_too_short_text"] = True
  333. vector_info["vector"] = all_items
  334. return vector_info
  335. def save_vector(self, vector_info: dict, save_path: str) -> None:
  336. with open(save_path, "w") as fout:
  337. fout.write(json.dumps(vector_info, ensure_ascii=False) + "\n")
  338. return
  339. def load_vector(self, data_path: str) -> dict:
  340. vector_info = None
  341. with open(data_path, "r") as fin:
  342. data = fin.readline()
  343. vector_info = json.loads(data)
  344. if (
  345. "flag_too_short_text" not in vector_info
  346. or "flag_save_bytes_vector" not in vector_info
  347. or "vector" not in vector_info
  348. ):
  349. logging.error("Invalid vector info.")
  350. return {"error": "Invalid vector info when load vector!"}
  351. if vector_info["flag_save_bytes_vector"]:
  352. vector_info["vector"] = self.retriever.decode_vector_store_from_bytes(
  353. vector_info["vector"]
  354. )
  355. return vector_info
  356. def format_key(self, key_list: str | list[str]) -> list[str]:
  357. """
  358. Formats the key list.
  359. Args:
  360. key_list (str|list[str]): A string or a list of strings representing the keys.
  361. Returns:
  362. list[str]: A list of formatted keys.
  363. """
  364. if key_list == "":
  365. return []
  366. if isinstance(key_list, list):
  367. key_list = [key.replace("\xa0", " ") for key in key_list]
  368. return key_list
  369. if isinstance(key_list, str):
  370. key_list = re.sub(r"[\t\n\r\f\v]", "", key_list)
  371. key_list = key_list.replace(",", ",").split(",")
  372. return key_list
  373. return []
  374. def mllm_pred(
  375. self,
  376. input: str | np.ndarray,
  377. key_list,
  378. **kwargs,
  379. ) -> dict:
  380. key_list = self.format_key(key_list)
  381. if len(key_list) == 0:
  382. return {"mllm_res": "Error:输入的key_list无效!"}
  383. if isinstance(input, list):
  384. logging.error("Input is a list, but it's not supported here.")
  385. return {"mllm_res": "Error:Input is a list, but it's not supported here!"}
  386. image_array_list = self.img_reader([input])
  387. if (
  388. isinstance(input, str)
  389. and input.endswith(".pdf")
  390. and len(image_array_list) > 1
  391. ):
  392. logging.error("The input with PDF should have only one page.")
  393. return {"mllm_res": "Error:The input with PDF should have only one page!"}
  394. for image_array in image_array_list:
  395. assert len(image_array.shape) == 3
  396. image_string = cv2.imencode(".jpg", image_array)[1].tostring()
  397. image_base64 = base64.b64encode(image_string).decode("utf-8")
  398. result = {}
  399. for key in key_list:
  400. prompt = (
  401. str(key)
  402. + "\n请用图片中完整出现的内容回答,可以是单词、短语或句子,针对问题回答尽可能详细和完整,并保持格式、单位、符号和标点都与图片中的文字内容完全一致。"
  403. )
  404. mllm_chat_bot_result = self.mllm_chat_bot.generate_chat_results(
  405. prompt=prompt, image=image_base64
  406. )
  407. if mllm_chat_bot_result is None:
  408. return {"mllm_res": "大模型调用失败"}
  409. result[key] = mllm_chat_bot_result
  410. return {"mllm_res": result}
  411. def generate_and_merge_chat_results(
  412. self, prompt: str, key_list: list, final_results: dict, failed_results: list
  413. ) -> None:
  414. """
  415. Generate and merge chat results into the final results dictionary.
  416. Args:
  417. prompt (str): The input prompt for the chat bot.
  418. key_list (list): A list of keys to track which results to merge.
  419. final_results (dict): The dictionary to store the final merged results.
  420. failed_results (list): A list of failed results to avoid merging.
  421. Returns:
  422. None
  423. """
  424. llm_result = self.chat_bot.generate_chat_results(prompt)
  425. if llm_result is None:
  426. logging.error(
  427. "chat bot error: \n [prompt:]\n %s\n [result:] %s\n"
  428. % (prompt, self.chat_bot.ERROR_MASSAGE)
  429. )
  430. return
  431. llm_result = self.chat_bot.fix_llm_result_format(llm_result)
  432. for key, value in llm_result.items():
  433. if value not in failed_results and key in key_list:
  434. key_list.remove(key)
  435. final_results[key] = value
  436. return
  437. def get_related_normal_text(
  438. self,
  439. use_vector_retrieval: bool,
  440. vector_info: dict,
  441. key_list: list[str],
  442. all_normal_text_list: list,
  443. min_characters: int,
  444. ) -> str:
  445. """
  446. Retrieve related normal text based on vector retrieval or all normal text list.
  447. Args:
  448. use_vector_retrieval (bool): Whether to use vector retrieval.
  449. vector_info (dict): Dictionary containing vector information.
  450. key_list (list[str]): List of keys to generate question keys.
  451. all_normal_text_list (list): List of normal text.
  452. min_characters (int): The minimum number of characters required for text processing, defaults to 3500.
  453. Returns:
  454. str: Related normal text.
  455. """
  456. if use_vector_retrieval and vector_info is not None:
  457. question_key_list = [f"{key}" for key in key_list]
  458. vector = vector_info["vector"]
  459. if not vector_info["flag_too_short_text"]:
  460. related_text = self.retriever.similarity_retrieval(
  461. question_key_list, vector, topk=50, min_characters=min_characters
  462. )
  463. else:
  464. if len(vector) > 0:
  465. related_text = "".join(vector)
  466. else:
  467. related_text = ""
  468. else:
  469. all_items = []
  470. for i, normal_text_dict in enumerate(all_normal_text_list):
  471. for type, text in normal_text_dict.items():
  472. all_items += [f"{type}:{text}\n"]
  473. related_text = "".join(all_items)
  474. if len(related_text) > min_characters:
  475. logging.warning(
  476. "The input text content is too long, the large language model may truncate it."
  477. )
  478. return related_text
  479. def ensemble_ocr_llm_mllm(
  480. self, key_list: list[str], ocr_llm_predict_dict: dict, mllm_predict_dict: dict
  481. ) -> dict:
  482. """
  483. Ensemble OCR_LLM and LMM predictions based on given key list.
  484. Args:
  485. key_list (list[str]): List of keys to retrieve predictions.
  486. ocr_llm_predict_dict (dict): Dictionary containing OCR LLM predictions.
  487. mllm_predict_dict (dict): Dictionary containing mLLM predictions.
  488. Returns:
  489. dict: A dictionary with final predictions.
  490. """
  491. final_predict_dict = {}
  492. for key in key_list:
  493. predict = ""
  494. ocr_llm_predict = ""
  495. mllm_predict = ""
  496. if key in ocr_llm_predict_dict:
  497. ocr_llm_predict = ocr_llm_predict_dict[key]
  498. if key in mllm_predict_dict:
  499. mllm_predict = mllm_predict_dict[key]
  500. if ocr_llm_predict != "" and mllm_predict != "":
  501. prompt = self.ensemble_pe.generate_prompt(
  502. key, ocr_llm_predict, mllm_predict
  503. )
  504. llm_result = self.chat_bot.generate_chat_results(prompt)
  505. if llm_result is not None:
  506. llm_result = self.chat_bot.fix_llm_result_format(llm_result)
  507. if key in llm_result:
  508. tmp = llm_result[key]
  509. if "B" in tmp:
  510. predict = mllm_predict
  511. else:
  512. predict = ocr_llm_predict
  513. else:
  514. predict = ocr_llm_predict
  515. elif key in ocr_llm_predict_dict:
  516. predict = ocr_llm_predict_dict[key]
  517. elif key in mllm_predict_dict:
  518. predict = mllm_predict_dict[key]
  519. if predict != "":
  520. final_predict_dict[key] = predict
  521. return final_predict_dict
  522. def chat(
  523. self,
  524. key_list: str | list[str],
  525. visual_info: dict,
  526. use_vector_retrieval: bool = True,
  527. vector_info: dict = None,
  528. min_characters: int = 3500,
  529. text_task_description: str = None,
  530. text_output_format: str = None,
  531. text_rules_str: str = None,
  532. text_few_shot_demo_text_content: str = None,
  533. text_few_shot_demo_key_value_list: str = None,
  534. table_task_description: str = None,
  535. table_output_format: str = None,
  536. table_rules_str: str = None,
  537. table_few_shot_demo_text_content: str = None,
  538. table_few_shot_demo_key_value_list: str = None,
  539. mllm_predict_dict: dict = None,
  540. mllm_integration_strategy: str = "integration",
  541. ) -> dict:
  542. """
  543. Generates chat results based on the provided key list and visual information.
  544. Args:
  545. key_list (str | list[str]): A single key or a list of keys to extract information.
  546. visual_info (dict): The visual information result.
  547. use_vector_retrieval (bool): Whether to use vector retrieval.
  548. vector_info (dict): The vector information for retrieval.
  549. min_characters (int): The minimum number of characters required for text processing, defaults to 3500.
  550. text_task_description (str): The description of the text task.
  551. text_output_format (str): The output format for text results.
  552. text_rules_str (str): The rules for generating text results.
  553. text_few_shot_demo_text_content (str): The text content for few-shot demos.
  554. text_few_shot_demo_key_value_list (str): The key-value list for few-shot demos.
  555. table_task_description (str): The description of the table task.
  556. table_output_format (str): The output format for table results.
  557. table_rules_str (str): The rules for generating table results.
  558. table_few_shot_demo_text_content (str): The text content for table few-shot demos.
  559. table_few_shot_demo_key_value_list (str): The key-value list for table few-shot demos.
  560. mllm_predict_dict (dict): The dictionary of mLLM predicts.
  561. mllm_integration_strategy(str): The integration strategy of mLLM and LLM, defaults to "integration", options are "integration", "llm_only" and "mllm_only".
  562. Returns:
  563. dict: A dictionary containing the chat results.
  564. """
  565. key_list = self.format_key(key_list)
  566. key_list_ori = key_list.copy()
  567. if len(key_list) == 0:
  568. return {"chat_res": "Error:输入的key_list无效!"}
  569. if not isinstance(visual_info, list):
  570. visual_info_list = [visual_info]
  571. else:
  572. visual_info_list = visual_info
  573. all_visual_info = self.merge_visual_info_list(visual_info_list)
  574. (
  575. all_normal_text_list,
  576. all_table_text_list,
  577. all_table_html_list,
  578. all_table_nei_text_list,
  579. ) = all_visual_info
  580. final_results = {}
  581. failed_results = ["大模型调用失败", "未知", "未找到关键信息", "None", ""]
  582. if len(key_list) > 0:
  583. related_text = self.get_related_normal_text(
  584. use_vector_retrieval,
  585. vector_info,
  586. key_list,
  587. all_normal_text_list,
  588. min_characters,
  589. )
  590. if len(related_text) > 0:
  591. prompt = self.text_pe.generate_prompt(
  592. related_text,
  593. key_list,
  594. task_description=text_task_description,
  595. output_format=text_output_format,
  596. rules_str=text_rules_str,
  597. few_shot_demo_text_content=text_few_shot_demo_text_content,
  598. few_shot_demo_key_value_list=text_few_shot_demo_key_value_list,
  599. )
  600. self.generate_and_merge_chat_results(
  601. prompt, key_list, final_results, failed_results
  602. )
  603. if len(key_list) > 0:
  604. for table_html, table_text, table_nei_text in zip(
  605. all_table_html_list, all_table_text_list, all_table_nei_text_list
  606. ):
  607. if len(table_html) <= min_characters - self.table_structure_len_max:
  608. for table_info in [table_html]:
  609. if len(key_list) > 0:
  610. if len(table_nei_text) > 0:
  611. table_info = (
  612. table_info + "\n 表格周围文字:" + table_nei_text
  613. )
  614. prompt = self.table_pe.generate_prompt(
  615. table_info,
  616. key_list,
  617. task_description=table_task_description,
  618. output_format=table_output_format,
  619. rules_str=table_rules_str,
  620. few_shot_demo_text_content=table_few_shot_demo_text_content,
  621. few_shot_demo_key_value_list=table_few_shot_demo_key_value_list,
  622. )
  623. self.generate_and_merge_chat_results(
  624. prompt, key_list, final_results, failed_results
  625. )
  626. if self.use_mllm_predict and mllm_predict_dict != "llm_only":
  627. if mllm_integration_strategy == "integration":
  628. final_predict_dict = self.ensemble_ocr_llm_mllm(
  629. key_list_ori, final_results, mllm_predict_dict
  630. )
  631. elif mllm_integration_strategy == "mllm_only":
  632. final_predict_dict = mllm_predict_dict
  633. else:
  634. return {
  635. "chat_res": f"Error:Unsupported mllm_integration_strategy {mllm_integration_strategy}, only support 'integration', 'llm_only' and 'mllm_only'!"
  636. }
  637. else:
  638. final_predict_dict = final_results
  639. return {"chat_res": final_predict_dict}
  640. def predict(self, *args, **kwargs) -> None:
  641. logging.error(
  642. "PP-ChatOCRv4-doc Pipeline do not support to call `predict()` directly! Please invoke `visual_predict`, `build_vector`, `chat` sequentially to obtain the result."
  643. )
  644. return