← back to Exo
chatgpt api repsonse streaming solves #20
8762effaf4392417d01c3e3f3f8d058d67b44ff8 · 2024-07-17 19:00:18 -0700 · Alex Cheema
Files touched
M exo/__init__.pyM exo/api/chatgpt_api.pyM exo/helpers.pyM setup.py
Diff
commit 8762effaf4392417d01c3e3f3f8d058d67b44ff8
Author: Alex Cheema <alexcheema123@gmail.com>
Date: Wed Jul 17 19:00:18 2024 -0700
chatgpt api repsonse streaming solves #20
---
exo/__init__.py | 2 +-
exo/api/chatgpt_api.py | 158 ++++++++++++++++++++++++++++++++++++++-----------
exo/helpers.py | 1 +
setup.py | 1 +
4 files changed, 126 insertions(+), 36 deletions(-)
diff --git a/exo/__init__.py b/exo/__init__.py
index 14f1e518..47a0b637 100644
--- a/exo/__init__.py
+++ b/exo/__init__.py
@@ -1 +1 @@
-from exo.helpers import DEBUG, DEBUG_DISCOVERY
+from exo.helpers import DEBUG, DEBUG_DISCOVERY, VERSION
diff --git a/exo/api/chatgpt_api.py b/exo/api/chatgpt_api.py
index b1968b98..f7b3d242 100644
--- a/exo/api/chatgpt_api.py
+++ b/exo/api/chatgpt_api.py
@@ -1,10 +1,12 @@
import uuid
import time
import asyncio
+import json
from transformers import AutoTokenizer
-from typing import List
+from typing import List, Literal, Union
from aiohttp import web
-from exo import DEBUG
+import aiohttp_cors
+from exo import DEBUG, VERSION
from exo.inference.shard import Shard
from exo.orchestration import Node
@@ -59,18 +61,77 @@ def resolve_tokenizer(model_id: str):
from exo.inference.mlx.sharded_utils import get_model_path, load_tokenizer
return load_tokenizer(get_model_path(model_id))
+def generate_completion(
+ chat_request: ChatCompletionRequest,
+ tokenizer,
+ prompt: str,
+ request_id: str,
+ tokens: List[int],
+ stream: bool,
+ finish_reason: Union[Literal["length", "stop"], None],
+ object_type: Literal["chat.completion", "text_completion"]
+ ) -> dict:
+ completion = {
+ "id": f"chatcmpl-{request_id}",
+ "object": object_type,
+ "created": int(time.time()),
+ "model": chat_request.model,
+ "system_fingerprint": f"exo_{VERSION}",
+ "choices": [
+ {
+ "index": 0,
+ "message": {
+ "role": "assistant",
+ "content": tokenizer.decode(tokens)
+ },
+ "logprobs": None,
+ "finish_reason": finish_reason,
+ }
+ ]
+ }
+
+ if not stream:
+ completion["usage"] = {
+ "prompt_tokens": len(tokenizer.encode(prompt)),
+ "completion_tokens": len(tokens),
+ "total_tokens": len(tokenizer.encode(prompt)) + len(tokens)
+ }
+
+ choice = completion["choices"][0]
+ if object_type.startswith("chat.completion"):
+ key_name = "delta" if stream else "message"
+ choice[key_name] = {"role": "assistant", "content": tokenizer.decode(tokens)}
+ elif object_type == "text_completion":
+ choice['text'] = tokenizer.decode(tokens)
+ else:
+ ValueError(f"Unsupported response type: {object_type}")
+
+ return completion
+
+
class ChatGPTAPI:
def __init__(self, node: Node, inference_engine_classname: str):
self.node = node
- self.app = web.Application()
- self.app.router.add_post('/v1/chat/completions', self.handle_post)
self.inference_engine_classname = inference_engine_classname
self.response_timeout_secs = 90
+ self.app = web.Application()
+ cors = aiohttp_cors.setup(self.app)
+ cors.add(self.app.router.add_post('/v1/chat/completions', self.handle_post), {
+ "*": aiohttp_cors.ResourceOptions(
+ allow_credentials=True,
+ expose_headers="*",
+ allow_headers="*",
+ allow_methods="*",
+ )
+ })
async def handle_post(self, request):
data = await request.json()
+ stream = data.get('stream', False)
messages = [Message(**msg) for msg in data['messages']]
chat_request = ChatCompletionRequest(data.get('model', 'llama-3-8b'), messages, data.get('temperature', 0.0))
+ if chat_request.model and chat_request.model.startswith("gpt-"): # to be compatible with ChatGPT tools, point all gpt- model requests to llama instead
+ chat_request.model = "llama-3-8b"
prompt = " ".join([msg.content for msg in chat_request.messages if msg.role == "user"])
shard = shard_mappings.get(chat_request.model, {}).get(self.inference_engine_classname)
if not shard:
@@ -83,6 +144,9 @@ class ChatGPTAPI:
chat_request.messages, tokenize=False, add_generation_prompt=True
)
+ callback_id = f"chatgpt-api-wait-response-{request_id}"
+ callback = self.node.on_token.register(callback_id)
+
if DEBUG >= 2: print(f"Sending prompt from ChatGPT api {request_id=} {shard=} {prompt=}")
try:
await self.node.process_prompt(shard, prompt, request_id=request_id)
@@ -92,40 +156,64 @@ class ChatGPTAPI:
traceback.print_exc()
return web.json_response({'detail': f"Error processing prompt (see logs): {str(e)}"}, status=500)
- callback_id = f"chatgpt-api-wait-response-{request_id}"
- callback = self.node.on_token.register(callback_id)
-
try:
if DEBUG >= 2: print(f"Waiting for response to finish. timeout={self.response_timeout_secs}s")
- _, result, _ = await callback.wait(lambda _request_id, tokens, is_finished: _request_id == request_id and is_finished, timeout=self.response_timeout_secs)
-
- eos_token_id = tokenizer.special_tokens_map.get("eos_token_id") if isinstance(tokenizer._tokenizer, AutoTokenizer) else tokenizer.eos_token_id
- if DEBUG >= 2: print(f"Checking if end of result {result[-1]=} is {eos_token_id=}")
- if result[-1] == eos_token_id:
- result = result[:-1]
-
- return web.json_response({
- "id": f"chatcmpl-{request_id}",
- "object": "chat.completion",
- "created": int(time.time()),
- "model": chat_request.model,
- "usage": {
- "prompt_tokens": len(tokenizer.encode(prompt)),
- "completion_tokens": len(result),
- "total_tokens": len(tokenizer.encode(prompt)) + len(result)
- },
- "choices": [
- {
- "message": {
- "role": "assistant",
- "content": tokenizer.decode(result)
- },
- "logprobs": None,
- "finish_reason": "stop",
- "index": 0
+
+ if stream:
+ response = web.StreamResponse(
+ status=200,
+ reason="OK",
+ headers={
+ "Content-Type": "application/json",
+ "Cache-Control": "no-cache",
+ # "Access-Control-Allow-Origin": "*",
+ # "Access-Control-Allow-Methods": "*",
+ # "Access-Control-Allow-Headers": "*",
}
- ]
- })
+ )
+ await response.prepare(request)
+
+ stream_task = None
+ last_tokens_len = 0
+ async def stream_result(request_id: str, tokens: List[int], is_finished: bool):
+ nonlocal last_tokens_len
+ prev_last_tokens_len = last_tokens_len
+ last_tokens_len = len(tokens)
+ new_tokens = tokens[prev_last_tokens_len:]
+ finish_reason = None
+ eos_token_id = tokenizer.special_tokens_map.get("eos_token_id") if isinstance(tokenizer._tokenizer, AutoTokenizer) else tokenizer.eos_token_id
+ if len(new_tokens) > 0 and new_tokens[-1] == eos_token_id:
+ new_tokens = new_tokens[:-1]
+ if is_finished:
+ finish_reason = "stop"
+ if is_finished and not finish_reason:
+ finish_reason = "length"
+
+ completion = generate_completion(chat_request, tokenizer, prompt, request_id, new_tokens, stream, finish_reason, "chat.completion")
+ if DEBUG >= 2: print(f"Streaming completion: {completion}")
+ await response.write(f"data: {json.dumps(completion)}\n\n".encode())
+ def on_result(_request_id: str, tokens: List[int], is_finished: bool):
+ nonlocal stream_task
+ stream_task = asyncio.create_task(stream_result(request_id, tokens, is_finished))
+
+ return _request_id == request_id and is_finished
+ _, tokens, _ = await callback.wait(on_result, timeout=self.response_timeout_secs)
+ if stream_task: # in case there is still a stream task running, wait for it to complete
+ if DEBUG >= 2: print(f"Pending stream task. Waiting for stream task to complete.")
+ await stream_task
+ await response.write_eof()
+ return response
+ else:
+ _, tokens, _ = await callback.wait(lambda _request_id, tokens, is_finished: _request_id == request_id and is_finished, timeout=self.response_timeout_secs)
+
+ finish_reason = "length"
+ eos_token_id = tokenizer.special_tokens_map.get("eos_token_id") if isinstance(tokenizer._tokenizer, AutoTokenizer) else tokenizer.eos_token_id
+ if DEBUG >= 2: print(f"Checking if end of tokens result {tokens[-1]=} is {eos_token_id=}")
+ if tokens[-1] == eos_token_id:
+ tokens = tokens[:-1]
+ finish_reason = "stop"
+
+ return web.json_response(generate_completion(chat_request, tokenizer, prompt, request_id, tokens, stream, finish_reason, "chat.completion"))
except asyncio.TimeoutError:
return web.json_response({'detail': "Response generation timed out"}, status=408)
finally:
diff --git a/exo/helpers.py b/exo/helpers.py
index b12b5d07..2d6eceea 100644
--- a/exo/helpers.py
+++ b/exo/helpers.py
@@ -4,6 +4,7 @@ from typing import Any, Callable, Coroutine, TypeVar, Optional, Dict, Generic, T
DEBUG = int(os.getenv("DEBUG", default="0"))
DEBUG_DISCOVERY = int(os.getenv("DEBUG_DISCOVERY", default="0"))
+VERSION = "0.0.1"
T = TypeVar('T')
K = TypeVar('K')
diff --git a/setup.py b/setup.py
index 9f229895..29e4b4f4 100644
--- a/setup.py
+++ b/setup.py
@@ -4,6 +4,7 @@ import sys
# Base requirements for all platforms
install_requires = [
"aiohttp==3.9.5",
+ "aiohttp_cors==0.7.0",
"grpcio==1.64.1",
"grpcio-tools==1.64.1",
"huggingface-hub==0.23.4",
← 5de2ea51 default to llama-3-8b and temperature=0 if not provided
·
back to Exo
·
support for streaming from non-tail nodes on chatgpt api, ad 8df2f4d8 →