doc_understanding.py 3.8 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109
  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 time
  15. from typing import Any, List
  16. from .....utils.deps import function_requires_deps, is_dep_available
  17. from ...infra import utils as serving_utils
  18. from ...infra.config import AppConfig
  19. from ...schemas.doc_understanding import (
  20. INFER_ENDPOINT,
  21. ImageContent,
  22. ImageUrl,
  23. InferRequest,
  24. Message,
  25. RoleType,
  26. TextContent,
  27. )
  28. from .._app import create_app, primary_operation
  29. if is_dep_available("fastapi"):
  30. from fastapi import FastAPI
  31. if is_dep_available("openai"):
  32. from openai.types.chat import ChatCompletion
  33. from openai.types.chat.chat_completion import Choice as ChatCompletionChoice
  34. from openai.types.chat.chat_completion_message import ChatCompletionMessage
  35. @function_requires_deps("fastapi", "openai")
  36. def create_pipeline_app(pipeline: Any, app_config: AppConfig) -> "FastAPI":
  37. app, ctx = create_app(
  38. pipeline=pipeline, app_config=app_config, app_aiohttp_session=True
  39. )
  40. @primary_operation(
  41. app,
  42. INFER_ENDPOINT,
  43. "infer",
  44. )
  45. async def _infer(request: InferRequest) -> "ChatCompletion":
  46. pipeline = ctx.pipeline
  47. def _process_messages(messages: List[Message]):
  48. system_message = ""
  49. user_message = ""
  50. image_url = ""
  51. for msg in messages:
  52. if msg.role == RoleType.SYSTEM:
  53. if isinstance(msg.content, list):
  54. for content in msg.content:
  55. if isinstance(content, TextContent):
  56. system_message = content.text
  57. break
  58. else:
  59. system_message = msg.content
  60. elif msg.role == RoleType.USER:
  61. if isinstance(msg.content, list):
  62. for content in msg.content:
  63. if isinstance(content, str):
  64. user_message = content
  65. else:
  66. if isinstance(content, TextContent):
  67. user_message = content.text
  68. elif isinstance(content, ImageContent):
  69. image_url = content.image_url
  70. if isinstance(image_url, ImageUrl):
  71. image_url = image_url.url
  72. else:
  73. user_message = msg.content
  74. return system_message, user_message, image_url
  75. system_message, user_message, image_url = _process_messages(request.messages)
  76. result = (
  77. await pipeline.infer(
  78. {"image": image_url, "query": user_message},
  79. )
  80. )[0]
  81. return ChatCompletion(
  82. id=serving_utils.generate_log_id(),
  83. model=request.model,
  84. choices=[
  85. ChatCompletionChoice(
  86. index=0,
  87. finish_reason="stop",
  88. message=ChatCompletionMessage(
  89. role="assistant",
  90. content=result["result"],
  91. ),
  92. )
  93. ],
  94. created=int(time.time()),
  95. object="chat.completion",
  96. )
  97. return app