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

fix py lock

Browse Source
This commit is contained in:
josc146 2023-05-21 13:46:54 +08:00
parent c02ce51d86
commit c3084a3290
1 changed files with 45 additions and 37 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

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

@ -51,63 +51,71 @@ async def completions(body: CompletionBody, request: Request):
while completion_lock.locked(): while completion_lock.locked():
await asyncio.sleep(0.1) await asyncio.sleep(0.1)
else: else:
with completion_lock: completion_lock.acquire()
set_rwkv_config(model, global_var.get(global_var.Model_Config)) set_rwkv_config(model, global_var.get(global_var.Model_Config))
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 await request.is_disconnected(): if await request.is_disconnected():
break break
yield json.dumps(
{
"response": response,
"model": "rwkv",
"choices": [
{
"delta": {"content": delta},
"index": 0,
"finish_reason": None,
}
],
}
)
yield json.dumps( yield json.dumps(
{ {
"response": response, "response": response,
"model": "rwkv", "model": "rwkv",
"choices": [ "choices": [
{ {
"delta": {}, "delta": {"content": delta},
"index": 0, "index": 0,
"finish_reason": "stop", "finish_reason": None,
} }
], ],
} }
) )
yield "[DONE]" if await request.is_disconnected():
else: completion_lock.release()
response = None return
for response, delta in rwkv_generate( yield json.dumps(
model, completion_text, stop="\n\nBob" {
):
pass
yield {
"response": response, "response": response,
"model": "rwkv", "model": "rwkv",
"choices": [ "choices": [
{ {
"message": { "delta": {},
"role": "assistant",
"content": response,
},
"index": 0, "index": 0,
"finish_reason": "stop", "finish_reason": "stop",
} }
], ],
} }
# torch_gc() )
yield "[DONE]"
else:
response = None
for response, delta in rwkv_generate(
model, completion_text, stop="\n\nBob"
):
if await request.is_disconnected():
break
if await request.is_disconnected():
completion_lock.release()
return
yield {
"response": response,
"model": "rwkv",
"choices": [
{
"message": {
"role": "assistant",
"content": response,
},
"index": 0,
"finish_reason": "stop",
}
],
}
# torch_gc()
completion_lock.release()
if body.stream: if body.stream:
return EventSourceResponse(eval_rwkv()) return EventSourceResponse(eval_rwkv())