serving_chat.py 15 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345
  1. import codecs
  2. import time
  3. from typing import AsyncGenerator, AsyncIterator, List, Optional, Union
  4. from fastapi import Request
  5. from loguru import logger
  6. from aphrodite.common.outputs import RequestOutput
  7. from aphrodite.common.utils import random_uuid
  8. from aphrodite.endpoints.openai.protocol import (
  9. ChatCompletionRequest,
  10. ChatCompletionResponse,
  11. ChatCompletionResponseChoice,
  12. ChatCompletionResponseStreamChoice,
  13. ChatCompletionStreamResponse,
  14. ChatMessage,
  15. DeltaMessage,
  16. ErrorResponse,
  17. UsageInfo,
  18. )
  19. from aphrodite.endpoints.openai.serving_engine import LoRA, OpenAIServing
  20. from aphrodite.engine.async_aphrodite import AsyncAphrodite
  21. from aphrodite.modeling.guided_decoding import (
  22. get_guided_decoding_logits_processor)
  23. class OpenAIServingChat(OpenAIServing):
  24. def __init__(self,
  25. engine: AsyncAphrodite,
  26. served_model_names: List[str],
  27. response_role: str,
  28. lora_modules: Optional[List[LoRA]] = None,
  29. chat_template=None):
  30. super().__init__(engine=engine,
  31. served_model_names=served_model_names,
  32. lora_modules=lora_modules)
  33. self.response_role = response_role
  34. self._load_chat_template(chat_template)
  35. async def create_chat_completion(
  36. self, request: ChatCompletionRequest, raw_request: Request
  37. ) -> Union[ErrorResponse, AsyncGenerator[str, None],
  38. ChatCompletionResponse]:
  39. """Completion API similar to OpenAI's API.
  40. See https://platform.openai.com/docs/api-reference/chat/create
  41. for the API specification. This API mimics the OpenAI
  42. ChatCompletion API.
  43. NOTE: Currently we do not support the following feature:
  44. - function_call (Users should implement this by themselves)
  45. """
  46. error_check_ret = await self._check_model(request)
  47. if error_check_ret is not None:
  48. return error_check_ret
  49. # Deal with list in messages.content
  50. # Just replace the content list with the very first text message
  51. for message in request.messages:
  52. if message.role == "user" and isinstance(message["content"], list):
  53. message["content"] = next((content["text"]
  54. for content in message["content"]
  55. if content["type"] == "text"), "")
  56. try:
  57. prompt = self.tokenizer.apply_chat_template(
  58. conversation=request.messages,
  59. tokenize=False,
  60. add_generation_prompt=request.add_generation_prompt)
  61. except Exception as e:
  62. logger.error(
  63. f"Error in applying chat template from request: {str(e)}")
  64. return self.create_error_response(str(e))
  65. request_id = f"cmpl-{random_uuid()}"
  66. try:
  67. # Tokenize/detokenize depending on prompt format (string/token list)
  68. prompt_ids, prompt_text = self._validate_prompt_and_tokenize(
  69. request, prompt=prompt)
  70. sampling_params = request.to_sampling_params(
  71. self.tokenizer.vocab_size)
  72. lora_request = self._maybe_get_lora(request)
  73. guided_decode_logits_processor = (
  74. await get_guided_decoding_logits_processor(
  75. request.guided_decoding_backend, request, await
  76. self.engine.get_tokenizer()))
  77. if guided_decode_logits_processor:
  78. sampling_params.logits_processors.append(
  79. guided_decode_logits_processor)
  80. except ValueError as e:
  81. return self.create_error_response(str(e))
  82. result_generator = self.engine.generate(prompt_text, sampling_params,
  83. request_id, prompt_ids,
  84. lora_request)
  85. # Streaming response
  86. if request.stream:
  87. return self.chat_completion_stream_generator(
  88. request, result_generator, request_id)
  89. else:
  90. try:
  91. return await self.chat_completion_full_generator(
  92. request, raw_request, result_generator, request_id)
  93. except ValueError as e:
  94. # TODO: Use an aphrodite-specific Validation Error
  95. return self.create_error_response(str(e))
  96. def get_chat_request_role(self, request: ChatCompletionRequest) -> str:
  97. if request.add_generation_prompt:
  98. return self.response_role
  99. else:
  100. return request.messages[-1]["role"]
  101. async def chat_completion_stream_generator(
  102. self, request: ChatCompletionRequest,
  103. result_generator: AsyncIterator[RequestOutput], request_id: str
  104. ) -> Union[ErrorResponse, AsyncGenerator[str, None]]:
  105. model_name = self.served_model_names[0]
  106. created_time = int(time.time())
  107. chunk_object_type = "chat.completion.chunk"
  108. first_iteration = True
  109. # Send response for each token for each request.n (index)
  110. previous_texts = [""] * request.n
  111. previous_num_tokens = [0] * request.n
  112. finish_reason_sent = [False] * request.n
  113. try:
  114. async for res in result_generator:
  115. res: RequestOutput
  116. # We need to do it here, because if there are exceptions in
  117. # the result_generator, it needs to be sent as the FIRST
  118. # response (by the try...catch).
  119. if first_iteration:
  120. # Send first response for each request.n (index) with
  121. # the role
  122. role = self.get_chat_request_role(request)
  123. for i in range(request.n):
  124. choice_data = ChatCompletionResponseStreamChoice(
  125. index=i,
  126. delta=DeltaMessage(role=role),
  127. logprobs=None,
  128. finish_reason=None)
  129. chunk = ChatCompletionStreamResponse(
  130. id=request_id,
  131. object=chunk_object_type,
  132. created=created_time,
  133. choices=[choice_data],
  134. model=model_name)
  135. data = chunk.model_dump_json(exclude_unset=True)
  136. yield f"data: {data}\n\n"
  137. # Send response to echo the input portion of the
  138. # last message
  139. if request.echo:
  140. last_msg_content = ""
  141. if request.messages and isinstance(
  142. request.messages,
  143. list) and request.messages[-1].get(
  144. "content") and request.messages[-1].get(
  145. "role") == role:
  146. last_msg_content = request.messages[-1]["content"]
  147. if last_msg_content:
  148. for i in range(request.n):
  149. choice_data = (
  150. ChatCompletionResponseStreamChoice(
  151. index=i,
  152. delta=DeltaMessage(
  153. content=last_msg_content),
  154. finish_reason=None))
  155. chunk = ChatCompletionStreamResponse(
  156. id=request_id,
  157. object=chunk_object_type,
  158. created=created_time,
  159. choices=[choice_data],
  160. logprobs=None,
  161. model=model_name)
  162. data = chunk.model_dump_json(
  163. exclude_unset=True)
  164. yield f"data: {data}\n\n"
  165. first_iteration = False
  166. for output in res.outputs:
  167. i = output.index
  168. if finish_reason_sent[i]:
  169. continue
  170. delta_token_ids = output.token_ids[previous_num_tokens[i]:]
  171. top_logprobs = output.logprobs[
  172. previous_num_tokens[i]:] if output.logprobs else None
  173. if request.logprobs:
  174. logprobs = self._create_logprobs(
  175. token_ids=delta_token_ids,
  176. top_logprobs=top_logprobs,
  177. num_output_top_logprobs=request.logprobs,
  178. initial_text_offset=len(previous_texts[i]),
  179. )
  180. else:
  181. logprobs = None
  182. delta_text = output.text[len(previous_texts[i]):]
  183. previous_texts[i] = output.text
  184. previous_num_tokens[i] = len(output.token_ids)
  185. if output.finish_reason is None:
  186. # Send token-by-token response for each request.n
  187. choice_data = ChatCompletionResponseStreamChoice(
  188. index=i,
  189. delta=DeltaMessage(content=delta_text),
  190. logprobs=logprobs,
  191. finish_reason=None)
  192. chunk = ChatCompletionStreamResponse(
  193. id=request_id,
  194. object=chunk_object_type,
  195. created=created_time,
  196. choices=[choice_data],
  197. model=model_name)
  198. data = chunk.model_dump_json(exclude_unset=True)
  199. yield f"data: {data}\n\n"
  200. else:
  201. # Send the finish response for each request.n only once
  202. prompt_tokens = len(res.prompt_token_ids)
  203. final_usage = UsageInfo(
  204. prompt_tokens=prompt_tokens,
  205. completion_tokens=previous_num_tokens[i],
  206. total_tokens=prompt_tokens +
  207. previous_num_tokens[i],
  208. )
  209. choice_data = ChatCompletionResponseStreamChoice(
  210. index=i,
  211. delta=DeltaMessage(content=delta_text),
  212. logprobs=logprobs,
  213. finish_reason=output.finish_reason,
  214. stop_reason=output.stop_reason)
  215. chunk = ChatCompletionStreamResponse(
  216. id=request_id,
  217. object=chunk_object_type,
  218. created=created_time,
  219. choices=[choice_data],
  220. model=model_name)
  221. if final_usage is not None:
  222. chunk.usage = final_usage
  223. data = chunk.model_dump_json(exclude_unset=True,
  224. exclude_none=True)
  225. yield f"data: {data}\n\n"
  226. finish_reason_sent[i] = True
  227. except ValueError as e:
  228. # TODO: Use an aphrodite-specific Validation Error
  229. data = self.create_streaming_error_response(str(e))
  230. yield f"data: {data}\n\n"
  231. # Send the final done message after all response.n are finished
  232. yield "data: [DONE]\n\n"
  233. async def chat_completion_full_generator(
  234. self, request: ChatCompletionRequest, raw_request: Request,
  235. result_generator: AsyncIterator[RequestOutput],
  236. request_id: str) -> Union[ErrorResponse, ChatCompletionResponse]:
  237. model_name = self.served_model_names[0]
  238. created_time = int(time.time())
  239. final_res: RequestOutput = None
  240. async for res in result_generator:
  241. if await raw_request.is_disconnected():
  242. # Abort the request if the client disconnects.
  243. await self.engine.abort(request_id)
  244. return self.create_error_response("Client disconnected")
  245. final_res = res
  246. assert final_res is not None
  247. choices = []
  248. role = self.get_chat_request_role(request)
  249. for output in final_res.outputs:
  250. token_ids = output.token_ids
  251. top_logprobs = output.logprobs
  252. if request.logprobs:
  253. logprobs = self._create_logprobs(
  254. token_ids=token_ids,
  255. top_logprobs=top_logprobs,
  256. num_output_top_logprobs=request.logprobs,
  257. )
  258. else:
  259. logprobs = None
  260. choice_data = ChatCompletionResponseChoice(
  261. index=output.index,
  262. message=ChatMessage(role=role, content=output.text),
  263. logprobs=logprobs,
  264. finish_reason=output.finish_reason,
  265. stop_reason=output.stop_reason,
  266. )
  267. choices.append(choice_data)
  268. if request.echo:
  269. last_msg_content = ""
  270. if request.messages and isinstance(
  271. request.messages, list) and request.messages[-1].get(
  272. "content") and request.messages[-1].get(
  273. "role") == role:
  274. last_msg_content = request.messages[-1]["content"]
  275. for choice in choices:
  276. full_message = last_msg_content + choice.message.content
  277. choice.message.content = full_message
  278. num_prompt_tokens = len(final_res.prompt_token_ids)
  279. num_generated_tokens = sum(
  280. len(output.token_ids) for output in final_res.outputs)
  281. usage = UsageInfo(
  282. prompt_tokens=num_prompt_tokens,
  283. completion_tokens=num_generated_tokens,
  284. total_tokens=num_prompt_tokens + num_generated_tokens,
  285. )
  286. response = ChatCompletionResponse(
  287. id=request_id,
  288. created=created_time,
  289. model=model_name,
  290. choices=choices,
  291. usage=usage,
  292. )
  293. return response
  294. def _load_chat_template(self, chat_template):
  295. if chat_template is not None:
  296. try:
  297. with open(chat_template, "r") as f:
  298. self.tokenizer.chat_template = f.read()
  299. except OSError:
  300. # If opening a file fails, set chat template to be args to
  301. # ensure we decode so our escape are interpreted correctly
  302. self.tokenizer.chat_template = codecs.decode(
  303. chat_template, "unicode_escape")
  304. logger.info("Using the supplied chat template.")
  305. elif self.tokenizer.chat_template is not None:
  306. logger.info("Using the default chat template")
  307. else:
  308. logger.warning(
  309. "No chat template provided. Chat API will not work.")