pipeline_v3.py 32 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723724725726727728729730731732733734735736737738739740741742743744745746747748749750751752753754755756757758759760761762763764765766767768769770771772773774775776777778779780781782783
  1. # Copyright (c) 2024 PaddlePaddle Authors. All Rights Reserved.
  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. import copy
  15. import json
  16. import os
  17. import re
  18. from typing import Any, Dict, List, Optional, Tuple, Union
  19. import numpy as np
  20. from ....utils import logging
  21. from ....utils.deps import pipeline_requires_extra
  22. from ....utils.file_interface import custom_open
  23. from ...common.batch_sampler import ImageBatchSampler
  24. from ...common.reader import ReadImage
  25. from ...utils.benchmark import benchmark
  26. from ...utils.hpi import HPIConfig
  27. from ...utils.pp_option import PaddlePredictorOption
  28. from ..components.chat_server import BaseChat
  29. from ..layout_parsing.result import LayoutParsingResult
  30. from .pipeline_base import PP_ChatOCR_Pipeline
  31. @benchmark.time_methods
  32. @pipeline_requires_extra("ie")
  33. class PP_ChatOCRv3_Pipeline(PP_ChatOCR_Pipeline):
  34. """PP-ChatOCR Pipeline"""
  35. entities = ["PP-ChatOCRv3-doc"]
  36. def __init__(
  37. self,
  38. config: Dict,
  39. device: str = None,
  40. pp_option: PaddlePredictorOption = None,
  41. use_hpip: bool = False,
  42. hpi_config: Optional[Union[Dict[str, Any], HPIConfig]] = None,
  43. initial_predictor: bool = True,
  44. ) -> None:
  45. """Initializes the pp-chatocrv3-doc pipeline.
  46. Args:
  47. config (Dict): Configuration dictionary containing various settings.
  48. device (str, optional): Device to run the predictions on. Defaults to None.
  49. pp_option (PaddlePredictorOption, optional): PaddlePredictor options. Defaults to None.
  50. use_hpip (bool, optional): Whether to use the high-performance
  51. inference plugin (HPIP) by default. Defaults to False.
  52. hpi_config (Optional[Union[Dict[str, Any], HPIConfig]], optional):
  53. The default high-performance inference configuration dictionary.
  54. Defaults to None.
  55. initial_predictor (bool, optional): Whether to initialize the predictor. Defaults to True.
  56. """
  57. super().__init__(
  58. device=device, pp_option=pp_option, use_hpip=use_hpip, hpi_config=hpi_config
  59. )
  60. self.pipeline_name = config["pipeline_name"]
  61. self.config = config
  62. self.use_layout_parser = config.get("use_layout_parser", True)
  63. self.layout_parsing_pipeline = None
  64. self.chat_bot = None
  65. self.retriever = None
  66. if initial_predictor:
  67. self.inintial_visual_predictor(config)
  68. self.inintial_chat_predictor(config)
  69. self.inintial_retriever_predictor(config)
  70. self.batch_sampler = ImageBatchSampler(batch_size=1)
  71. self.img_reader = ReadImage(format="BGR")
  72. self.table_structure_len_max = 500
  73. def inintial_visual_predictor(self, config: dict) -> None:
  74. """
  75. Initializes the visual predictor with the given configuration.
  76. Args:
  77. config (dict): The configuration dictionary containing the necessary
  78. parameters for initializing the predictor.
  79. Returns:
  80. None
  81. """
  82. self.use_layout_parser = config.get("use_layout_parser", True)
  83. if self.use_layout_parser:
  84. layout_parsing_config = config.get("SubPipelines", {}).get(
  85. "LayoutParser",
  86. {"pipeline_config_error": "config error for layout_parsing_pipeline!"},
  87. )
  88. self.layout_parsing_pipeline = self.create_pipeline(layout_parsing_config)
  89. return
  90. def inintial_retriever_predictor(self, config: dict) -> None:
  91. """
  92. Initializes the retriever predictor with the given configuration.
  93. Args:
  94. config (dict): The configuration dictionary containing the necessary
  95. parameters for initializing the predictor.
  96. Returns:
  97. None
  98. """
  99. from .. import create_retriever
  100. retriever_config = config.get("SubModules", {}).get(
  101. "LLM_Retriever",
  102. {"retriever_config_error": "config error for llm retriever!"},
  103. )
  104. self.retriever = create_retriever(retriever_config)
  105. def inintial_chat_predictor(self, config: dict) -> None:
  106. """
  107. Initializes the chat predictor with the given configuration.
  108. Args:
  109. config (dict): The configuration dictionary containing the necessary
  110. parameters for initializing the predictor.
  111. Returns:
  112. None
  113. """
  114. from .. import create_chat_bot
  115. chat_bot_config = config.get("SubModules", {}).get(
  116. "LLM_Chat",
  117. {"chat_bot_config_error": "config error for llm chat bot!"},
  118. )
  119. self.chat_bot = create_chat_bot(chat_bot_config)
  120. from .. import create_prompt_engineering
  121. text_pe_config = (
  122. config.get("SubModules", {})
  123. .get("PromptEngneering", {})
  124. .get(
  125. "KIE_CommonText",
  126. {"pe_config_error": "config error for text_pe!"},
  127. )
  128. )
  129. self.text_pe = create_prompt_engineering(text_pe_config)
  130. table_pe_config = (
  131. config.get("SubModules", {})
  132. .get("PromptEngneering", {})
  133. .get(
  134. "KIE_Table",
  135. {"pe_config_error": "config error for table_pe!"},
  136. )
  137. )
  138. self.table_pe = create_prompt_engineering(table_pe_config)
  139. return
  140. def decode_visual_result(self, layout_parsing_result: LayoutParsingResult) -> dict:
  141. """
  142. Decodes the visual result from the layout parsing result.
  143. Args:
  144. layout_parsing_result (LayoutParsingResult): The result of layout parsing.
  145. Returns:
  146. dict: The decoded visual information.
  147. """
  148. normal_text_dict = {}
  149. parsing_res_list = layout_parsing_result["parsing_res_list"]
  150. for pno in range(len(parsing_res_list)):
  151. label = parsing_res_list[pno]["block_label"]
  152. content = parsing_res_list[pno]["block_content"]
  153. if label in ["table", "formula"]:
  154. continue
  155. key = f"words in {label}"
  156. if key not in normal_text_dict:
  157. normal_text_dict[key] = content
  158. else:
  159. normal_text_dict[key] += f"\n {content}"
  160. table_res_list = layout_parsing_result["table_res_list"]
  161. table_text_list = []
  162. table_html_list = []
  163. for table_res in table_res_list:
  164. table_html_list.append(table_res["pred_html"])
  165. single_table_text = " ".join(table_res["table_ocr_pred"]["rec_texts"])
  166. table_text_list.append(single_table_text)
  167. visual_info = {}
  168. visual_info["normal_text_dict"] = normal_text_dict
  169. visual_info["table_text_list"] = table_text_list
  170. visual_info["table_html_list"] = table_html_list
  171. return visual_info
  172. # Function to perform visual prediction on input images
  173. def visual_predict(
  174. self,
  175. input: Union[str, List[str], np.ndarray, List[np.ndarray]],
  176. use_doc_orientation_classify: Optional[bool] = None,
  177. use_doc_unwarping: Optional[bool] = None,
  178. use_seal_recognition: Optional[bool] = None,
  179. use_table_recognition: Optional[bool] = None,
  180. layout_threshold: Optional[Union[float, dict]] = None,
  181. layout_nms: Optional[bool] = None,
  182. layout_unclip_ratio: Optional[Union[float, Tuple[float, float], dict]] = None,
  183. layout_merge_bboxes_mode: Optional[str] = None,
  184. text_det_limit_side_len: Optional[int] = None,
  185. text_det_limit_type: Optional[str] = None,
  186. text_det_thresh: Optional[float] = None,
  187. text_det_box_thresh: Optional[float] = None,
  188. text_det_unclip_ratio: Optional[float] = None,
  189. text_rec_score_thresh: Optional[float] = None,
  190. seal_det_limit_side_len: Optional[int] = None,
  191. seal_det_limit_type: Optional[str] = None,
  192. seal_det_thresh: Optional[float] = None,
  193. seal_det_box_thresh: Optional[float] = None,
  194. seal_det_unclip_ratio: Optional[float] = None,
  195. seal_rec_score_thresh: Optional[float] = None,
  196. **kwargs,
  197. ) -> dict:
  198. """
  199. This function takes an input image or a list of images and performs various visual
  200. prediction tasks such as document orientation classification, document unwarping,
  201. general OCR, seal recognition, and table recognition based on the provided flags.
  202. Args:
  203. input (Union[str, list[str], np.ndarray, list[np.ndarray]]): Input image path, list of image paths,
  204. numpy array of an image, or list of numpy arrays.
  205. use_doc_orientation_classify (bool): Flag to use document orientation classification.
  206. use_doc_unwarping (bool): Flag to use document unwarping.
  207. use_seal_recognition (bool): Flag to use seal recognition.
  208. use_table_recognition (bool): Flag to use table recognition.
  209. layout_threshold (Optional[float]): The threshold value to filter out low-confidence predictions. Default is None.
  210. layout_nms (bool, optional): Whether to use layout-aware NMS. Defaults to False.
  211. layout_unclip_ratio (Optional[Union[float, Tuple[float, float]]], optional): The ratio of unclipping the bounding box.
  212. Defaults to None.
  213. If it's a single number, then both width and height are used.
  214. If it's a tuple of two numbers, then they are used separately for width and height respectively.
  215. If it's None, then no unclipping will be performed.
  216. layout_merge_bboxes_mode (Optional[str], optional): The mode for merging bounding boxes. Defaults to None.
  217. text_det_limit_side_len (Optional[int]): Maximum side length for text detection.
  218. text_det_limit_type (Optional[str]): Type of limit to apply for text detection.
  219. text_det_thresh (Optional[float]): Threshold for text detection.
  220. text_det_box_thresh (Optional[float]): Threshold for text detection boxes.
  221. text_det_unclip_ratio (Optional[float]): Ratio for unclipping text detection boxes.
  222. text_rec_score_thresh (Optional[float]): Score threshold for text recognition.
  223. seal_det_limit_side_len (Optional[int]): Maximum side length for seal detection.
  224. seal_det_limit_type (Optional[str]): Type of limit to apply for seal detection.
  225. seal_det_thresh (Optional[float]): Threshold for seal detection.
  226. seal_det_box_thresh (Optional[float]): Threshold for seal detection boxes.
  227. seal_det_unclip_ratio (Optional[float]): Ratio for unclipping seal detection boxes.
  228. seal_rec_score_thresh (Optional[float]): Score threshold for seal recognition.
  229. **kwargs: Additional keyword arguments.
  230. Returns:
  231. dict: A dictionary containing the layout parsing result and visual information.
  232. """
  233. if self.use_layout_parser == False:
  234. logging.error("The models for layout parser are not initialized.")
  235. yield {"error": "The models for layout parser are not initialized."}
  236. if self.layout_parsing_pipeline is None:
  237. logging.warning(
  238. "The layout parsing pipeline is not initialized, will initialize it now."
  239. )
  240. self.inintial_visual_predictor(self.config)
  241. for layout_parsing_result in self.layout_parsing_pipeline.predict(
  242. input,
  243. use_doc_orientation_classify=use_doc_orientation_classify,
  244. use_doc_unwarping=use_doc_unwarping,
  245. use_seal_recognition=use_seal_recognition,
  246. use_table_recognition=use_table_recognition,
  247. layout_threshold=layout_threshold,
  248. layout_nms=layout_nms,
  249. layout_unclip_ratio=layout_unclip_ratio,
  250. layout_merge_bboxes_mode=layout_merge_bboxes_mode,
  251. text_det_limit_side_len=text_det_limit_side_len,
  252. text_det_limit_type=text_det_limit_type,
  253. text_det_thresh=text_det_thresh,
  254. text_det_box_thresh=text_det_box_thresh,
  255. text_det_unclip_ratio=text_det_unclip_ratio,
  256. text_rec_score_thresh=text_rec_score_thresh,
  257. seal_det_box_thresh=seal_det_box_thresh,
  258. seal_det_limit_side_len=seal_det_limit_side_len,
  259. seal_det_limit_type=seal_det_limit_type,
  260. seal_det_thresh=seal_det_thresh,
  261. seal_det_unclip_ratio=seal_det_unclip_ratio,
  262. seal_rec_score_thresh=seal_rec_score_thresh,
  263. ):
  264. visual_info = self.decode_visual_result(layout_parsing_result)
  265. visual_predict_res = {
  266. "layout_parsing_result": layout_parsing_result,
  267. "visual_info": visual_info,
  268. }
  269. yield visual_predict_res
  270. def save_visual_info_list(self, visual_info: dict, save_path: str) -> None:
  271. """
  272. Save the visual info list to the specified file path.
  273. Args:
  274. visual_info (dict): The visual info result, which can be a single object or a list of objects.
  275. save_path (str): The file path to save the visual info list.
  276. Returns:
  277. None
  278. """
  279. if not isinstance(visual_info, list):
  280. visual_info_list = [visual_info]
  281. else:
  282. visual_info_list = visual_info
  283. directory = os.path.dirname(save_path)
  284. if not os.path.exists(directory):
  285. os.makedirs(directory)
  286. with custom_open(save_path, "w") as fout:
  287. fout.write(json.dumps(visual_info_list, ensure_ascii=False) + "\n")
  288. return
  289. def load_visual_info_list(self, data_path: str) -> List[dict]:
  290. """
  291. Loads visual info list from a JSON file.
  292. Args:
  293. data_path (str): The path to the JSON file containing visual info.
  294. Returns:
  295. list[dict]: A list of dict objects parsed from the JSON file.
  296. """
  297. with custom_open(data_path, "r") as fin:
  298. data = fin.readline()
  299. visual_info_list = json.loads(data)
  300. return visual_info_list
  301. def merge_visual_info_list(
  302. self, visual_info_list: List[dict]
  303. ) -> Tuple[list, list, list]:
  304. """
  305. Merge visual info lists.
  306. Args:
  307. visual_info_list (list[dict]): A list of visual info results.
  308. Returns:
  309. tuple[list, list, list]: A tuple containing four lists, one for normal text dicts,
  310. one for table text lists, one for table HTML lists.
  311. """
  312. all_normal_text_list = []
  313. all_table_text_list = []
  314. all_table_html_list = []
  315. for single_visual_info in visual_info_list:
  316. normal_text_dict = single_visual_info["normal_text_dict"]
  317. for key in normal_text_dict:
  318. normal_text_dict[key] = normal_text_dict[key].replace("\n", "")
  319. table_text_list = single_visual_info["table_text_list"]
  320. table_html_list = single_visual_info["table_html_list"]
  321. all_normal_text_list.append(normal_text_dict)
  322. all_table_text_list.extend(table_text_list)
  323. all_table_html_list.extend(table_html_list)
  324. return (all_normal_text_list, all_table_text_list, all_table_html_list)
  325. def build_vector(
  326. self,
  327. visual_info: dict,
  328. min_characters: int = 3500,
  329. block_size: int = 300,
  330. flag_save_bytes_vector: bool = False,
  331. retriever_config: dict = None,
  332. ) -> dict:
  333. """
  334. Build a vector representation from visual information.
  335. Args:
  336. visual_info (dict): The visual information input, can be a single instance or a list of instances.
  337. min_characters (int): The minimum number of characters required for text processing, defaults to 3500.
  338. block_size (int): The size of each chunk to split the text into.
  339. flag_save_bytes_vector (bool): Whether to save the vector as bytes, defaults to False.
  340. retriever_config (dict): The configuration for the retriever, defaults to None.
  341. Returns:
  342. dict: A dictionary containing the vector info and a flag indicating if the text is too short.
  343. """
  344. if not isinstance(visual_info, list):
  345. visual_info_list = [visual_info]
  346. else:
  347. visual_info_list = visual_info
  348. if retriever_config is not None:
  349. from .. import create_retriever
  350. retriever = create_retriever(retriever_config)
  351. else:
  352. if self.retriever is None:
  353. logging.warning(
  354. "The retriever is not initialized,will initialize it now."
  355. )
  356. self.inintial_retriever_predictor(self.config)
  357. retriever = self.retriever
  358. all_visual_info = self.merge_visual_info_list(visual_info_list)
  359. (
  360. all_normal_text_list,
  361. all_table_text_list,
  362. all_table_html_list,
  363. ) = all_visual_info
  364. vector_info = {}
  365. all_items = []
  366. for i, normal_text_dict in enumerate(all_normal_text_list):
  367. for type, text in normal_text_dict.items():
  368. all_items += [f"{type}:{text}\n"]
  369. for table_html, table_text in zip(all_table_html_list, all_table_text_list):
  370. if len(table_html) > min_characters - self.table_structure_len_max:
  371. all_items += [f"table:{table_text}"]
  372. all_text_str = "".join(all_items)
  373. vector_info["flag_save_bytes_vector"] = False
  374. if len(all_text_str) > min_characters:
  375. vector_info["flag_too_short_text"] = False
  376. vector_info["model_name"] = retriever.model_name
  377. vector_info["block_size"] = block_size
  378. vector_info["vector"] = retriever.generate_vector_database(
  379. all_items, block_size=block_size
  380. )
  381. if flag_save_bytes_vector:
  382. vector_info["vector"] = retriever.encode_vector_store_to_bytes(
  383. vector_info["vector"]
  384. )
  385. vector_info["flag_save_bytes_vector"] = True
  386. else:
  387. vector_info["flag_too_short_text"] = True
  388. vector_info["vector"] = all_items
  389. return vector_info
  390. def save_vector(
  391. self, vector_info: dict, save_path: str, retriever_config: dict = None
  392. ) -> None:
  393. directory = os.path.dirname(save_path)
  394. if not os.path.exists(directory):
  395. os.makedirs(directory)
  396. if retriever_config is not None:
  397. from .. import create_retriever
  398. retriever = create_retriever(retriever_config)
  399. else:
  400. if self.retriever is None:
  401. logging.warning(
  402. "The retriever is not initialized,will initialize it now."
  403. )
  404. self.inintial_retriever_predictor(self.config)
  405. retriever = self.retriever
  406. vector_info_data = copy.deepcopy(vector_info)
  407. if (
  408. not vector_info["flag_too_short_text"]
  409. and not vector_info["flag_save_bytes_vector"]
  410. ):
  411. vector_info_data["vector"] = retriever.encode_vector_store_to_bytes(
  412. vector_info_data["vector"]
  413. )
  414. vector_info_data["flag_save_bytes_vector"] = True
  415. with custom_open(save_path, "w") as fout:
  416. fout.write(json.dumps(vector_info_data, ensure_ascii=False) + "\n")
  417. return
  418. def load_vector(self, data_path: str, retriever_config: dict = None) -> dict:
  419. vector_info = None
  420. if retriever_config is not None:
  421. from .. import create_retriever
  422. retriever = create_retriever(retriever_config)
  423. else:
  424. if self.retriever is None:
  425. logging.warning(
  426. "The retriever is not initialized,will initialize it now."
  427. )
  428. self.inintial_retriever_predictor(self.config)
  429. retriever = self.retriever
  430. with open(data_path, "r") as fin:
  431. data = fin.readline()
  432. vector_info = json.loads(data)
  433. if (
  434. "flag_too_short_text" not in vector_info
  435. or "flag_save_bytes_vector" not in vector_info
  436. or "vector" not in vector_info
  437. ):
  438. logging.error("Invalid vector info.")
  439. return {"error": "Invalid vector info when load vector!"}
  440. if vector_info["flag_save_bytes_vector"]:
  441. vector_info["vector"] = retriever.decode_vector_store_from_bytes(
  442. vector_info["vector"]
  443. )
  444. vector_info["flag_save_bytes_vector"] = False
  445. return vector_info
  446. def format_key(self, key_list: Union[str, List[str]]) -> List[str]:
  447. """
  448. Formats the key list.
  449. Args:
  450. key_list (str|list[str]): A string or a list of strings representing the keys.
  451. Returns:
  452. list[str]: A list of formatted keys.
  453. """
  454. if key_list == "":
  455. return []
  456. if isinstance(key_list, list):
  457. return key_list
  458. if isinstance(key_list, str):
  459. key_list = re.sub(r"[\t\n\r\f\v]", "", key_list)
  460. key_list = key_list.replace(",", ",").split(",")
  461. return key_list
  462. return []
  463. def generate_and_merge_chat_results(
  464. self,
  465. chat_bot: BaseChat,
  466. prompt: str,
  467. key_list: list,
  468. final_results: dict,
  469. failed_results: list,
  470. ) -> None:
  471. """
  472. Generate and merge chat results into the final results dictionary.
  473. Args:
  474. prompt (str): The input prompt for the chat bot.
  475. key_list (list): A list of keys to track which results to merge.
  476. final_results (dict): The dictionary to store the final merged results.
  477. failed_results (list): A list of failed results to avoid merging.
  478. Returns:
  479. None
  480. """
  481. llm_result = chat_bot.generate_chat_results(prompt)
  482. llm_result_content = llm_result["content"]
  483. llm_result_reasoning_content = llm_result["reasoning_content"]
  484. if llm_result_reasoning_content is not None:
  485. if "reasoning_content" not in final_results:
  486. final_results["reasoning_content"] = [llm_result_reasoning_content]
  487. else:
  488. final_results["reasoning_content"].append(llm_result_reasoning_content)
  489. if llm_result_content is None:
  490. logging.error(
  491. "chat bot error: \n [prompt:]\n %s\n [result:] %s\n"
  492. % (prompt, chat_bot.ERROR_MASSAGE)
  493. )
  494. return
  495. llm_result_content = chat_bot.fix_llm_result_format(llm_result_content)
  496. for key, value in llm_result_content.items():
  497. if value not in failed_results and key in key_list:
  498. key_list.remove(key)
  499. final_results[key] = value
  500. return
  501. def get_related_normal_text(
  502. self,
  503. retriever_config: dict,
  504. use_vector_retrieval: bool,
  505. vector_info: dict,
  506. key_list: List[str],
  507. all_normal_text_list: list,
  508. min_characters: int,
  509. ) -> str:
  510. """
  511. Retrieve related normal text based on vector retrieval or all normal text list.
  512. Args:
  513. retriever_config (dict): Configuration for the retriever.
  514. use_vector_retrieval (bool): Whether to use vector retrieval.
  515. vector_info (dict): Dictionary containing vector information.
  516. key_list (list[str]): List of keys to generate question keys.
  517. all_normal_text_list (list): List of normal text.
  518. min_characters (int): The minimum number of characters required for text processing, defaults to 3500.
  519. Returns:
  520. str: Related normal text.
  521. """
  522. if use_vector_retrieval and vector_info is not None:
  523. if retriever_config is not None:
  524. from .. import create_retriever
  525. retriever = create_retriever(retriever_config)
  526. else:
  527. if self.retriever is None:
  528. logging.warning(
  529. "The retriever is not initialized,will initialize it now."
  530. )
  531. self.inintial_retriever_predictor(self.config)
  532. retriever = self.retriever
  533. question_key_list = [f"{key}" for key in key_list]
  534. vector = vector_info["vector"]
  535. if not vector_info["flag_too_short_text"]:
  536. assert (
  537. vector_info["model_name"] == retriever.model_name
  538. ), f"The vector model name ({vector_info['model_name']}) does not match the retriever model name ({retriever.model_name}). Please check your retriever config."
  539. if vector_info["flag_save_bytes_vector"]:
  540. vector = retriever.decode_vector_store_from_bytes(vector)
  541. related_text = retriever.similarity_retrieval(
  542. question_key_list, vector, topk=50, min_characters=min_characters
  543. )
  544. else:
  545. if len(vector) > 0:
  546. related_text = "".join(vector)
  547. else:
  548. related_text = ""
  549. else:
  550. all_items = []
  551. for i, normal_text_dict in enumerate(all_normal_text_list):
  552. for type, text in normal_text_dict.items():
  553. all_items += [f"{type}:{text}\n"]
  554. related_text = "".join(all_items)
  555. if len(related_text) > min_characters:
  556. logging.warning(
  557. "The input text content is too long, the large language model may truncate it."
  558. )
  559. return related_text
  560. def chat(
  561. self,
  562. key_list: Union[str, List[str]],
  563. visual_info: List[dict],
  564. use_vector_retrieval: bool = True,
  565. vector_info: dict = None,
  566. min_characters: int = 3500,
  567. text_task_description: str = None,
  568. text_output_format: str = None,
  569. text_rules_str: str = None,
  570. text_few_shot_demo_text_content: str = None,
  571. text_few_shot_demo_key_value_list: str = None,
  572. table_task_description: str = None,
  573. table_output_format: str = None,
  574. table_rules_str: str = None,
  575. table_few_shot_demo_text_content: str = None,
  576. table_few_shot_demo_key_value_list: str = None,
  577. chat_bot_config: dict = None,
  578. retriever_config: dict = None,
  579. ) -> dict:
  580. """
  581. Generates chat results based on the provided key list and visual information.
  582. Args:
  583. key_list (Union[str, list[str]]): A single key or a list of keys to extract information.
  584. visual_info (dict): The visual information result.
  585. use_vector_retrieval (bool): Whether to use vector retrieval.
  586. vector_info (dict): The vector information for retrieval.
  587. min_characters (int): The minimum number of characters required.
  588. text_task_description (str): The description of the text task.
  589. text_output_format (str): The output format for text results.
  590. text_rules_str (str): The rules for generating text results.
  591. text_few_shot_demo_text_content (str): The text content for few-shot demos.
  592. text_few_shot_demo_key_value_list (str): The key-value list for few-shot demos.
  593. table_task_description (str): The description of the table task.
  594. table_output_format (str): The output format for table results.
  595. table_rules_str (str): The rules for generating table results.
  596. table_few_shot_demo_text_content (str): The text content for table few-shot demos.
  597. table_few_shot_demo_key_value_list (str): The key-value list for table few-shot demos.
  598. chat_bot_config(dict): The parameters for LLM chatbot, including api_type, api_key... refer to config file for more details.
  599. retriever_config (dict): The parameters for LLM retriever, including api_type, api_key... refer to config file for more details.
  600. Returns:
  601. dict: A dictionary containing the chat results.
  602. """
  603. key_list = self.format_key(key_list)
  604. if len(key_list) == 0:
  605. return {"chat_res": "Error:输入的key_list无效!"}
  606. if not isinstance(visual_info, list):
  607. visual_info_list = [visual_info]
  608. else:
  609. visual_info_list = visual_info
  610. if self.chat_bot is None:
  611. logging.warning(
  612. "The LLM chat bot is not initialized,will initialize it now."
  613. )
  614. self.inintial_chat_predictor(self.config)
  615. if chat_bot_config is not None:
  616. from .. import create_chat_bot
  617. chat_bot = create_chat_bot(chat_bot_config)
  618. else:
  619. chat_bot = self.chat_bot
  620. all_visual_info = self.merge_visual_info_list(visual_info_list)
  621. (
  622. all_normal_text_list,
  623. all_table_text_list,
  624. all_table_html_list,
  625. ) = all_visual_info
  626. final_results = {}
  627. failed_results = ["大模型调用失败", "未知", "未找到关键信息", "None", ""]
  628. if len(key_list) > 0:
  629. for table_html, table_text in zip(all_table_html_list, all_table_text_list):
  630. if len(table_html) <= min_characters - self.table_structure_len_max:
  631. for table_info in [table_html]:
  632. if len(key_list) > 0:
  633. prompt = self.table_pe.generate_prompt(
  634. table_info,
  635. key_list,
  636. task_description=table_task_description,
  637. output_format=table_output_format,
  638. rules_str=table_rules_str,
  639. few_shot_demo_text_content=table_few_shot_demo_text_content,
  640. few_shot_demo_key_value_list=table_few_shot_demo_key_value_list,
  641. )
  642. self.generate_and_merge_chat_results(
  643. chat_bot,
  644. prompt,
  645. key_list,
  646. final_results,
  647. failed_results,
  648. )
  649. if len(key_list) > 0:
  650. related_text = self.get_related_normal_text(
  651. retriever_config,
  652. use_vector_retrieval,
  653. vector_info,
  654. key_list,
  655. all_normal_text_list,
  656. min_characters,
  657. )
  658. if len(related_text) > 0:
  659. prompt = self.text_pe.generate_prompt(
  660. related_text,
  661. key_list,
  662. task_description=text_task_description,
  663. output_format=text_output_format,
  664. rules_str=text_rules_str,
  665. few_shot_demo_text_content=text_few_shot_demo_text_content,
  666. few_shot_demo_key_value_list=text_few_shot_demo_key_value_list,
  667. )
  668. self.generate_and_merge_chat_results(
  669. chat_bot, prompt, key_list, final_results, failed_results
  670. )
  671. return {"chat_res": final_results}
  672. def predict(self, *args, **kwargs) -> None:
  673. logging.error(
  674. "PP-ChatOCRv3-doc Pipeline do not support to call `predict()` directly! Please invoke `visual_predict`, `build_vector`, `chat` sequentially to obtain the result."
  675. )
  676. return