ppchatocrv3.py 26 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710
  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. import os
  15. import re
  16. import json
  17. import numpy as np
  18. from .utils import *
  19. from copy import deepcopy
  20. from ...components import *
  21. from ..ocr import OCRPipeline
  22. from ....utils import logging
  23. from ...results import *
  24. from ...components.llm import ErnieBot
  25. from ...utils.io import ImageReader, PDFReader
  26. from ..table_recognition import _TableRecPipeline
  27. from ...components.llm import create_llm_api, ErnieBot
  28. from ....utils.file_interface import read_yaml_file
  29. from ..table_recognition.utils import convert_4point2rect, get_ori_coordinate_for_table
  30. PROMPT_FILE = os.path.join(os.path.dirname(__file__), "ch_prompt.yaml")
  31. class PPChatOCRPipeline(_TableRecPipeline):
  32. """PP-ChatOCRv3 Pileline"""
  33. entities = "PP-ChatOCRv3-doc"
  34. def __init__(
  35. self,
  36. layout_model,
  37. text_det_model,
  38. text_rec_model,
  39. table_model,
  40. doc_image_ori_cls_model=None,
  41. doc_image_unwarp_model=None,
  42. seal_text_det_model=None,
  43. llm_name="ernie-3.5",
  44. llm_params={},
  45. task_prompt_yaml=None,
  46. user_prompt_yaml=None,
  47. layout_batch_size=1,
  48. text_det_batch_size=1,
  49. text_rec_batch_size=1,
  50. table_batch_size=1,
  51. doc_image_ori_cls_batch_size=1,
  52. doc_image_unwarp_batch_size=1,
  53. seal_text_det_batch_size=1,
  54. recovery=True,
  55. device=None,
  56. predictor_kwargs=None,
  57. ):
  58. super().__init__(
  59. predictor_kwargs=predictor_kwargs,
  60. )
  61. self._build_predictor(
  62. layout_model=layout_model,
  63. text_det_model=text_det_model,
  64. text_rec_model=text_rec_model,
  65. table_model=table_model,
  66. doc_image_ori_cls_model=doc_image_ori_cls_model,
  67. doc_image_unwarp_model=doc_image_unwarp_model,
  68. seal_text_det_model=seal_text_det_model,
  69. llm_name=llm_name,
  70. llm_params=llm_params,
  71. )
  72. self.set_predictor(
  73. layout_batch_size=layout_batch_size,
  74. text_det_batch_size=text_det_batch_size,
  75. text_rec_batch_size=text_rec_batch_size,
  76. table_batch_size=table_batch_size,
  77. doc_image_ori_cls_batch_size=doc_image_ori_cls_batch_size,
  78. doc_image_unwarp_batch_size=doc_image_unwarp_batch_size,
  79. seal_text_det_batch_size=seal_text_det_batch_size,
  80. device=device,
  81. )
  82. # get base prompt from yaml info
  83. if task_prompt_yaml:
  84. self.task_prompt_dict = read_yaml_file(task_prompt_yaml)
  85. else:
  86. self.task_prompt_dict = read_yaml_file(
  87. PROMPT_FILE
  88. ) # get user prompt from yaml info
  89. if user_prompt_yaml:
  90. self.user_prompt_dict = read_yaml_file(user_prompt_yaml)
  91. else:
  92. self.user_prompt_dict = None
  93. self.recovery = recovery
  94. self.visual_info = None
  95. self.vector = None
  96. self.visual_flag = False
  97. def _build_predictor(
  98. self,
  99. layout_model,
  100. text_det_model,
  101. text_rec_model,
  102. table_model,
  103. llm_name,
  104. llm_params,
  105. seal_text_det_model=None,
  106. doc_image_ori_cls_model=None,
  107. doc_image_unwarp_model=None,
  108. ):
  109. super()._build_predictor(
  110. layout_model, text_det_model, text_rec_model, table_model
  111. )
  112. if seal_text_det_model:
  113. self.curve_pipeline = self._create(
  114. pipeline=OCRPipeline,
  115. text_det_model=seal_text_det_model,
  116. text_rec_model=text_rec_model,
  117. )
  118. else:
  119. self.curve_pipeline = None
  120. if doc_image_ori_cls_model:
  121. self.oricls_predictor = self._create(doc_image_ori_cls_model)
  122. else:
  123. self.oricls_predictor = None
  124. if doc_image_unwarp_model:
  125. self.uvdoc_predictor = self._create(doc_image_unwarp_model)
  126. else:
  127. self.uvdoc_predictor = None
  128. self.img_reader = ReadImage(format="RGB")
  129. self.llm_api = create_llm_api(
  130. llm_name,
  131. llm_params,
  132. )
  133. self.cropper = CropByBoxes()
  134. def set_predictor(
  135. self,
  136. layout_batch_size=None,
  137. text_det_batch_size=None,
  138. text_rec_batch_size=None,
  139. table_batch_size=None,
  140. doc_image_ori_cls_batch_size=None,
  141. doc_image_unwarp_batch_size=None,
  142. seal_text_det_batch_size=None,
  143. device=None,
  144. ):
  145. if text_det_batch_size and text_det_batch_size > 1:
  146. logging.warning(
  147. f"text det model only support batch_size=1 now,the setting of text_det_batch_size={text_det_batch_size} will not using! "
  148. )
  149. if layout_batch_size:
  150. self.layout_predictor.set_predictor(batch_size=layout_batch_size)
  151. if text_rec_batch_size:
  152. self.ocr_pipeline.text_rec_model.set_predictor(
  153. batch_size=text_rec_batch_size
  154. )
  155. if table_batch_size:
  156. self.table_predictor.set_predictor(batch_size=table_batch_size)
  157. if self.curve_pipeline and seal_text_det_batch_size:
  158. self.curve_pipeline.text_det_model.set_predictor(
  159. batch_size=seal_text_det_batch_size
  160. )
  161. if self.oricls_predictor and doc_image_ori_cls_batch_size:
  162. self.oricls_predictor.set_predictor(batch_size=doc_image_ori_cls_batch_size)
  163. if self.uvdoc_predictor and doc_image_unwarp_batch_size:
  164. self.uvdoc_predictor.set_predictor(batch_size=doc_image_unwarp_batch_size)
  165. if device:
  166. if self.curve_pipeline:
  167. self.curve_pipeline.set_predictor(device=device)
  168. if self.oricls_predictor:
  169. self.oricls_predictor.set_predictor(device=device)
  170. if self.uvdoc_predictor:
  171. self.uvdoc_predictor.set_predictor(device=device)
  172. self.layout_batch_size.set_predictor(device=device)
  173. self.ocr_pipeline.set_predictor(device=device)
  174. def predict(
  175. self,
  176. input,
  177. use_doc_image_ori_cls_model=True,
  178. use_doc_image_unwarp_model=True,
  179. use_seal_text_det_model=True,
  180. recovery=True,
  181. **kwargs,
  182. ):
  183. self.set_predictor(**kwargs)
  184. visual_info = {"ocr_text": [], "table_html": [], "table_text": []}
  185. # get all visual result
  186. visual_result = list(
  187. self.get_visual_result(
  188. input,
  189. use_doc_image_ori_cls_model=use_doc_image_ori_cls_model,
  190. use_doc_image_unwarp_model=use_doc_image_unwarp_model,
  191. use_seal_text_det_model=use_seal_text_det_model,
  192. recovery=recovery,
  193. )
  194. )
  195. # decode visual result to get table_html, table_text, ocr_text
  196. ocr_text, table_text, table_html = self.decode_visual_result(visual_result)
  197. visual_info["ocr_text"] = ocr_text
  198. visual_info["table_html"] = table_html
  199. visual_info["table_text"] = table_text
  200. visual_info = VisualInfoResult(visual_info)
  201. # for local user save visual info in self
  202. self.visual_info = visual_info
  203. self.visual_flag = True
  204. return visual_result, visual_info
  205. def get_visual_result(
  206. self,
  207. inputs,
  208. use_doc_image_ori_cls_model=True,
  209. use_doc_image_unwarp_model=True,
  210. use_seal_text_det_model=True,
  211. recovery=True,
  212. ):
  213. # get oricls and uvdoc results
  214. img_info_list = list(self.img_reader(inputs))[0]
  215. oricls_results = []
  216. if self.oricls_predictor and use_doc_image_ori_cls_model:
  217. oricls_results = get_oriclas_results(img_info_list, self.oricls_predictor)
  218. uvdoc_results = []
  219. if self.uvdoc_predictor and use_doc_image_unwarp_model:
  220. uvdoc_results = get_uvdoc_results(img_info_list, self.uvdoc_predictor)
  221. img_list = [img_info["img"] for img_info in img_info_list]
  222. for idx, (img_info, layout_pred) in enumerate(
  223. zip(img_info_list, self.layout_predictor(img_list))
  224. ):
  225. single_img_res = {
  226. "input_path": "",
  227. "layout_result": DetResult({}),
  228. "ocr_result": OCRResult({}),
  229. "table_ocr_result": [],
  230. "table_result": StructureTableResult([]),
  231. "structure_result": [],
  232. "oricls_result": TopkResult({}),
  233. "uvdoc_result": DocTrResult({}),
  234. "curve_result": [],
  235. }
  236. # update oricls and uvdoc result
  237. if oricls_results:
  238. single_img_res["oricls_result"] = oricls_results[idx]
  239. if uvdoc_results:
  240. single_img_res["uvdoc_result"] = uvdoc_results[idx]
  241. # update layout result
  242. single_img_res["input_path"] = layout_pred["input_path"]
  243. single_img_res["layout_result"] = layout_pred
  244. single_img = img_info["img"]
  245. table_subs = []
  246. curve_subs = []
  247. structure_res = []
  248. ocr_res_with_layout = []
  249. if len(layout_pred["boxes"]) > 0:
  250. subs_of_img = list(self._crop_by_boxes(layout_pred))
  251. # get cropped images
  252. for sub in subs_of_img:
  253. box = sub["box"]
  254. xmin, ymin, xmax, ymax = [int(i) for i in box]
  255. mask_flag = True
  256. if sub["label"].lower() == "table":
  257. table_subs.append(sub)
  258. elif sub["label"].lower() == "seal":
  259. curve_subs.append(sub)
  260. else:
  261. if self.recovery and recovery:
  262. # TODO: Why use the entire image?
  263. wht_im = (
  264. np.ones(single_img.shape, dtype=single_img.dtype) * 255
  265. )
  266. wht_im[ymin:ymax, xmin:xmax, :] = sub["img"]
  267. sub_ocr_res = get_ocr_res(self.ocr_pipeline, wht_im)
  268. else:
  269. sub_ocr_res = get_ocr_res(self.ocr_pipeline, sub)
  270. sub_ocr_res["dt_polys"] = get_ori_coordinate_for_table(
  271. xmin, ymin, sub_ocr_res["dt_polys"]
  272. )
  273. layout_label = sub["label"].lower()
  274. if sub_ocr_res and sub["label"].lower() in [
  275. "image",
  276. "figure",
  277. "img",
  278. "fig",
  279. ]:
  280. mask_flag = False
  281. else:
  282. ocr_res_with_layout.append(sub_ocr_res)
  283. structure_res.append(
  284. {
  285. "layout_bbox": box,
  286. f"{layout_label}": "\n".join(
  287. sub_ocr_res["rec_text"]
  288. ),
  289. }
  290. )
  291. if mask_flag:
  292. single_img[ymin:ymax, xmin:xmax, :] = 255
  293. curve_pipeline = self.ocr_pipeline
  294. if self.curve_pipeline and use_seal_text_det_model:
  295. curve_pipeline = self.curve_pipeline
  296. all_curve_res = get_ocr_res(curve_pipeline, curve_subs)
  297. single_img_res["curve_result"] = all_curve_res
  298. if isinstance(all_curve_res, dict):
  299. all_curve_res = [all_curve_res]
  300. for sub, curve_res in zip(curve_subs, all_curve_res):
  301. structure_res.append(
  302. {
  303. "layout_bbox": sub["box"],
  304. "印章": "".join(curve_res["rec_text"]),
  305. }
  306. )
  307. ocr_res = get_ocr_res(self.ocr_pipeline, single_img)
  308. ocr_res["input_path"] = layout_pred["input_path"]
  309. all_table_res, _ = self.get_table_result(table_subs)
  310. for idx, single_dt_poly in enumerate(ocr_res["dt_polys"]):
  311. structure_res.append(
  312. {
  313. "layout_bbox": convert_4point2rect(single_dt_poly),
  314. "words in text block": ocr_res["rec_text"][idx],
  315. }
  316. )
  317. # update ocr result
  318. for layout_ocr_res in ocr_res_with_layout:
  319. ocr_res["dt_polys"].extend(layout_ocr_res["dt_polys"])
  320. ocr_res["rec_text"].extend(layout_ocr_res["rec_text"])
  321. ocr_res["input_path"] = single_img_res["input_path"]
  322. all_table_ocr_res = []
  323. # get table text from html
  324. structure_res_table, all_table_ocr_res = get_table_text_from_html(
  325. all_table_res
  326. )
  327. structure_res.extend(structure_res_table)
  328. # sort the layout result by the left top point of the box
  329. structure_res = sorted_layout_boxes(structure_res, w=single_img.shape[1])
  330. structure_res = [LayoutStructureResult(item) for item in structure_res]
  331. single_img_res["table_result"] = all_table_res
  332. single_img_res["ocr_result"] = ocr_res
  333. single_img_res["table_ocr_result"] = all_table_ocr_res
  334. single_img_res["structure_result"] = structure_res
  335. yield VisualResult(single_img_res)
  336. def decode_visual_result(self, visual_result):
  337. ocr_text = []
  338. table_text_list = []
  339. table_html = []
  340. for single_img_pred in visual_result:
  341. layout_res = single_img_pred["structure_result"]
  342. layout_res_copy = deepcopy(layout_res)
  343. # layout_res is [{"layout_bbox": [x1, y1, x2, y2], "layout": "single","words in text block":"xxx"}, {"layout_bbox": [x1, y1, x2, y2], "layout": "double","印章":"xxx"}
  344. ocr_res = {}
  345. for block in layout_res_copy:
  346. block.pop("layout_bbox")
  347. block.pop("layout")
  348. for layout_type, text in block.items():
  349. if text == "":
  350. continue
  351. # Table results are used separately
  352. if layout_type == "table":
  353. continue
  354. if layout_type not in ocr_res:
  355. ocr_res[layout_type] = text
  356. else:
  357. ocr_res[layout_type] += f"\n {text}"
  358. single_table_text = " ".join(single_img_pred["table_ocr_result"])
  359. for table_pred in single_img_pred["table_result"]:
  360. html = table_pred["html"]
  361. table_html.append(html)
  362. if ocr_res:
  363. ocr_text.append(ocr_res)
  364. table_text_list.append(single_table_text)
  365. return ocr_text, table_text_list, table_html
  366. def get_vector_text(
  367. self,
  368. llm_name=None,
  369. llm_params={},
  370. visual_info=None,
  371. min_characters=3500,
  372. llm_request_interval=1.0,
  373. ):
  374. """get vector for ocr"""
  375. if isinstance(self.llm_api, ErnieBot):
  376. get_vector_flag = True
  377. else:
  378. logging.warning("Do not use ErnieBot, will not get vector text.")
  379. get_vector_flag = False
  380. if not any([visual_info, self.visual_info]):
  381. return VectorResult({"vector": None})
  382. if visual_info:
  383. # use for serving or local
  384. _visual_info = visual_info
  385. else:
  386. # use for local
  387. _visual_info = self.visual_info
  388. ocr_text = _visual_info["ocr_text"]
  389. html_list = _visual_info["table_html"]
  390. table_text_list = _visual_info["table_text"]
  391. # add table text to ocr text
  392. for html, table_text_rec in zip(html_list, table_text_list):
  393. if len(html) > 3000:
  394. ocr_text.append({"table": table_text_rec})
  395. ocr_all_result = "".join(["\n".join(e.values()) for e in ocr_text])
  396. if len(ocr_all_result) > min_characters and get_vector_flag:
  397. if visual_info and llm_name:
  398. # for serving or local
  399. llm_api = create_llm_api(llm_name, llm_params)
  400. text_result = llm_api.get_vector(ocr_text, llm_request_interval)
  401. else:
  402. # for local
  403. text_result = self.llm_api.get_vector(ocr_text, llm_request_interval)
  404. else:
  405. text_result = str(ocr_text)
  406. self.visual_flag = False
  407. return VectorResult({"vector": text_result})
  408. def get_retrieval_text(
  409. self,
  410. key_list,
  411. visual_info=None,
  412. vector=None,
  413. llm_name=None,
  414. llm_params={},
  415. llm_request_interval=0.1,
  416. ):
  417. if not any([visual_info, vector, self.visual_info, self.vector]):
  418. return RetrievalResult({"retrieval": None})
  419. key_list = format_key(key_list)
  420. is_seving = visual_info and llm_name
  421. if self.visual_flag and not is_seving:
  422. self.vector = self.get_vector_text()
  423. if not any([vector, self.vector]):
  424. logging.warning(
  425. "The vector library is not created, and is being created automatically"
  426. )
  427. if is_seving:
  428. # for serving
  429. vector = self.get_vector_text(
  430. llm_name=llm_name, llm_params=llm_params, visual_info=visual_info
  431. )
  432. else:
  433. self.vector = self.get_vector_text()
  434. if vector and llm_name:
  435. _vector = vector["vector"]
  436. llm_api = create_llm_api(llm_name, llm_params)
  437. retrieval = llm_api.caculate_similar(
  438. vector=_vector,
  439. key_list=key_list,
  440. llm_params=llm_params,
  441. sleep_time=llm_request_interval,
  442. )
  443. else:
  444. _vector = self.vector["vector"]
  445. retrieval = self.llm_api.caculate_similar(
  446. vector=_vector, key_list=key_list, sleep_time=llm_request_interval
  447. )
  448. return RetrievalResult({"retrieval": retrieval})
  449. def chat(
  450. self,
  451. key_list,
  452. vector=None,
  453. visual_info=None,
  454. retrieval_result=None,
  455. user_task_description="",
  456. rules="",
  457. few_shot="",
  458. use_vector=True,
  459. save_prompt=False,
  460. llm_name="ernie-3.5",
  461. llm_params={},
  462. ):
  463. """
  464. chat with key
  465. """
  466. if not any(
  467. [vector, visual_info, retrieval_result, self.visual_info, self.vector]
  468. ):
  469. return ChatResult(
  470. {"chat_res": "请先完成图像解析再开始再对话", "prompt": ""}
  471. )
  472. key_list = format_key(key_list)
  473. # first get from table, then get from text in table, last get from all ocr
  474. if visual_info:
  475. # use for serving or local
  476. _visual_info = visual_info
  477. else:
  478. # use for local
  479. _visual_info = self.visual_info
  480. ocr_text = _visual_info["ocr_text"]
  481. html_list = _visual_info["table_html"]
  482. table_text_list = _visual_info["table_text"]
  483. prompt_res = {"ocr_prompt": "str", "table_prompt": [], "html_prompt": []}
  484. final_results = {}
  485. failed_results = ["大模型调用失败", "未知", "未找到关键信息", "None", ""]
  486. if html_list:
  487. prompt_list = self.get_prompt_for_table(
  488. html_list, key_list, rules, few_shot
  489. )
  490. prompt_res["html_prompt"] = prompt_list
  491. for prompt, table_text in zip(prompt_list, table_text_list):
  492. logging.debug(prompt)
  493. res = self.get_llm_result(prompt)
  494. # TODO: why use one html but the whole table_text in next step
  495. if list(res.values())[0] in failed_results:
  496. logging.info(
  497. "table html sequence is too much longer, using ocr directly"
  498. )
  499. prompt = self.get_prompt_for_ocr(
  500. table_text, key_list, rules, few_shot, user_task_description
  501. )
  502. logging.debug(prompt)
  503. prompt_res["table_prompt"].append(prompt)
  504. res = self.get_llm_result(prompt)
  505. for key, value in res.items():
  506. if value not in failed_results and key in key_list:
  507. key_list.remove(key)
  508. final_results[key] = value
  509. if len(key_list) > 0:
  510. logging.info("get result from ocr")
  511. if retrieval_result:
  512. ocr_text = retrieval_result.get("retrieval")
  513. elif use_vector and any([visual_info, vector]):
  514. # for serving or local
  515. ocr_text = self.get_retrieval_text(
  516. key_list=key_list,
  517. visual_info=visual_info,
  518. vector=vector,
  519. llm_name=llm_name,
  520. llm_params=llm_params,
  521. )["retrieval"]
  522. else:
  523. # for local
  524. ocr_text = self.get_retrieval_text(key_list=key_list)["retrieval"]
  525. prompt = self.get_prompt_for_ocr(
  526. ocr_text,
  527. key_list,
  528. rules,
  529. few_shot,
  530. user_task_description,
  531. )
  532. logging.debug(prompt)
  533. prompt_res["ocr_prompt"] = prompt
  534. res = self.get_llm_result(prompt)
  535. if res:
  536. final_results.update(res)
  537. if not res and not final_results:
  538. final_results = self.llm_api.ERROR_MASSAGE
  539. if save_prompt:
  540. return ChatResult({"chat_res": final_results, "prompt": prompt_res})
  541. else:
  542. return ChatResult({"chat_res": final_results, "prompt": ""})
  543. def get_llm_result(self, prompt):
  544. """get llm result and decode to dict"""
  545. llm_result = self.llm_api.pred(prompt)
  546. # when the llm pred failed, return None
  547. if not llm_result:
  548. return None
  549. if "json" in llm_result or "```" in llm_result:
  550. llm_result = (
  551. llm_result.replace("```", "").replace("json", "").replace("/n", "")
  552. )
  553. llm_result = llm_result.replace("[", "").replace("]", "")
  554. try:
  555. llm_result = json.loads(llm_result)
  556. llm_result_final = {}
  557. for key in llm_result:
  558. value = llm_result[key]
  559. if isinstance(value, list):
  560. if len(value) > 0:
  561. llm_result_final[key] = value[0]
  562. else:
  563. llm_result_final[key] = value
  564. return llm_result_final
  565. except:
  566. results = (
  567. llm_result.replace("\n", "")
  568. .replace(" ", "")
  569. .replace("{", "")
  570. .replace("}", "")
  571. )
  572. if not results.endswith('"'):
  573. results = results + '"'
  574. pattern = r'"(.*?)": "([^"]*)"'
  575. matches = re.findall(pattern, str(results))
  576. llm_result = {k: v for k, v in matches}
  577. return llm_result
  578. def get_prompt_for_table(self, table_result, key_list, rules="", few_shot=""):
  579. """get prompt for table"""
  580. prompt_key_information = []
  581. merge_table = ""
  582. for idx, result in enumerate(table_result):
  583. if len(merge_table + result) < 2000:
  584. merge_table += result
  585. if len(merge_table + result) > 2000 or idx == len(table_result) - 1:
  586. single_prompt = self.get_kie_prompt(
  587. merge_table,
  588. key_list,
  589. rules_str=rules,
  590. few_shot_demo_str=few_shot,
  591. prompt_type="table",
  592. )
  593. prompt_key_information.append(single_prompt)
  594. merge_table = ""
  595. return prompt_key_information
  596. def get_prompt_for_ocr(
  597. self,
  598. ocr_result,
  599. key_list,
  600. rules="",
  601. few_shot="",
  602. user_task_description="",
  603. ):
  604. """get prompt for ocr"""
  605. prompt_key_information = self.get_kie_prompt(
  606. ocr_result, key_list, user_task_description, rules, few_shot
  607. )
  608. return prompt_key_information
  609. def get_kie_prompt(
  610. self,
  611. text_result,
  612. key_list,
  613. user_task_description="",
  614. rules_str="",
  615. few_shot_demo_str="",
  616. prompt_type="common",
  617. ):
  618. """get_kie_prompt"""
  619. if prompt_type == "table":
  620. task_description = self.task_prompt_dict["kie_table_prompt"][
  621. "task_description"
  622. ]
  623. else:
  624. task_description = self.task_prompt_dict["kie_common_prompt"][
  625. "task_description"
  626. ]
  627. output_format = self.task_prompt_dict["kie_common_prompt"]["output_format"]
  628. if len(user_task_description) > 0:
  629. task_description = user_task_description
  630. task_description = task_description + output_format
  631. few_shot_demo_key_value = ""
  632. if self.user_prompt_dict:
  633. logging.info("======= common use custom ========")
  634. task_description = self.user_prompt_dict["task_description"]
  635. rules_str = self.user_prompt_dict["rules_str"]
  636. few_shot_demo_str = self.user_prompt_dict["few_shot_demo_str"]
  637. few_shot_demo_key_value = self.user_prompt_dict["few_shot_demo_key_value"]
  638. prompt = f"""{task_description}{rules_str}{few_shot_demo_str}{few_shot_demo_key_value}"""
  639. if prompt_type == "table":
  640. prompt += f"""\n结合上面,下面正式开始:\
  641. 表格内容:```{text_result}```\
  642. 关键词列表:[{key_list}]。""".replace(
  643. " ", ""
  644. )
  645. else:
  646. prompt += f"""\n结合上面的例子,下面正式开始:\
  647. OCR文字:```{text_result}```\
  648. 关键词列表:[{key_list}]。""".replace(
  649. " ", ""
  650. )
  651. return prompt