theluyuan/RWKV-Runner
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

improve compatible API

Browse Source
This commit is contained in:
josc146 2023-05-22 11:24:57 +08:00
parent 85493da730
commit 375af3bc1a
1 changed files with 7 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

9
backend-python/routes/completion.py
View File

@ -21,6 +21,7 @@ class ChatCompletionBody(ModelConfigBody):
messages: List[Message] messages: List[Message]
model: str = "rwkv" model: str = "rwkv"
stream: bool = False stream: bool = False
stop: str = None
completion_lock = Lock() completion_lock = Lock()
@ -70,7 +71,9 @@ async def chat_completions(body: ChatCompletionBody, request: Request):
set_rwkv_config(model, body) set_rwkv_config(model, body)
if body.stream: if body.stream:
for response, delta in rwkv_generate( for response, delta in rwkv_generate(
model, completion_text, stop="\n\nBob" model,
completion_text,
stop="\n\nBob" if body.stop is None else body.stop,
): ):
if await request.is_disconnected(): if await request.is_disconnected():
break break
@ -107,7 +110,9 @@ async def chat_completions(body: ChatCompletionBody, request: Request):
else: else:
response = None response = None
for response, delta in rwkv_generate( for response, delta in rwkv_generate(
model, completion_text, stop="\n\nBob" model,
completion_text,
stop="\n\nBob" if body.stop is None else body.stop,
): ):
if await request.is_disconnected(): if await request.is_disconnected():
break break