ppchatocrv3.py 15 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447
  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 asyncio
  15. import os
  16. from typing import Awaitable, Final, List, Literal, Optional, Tuple, Union
  17. import cv2
  18. import numpy as np
  19. from fastapi import FastAPI, HTTPException
  20. from numpy.typing import ArrayLike
  21. from pydantic import BaseModel, Field
  22. from typing_extensions import Annotated, TypeAlias, assert_never
  23. from .....utils import logging
  24. from .... import results
  25. from ...ppchatocrv3 import PPChatOCRPipeline
  26. from .. import file_storage
  27. from .. import utils as serving_utils
  28. from ..app import AppConfig, create_app
  29. from ..models import Response, ResultResponse
  30. _DEFAULT_MAX_IMG_SIZE: Final[Tuple[int, int]] = (2000, 2000)
  31. _DEFAULT_MAX_NUM_IMGS: Final[int] = 10
  32. FileType: TypeAlias = Literal[0, 1]
  33. class InferenceParams(BaseModel):
  34. maxLongSide: Optional[Annotated[int, Field(gt=0)]] = None
  35. class AnalyzeImagesRequest(BaseModel):
  36. file: str
  37. fileType: Optional[FileType] = None
  38. useImgOrientationCls: bool = True
  39. useImgUnwrapping: bool = True
  40. useSealTextDet: bool = True
  41. inferenceParams: Optional[InferenceParams] = None
  42. Point: TypeAlias = Annotated[List[int], Field(min_length=2, max_length=2)]
  43. Polygon: TypeAlias = Annotated[List[Point], Field(min_length=3)]
  44. BoundingBox: TypeAlias = Annotated[List[float], Field(min_length=4, max_length=4)]
  45. class Text(BaseModel):
  46. poly: Polygon
  47. text: str
  48. score: float
  49. class Table(BaseModel):
  50. bbox: BoundingBox
  51. html: str
  52. class VisionResult(BaseModel):
  53. texts: List[Text]
  54. tables: List[Table]
  55. inputImage: str
  56. ocrImage: str
  57. layoutImage: str
  58. class AnalyzeImagesResult(BaseModel):
  59. visionResults: List[VisionResult]
  60. visionInfo: dict
  61. class QianfanParams(BaseModel):
  62. apiKey: str
  63. secretKey: str
  64. apiType: Literal["qianfan"] = "qianfan"
  65. class AIStudioParams(BaseModel):
  66. accessToken: str
  67. apiType: Literal["aistudio"] = "aistudio"
  68. LLMName: TypeAlias = Literal[
  69. "ernie-3.5",
  70. "ernie-3.5-8k",
  71. "ernie-lite",
  72. "ernie-4.0",
  73. "ernie-4.0-turbo-8k",
  74. "ernie-speed",
  75. "ernie-speed-128k",
  76. "ernie-tiny-8k",
  77. "ernie-char-8k",
  78. ]
  79. LLMParams: TypeAlias = Union[QianfanParams, AIStudioParams]
  80. class BuildVectorStoreRequest(BaseModel):
  81. visionInfo: dict
  82. minChars: Optional[int] = None
  83. llmRequestInterval: Optional[float] = None
  84. llmName: Optional[LLMName] = None
  85. llmParams: Optional[Annotated[LLMParams, Field(discriminator="apiType")]] = None
  86. class BuildVectorStoreResult(BaseModel):
  87. vectorStore: str
  88. class RetrieveKnowledgeRequest(BaseModel):
  89. keys: List[str]
  90. vectorStore: str
  91. llmName: Optional[LLMName] = None
  92. llmParams: Optional[Annotated[LLMParams, Field(discriminator="apiType")]] = None
  93. class RetrieveKnowledgeResult(BaseModel):
  94. retrievalResult: str
  95. class ChatRequest(BaseModel):
  96. keys: List[str]
  97. visionInfo: dict
  98. vectorStore: Optional[str] = None
  99. retrievalResult: Optional[str] = None
  100. taskDescription: Optional[str] = None
  101. rules: Optional[str] = None
  102. fewShot: Optional[str] = None
  103. llmName: Optional[LLMName] = None
  104. llmParams: Optional[Annotated[LLMParams, Field(discriminator="apiType")]] = None
  105. returnPrompts: bool = False
  106. class Prompts(BaseModel):
  107. ocr: str
  108. table: Optional[str] = None
  109. html: Optional[str] = None
  110. class ChatResult(BaseModel):
  111. chatResult: dict
  112. prompts: Optional[Prompts] = None
  113. def _llm_params_to_dict(llm_params: LLMParams) -> dict:
  114. if llm_params.apiType == "qianfan":
  115. return {
  116. "api_type": "qianfan",
  117. "ak": llm_params.apiKey,
  118. "sk": llm_params.secretKey,
  119. }
  120. if llm_params.apiType == "aistudio":
  121. return {"api_type": "aistudio", "access_token": llm_params.accessToken}
  122. else:
  123. assert_never(llm_params.apiType)
  124. def _postprocess_image(
  125. img: ArrayLike,
  126. request_id: str,
  127. filename: str,
  128. file_storage_config: file_storage.FileStorageConfig,
  129. ) -> str:
  130. key = f"{request_id}/{filename}"
  131. ext = os.path.splitext(filename)[1]
  132. img = np.asarray(img)
  133. _, encoded_img = cv2.imencode(ext, img)
  134. encoded_img = encoded_img.tobytes()
  135. return file_storage.postprocess_file(
  136. encoded_img, config=file_storage_config, key=key
  137. )
  138. def create_pipeline_app(pipeline: PPChatOCRPipeline, app_config: AppConfig) -> FastAPI:
  139. app, ctx = create_app(
  140. pipeline=pipeline, app_config=app_config, app_aiohttp_session=True
  141. )
  142. if "file_storage_config" in ctx.extra:
  143. ctx.extra["file_storage_config"] = file_storage.parse_file_storage_config(
  144. ctx.extra["file_storage_config"]
  145. )
  146. else:
  147. ctx.extra["file_storage_config"] = file_storage.InMemoryStorageConfig()
  148. ctx.extra.setdefault("max_img_size", _DEFAULT_MAX_IMG_SIZE)
  149. ctx.extra.setdefault("max_num_imgs", _DEFAULT_MAX_NUM_IMGS)
  150. @app.post(
  151. "/chatocr-vision",
  152. operation_id="analyzeImages",
  153. responses={422: {"model": Response}},
  154. )
  155. async def _analyze_images(
  156. request: AnalyzeImagesRequest,
  157. ) -> ResultResponse[AnalyzeImagesResult]:
  158. pipeline = ctx.pipeline
  159. aiohttp_session = ctx.aiohttp_session
  160. request_id = serving_utils.generate_request_id()
  161. if request.fileType is None:
  162. if serving_utils.is_url(request.file):
  163. try:
  164. file_type = serving_utils.infer_file_type(request.file)
  165. except Exception as e:
  166. logging.exception(e)
  167. raise HTTPException(
  168. status_code=422,
  169. detail="The file type cannot be inferred from the URL. Please specify the file type explicitly.",
  170. )
  171. else:
  172. raise HTTPException(status_code=422, detail="Unknown file type")
  173. else:
  174. file_type = "PDF" if request.fileType == 0 else "IMAGE"
  175. if request.inferenceParams:
  176. max_long_side = request.inferenceParams.maxLongSide
  177. if max_long_side:
  178. raise HTTPException(
  179. status_code=422,
  180. detail="`max_long_side` is currently not supported.",
  181. )
  182. try:
  183. file_bytes = await serving_utils.get_raw_bytes(
  184. request.file, aiohttp_session
  185. )
  186. images = await serving_utils.call_async(
  187. serving_utils.file_to_images,
  188. file_bytes,
  189. file_type,
  190. max_img_size=ctx.extra["max_img_size"],
  191. max_num_imgs=ctx.extra["max_num_imgs"],
  192. )
  193. result = await pipeline.call(
  194. pipeline.pipeline.visual_predict,
  195. images,
  196. use_doc_image_ori_cls_model=request.useImgOrientationCls,
  197. use_doc_image_unwarp_model=request.useImgUnwrapping,
  198. use_seal_text_det_model=request.useSealTextDet,
  199. )
  200. vision_results: List[VisionResult] = []
  201. for i, (img, item) in enumerate(zip(images, result[0])):
  202. pp_img_futures: List[Awaitable] = []
  203. future = serving_utils.call_async(
  204. _postprocess_image,
  205. img,
  206. request_id=request_id,
  207. filename=f"input_image_{i}.jpg",
  208. file_storage_config=ctx.extra["file_storage_config"],
  209. )
  210. pp_img_futures.append(future)
  211. future = serving_utils.call_async(
  212. _postprocess_image,
  213. item["ocr_result"].img,
  214. request_id=request_id,
  215. filename=f"ocr_image_{i}.jpg",
  216. file_storage_config=ctx.extra["file_storage_config"],
  217. )
  218. pp_img_futures.append(future)
  219. future = serving_utils.call_async(
  220. _postprocess_image,
  221. item["layout_result"].img,
  222. request_id=request_id,
  223. filename=f"layout_image_{i}.jpg",
  224. file_storage_config=ctx.extra["file_storage_config"],
  225. )
  226. pp_img_futures.append(future)
  227. texts: List[Text] = []
  228. for poly, text, score in zip(
  229. item["ocr_result"]["dt_polys"],
  230. item["ocr_result"]["rec_text"],
  231. item["ocr_result"]["rec_score"],
  232. ):
  233. texts.append(Text(poly=poly, text=text, score=score))
  234. tables = [
  235. Table(bbox=r["layout_bbox"], html=r["html"])
  236. for r in item["table_result"]
  237. ]
  238. input_img, ocr_img, layout_img = await asyncio.gather(*pp_img_futures)
  239. vision_result = VisionResult(
  240. texts=texts,
  241. tables=tables,
  242. inputImage=input_img,
  243. ocrImage=ocr_img,
  244. layoutImage=layout_img,
  245. )
  246. vision_results.append(vision_result)
  247. return ResultResponse(
  248. logId=serving_utils.generate_log_id(),
  249. errorCode=0,
  250. errorMsg="Success",
  251. result=AnalyzeImagesResult(
  252. visionResults=vision_results,
  253. visionInfo=result[1],
  254. ),
  255. )
  256. except Exception as e:
  257. logging.exception(e)
  258. raise HTTPException(status_code=500, detail="Internal server error")
  259. @app.post(
  260. "/chatocr-vector",
  261. operation_id="buildVectorStore",
  262. responses={422: {"model": Response}},
  263. )
  264. async def _build_vector_store(
  265. request: BuildVectorStoreRequest,
  266. ) -> ResultResponse[BuildVectorStoreResult]:
  267. pipeline = ctx.pipeline
  268. try:
  269. kwargs = {"visual_info": results.VisualInfoResult(request.visionInfo)}
  270. if request.minChars is not None:
  271. kwargs["min_characters"] = request.minChars
  272. if request.llmRequestInterval is not None:
  273. kwargs["llm_request_interval"] = request.llmRequestInterval
  274. if request.llmName is not None:
  275. kwargs["llm_name"] = request.llmName
  276. if request.llmParams is not None:
  277. kwargs["llm_params"] = _llm_params_to_dict(request.llmParams)
  278. result = await serving_utils.call_async(
  279. pipeline.pipeline.build_vector, **kwargs
  280. )
  281. return ResultResponse(
  282. logId=serving_utils.generate_log_id(),
  283. errorCode=0,
  284. errorMsg="Success",
  285. result=BuildVectorStoreResult(vectorStore=result["vector"]),
  286. )
  287. except Exception as e:
  288. logging.exception(e)
  289. raise HTTPException(status_code=500, detail="Internal server error")
  290. @app.post(
  291. "/chatocr-retrieval",
  292. operation_id="retrieveKnowledge",
  293. responses={422: {"model": Response}},
  294. )
  295. async def _retrieve_knowledge(
  296. request: RetrieveKnowledgeRequest,
  297. ) -> ResultResponse[RetrieveKnowledgeResult]:
  298. pipeline = ctx.pipeline
  299. try:
  300. kwargs = {
  301. "key_list": request.keys,
  302. "vector": results.VectorResult({"vector": request.vectorStore}),
  303. }
  304. if request.llmName is not None:
  305. kwargs["llm_name"] = request.llmName
  306. if request.llmParams is not None:
  307. kwargs["llm_params"] = _llm_params_to_dict(request.llmParams)
  308. result = await serving_utils.call_async(
  309. pipeline.pipeline.retrieval, **kwargs
  310. )
  311. return ResultResponse(
  312. logId=serving_utils.generate_log_id(),
  313. errorCode=0,
  314. errorMsg="Success",
  315. result=RetrieveKnowledgeResult(retrievalResult=result["retrieval"]),
  316. )
  317. except Exception as e:
  318. logging.exception(e)
  319. raise HTTPException(status_code=500, detail="Internal server error")
  320. @app.post(
  321. "/chatocr-chat",
  322. operation_id="chat",
  323. responses={422: {"model": Response}},
  324. response_model_exclude_none=True,
  325. )
  326. async def _chat(
  327. request: ChatRequest,
  328. ) -> ResultResponse[ChatResult]:
  329. pipeline = ctx.pipeline
  330. try:
  331. kwargs = {
  332. "key_list": request.keys,
  333. "visual_info": results.VisualInfoResult(request.visionInfo),
  334. }
  335. if request.vectorStore is not None:
  336. kwargs["vector"] = results.VectorResult({"vector": request.vectorStore})
  337. if request.retrievalResult is not None:
  338. kwargs["retrieval_result"] = results.RetrievalResult(
  339. {"retrieval": request.retrievalResult}
  340. )
  341. if request.taskDescription is not None:
  342. kwargs["user_task_description"] = request.taskDescription
  343. if request.rules is not None:
  344. kwargs["rules"] = request.rules
  345. if request.fewShot is not None:
  346. kwargs["few_shot"] = request.fewShot
  347. if request.llmName is not None:
  348. kwargs["llm_name"] = request.llmName
  349. if request.llmParams is not None:
  350. kwargs["llm_params"] = _llm_params_to_dict(request.llmParams)
  351. kwargs["save_prompt"] = request.returnPrompts
  352. result = await serving_utils.call_async(pipeline.pipeline.chat, **kwargs)
  353. if result["prompt"]:
  354. prompts = Prompts(
  355. ocr=result["prompt"]["ocr_prompt"],
  356. table=result["prompt"]["table_prompt"] or None,
  357. html=result["prompt"]["html_prompt"] or None,
  358. )
  359. else:
  360. prompts = None
  361. chat_result = ChatResult(
  362. chatResult=result["chat_res"],
  363. prompts=prompts,
  364. )
  365. return ResultResponse(
  366. logId=serving_utils.generate_log_id(),
  367. errorCode=0,
  368. errorMsg="Success",
  369. result=chat_result,
  370. )
  371. except Exception as e:
  372. logging.exception(e)
  373. raise HTTPException(status_code=500, detail="Internal server error")
  374. return app