reformat file proxy_llm.py

This commit is contained in:
xuyuan23 2023-06-13 17:42:11 +08:00
parent 6a8ee91834
commit 24457dc286

View File

@ -66,7 +66,7 @@ def proxyllm_generate_stream(model, tokenizer, params, device, context_len=2048)
"messages": history, "messages": history,
"temperature": params.get("temperature"), "temperature": params.get("temperature"),
"max_tokens": params.get("max_new_tokens"), "max_tokens": params.get("max_new_tokens"),
"stream": True "stream": True,
} }
res = requests.post( res = requests.post(
@ -76,12 +76,12 @@ def proxyllm_generate_stream(model, tokenizer, params, device, context_len=2048)
text = "" text = ""
for line in res.iter_lines(): for line in res.iter_lines():
if line: if line:
json_data = line.split(b': ', 1)[1] json_data = line.split(b": ", 1)[1]
decoded_line = json_data.decode("utf-8") decoded_line = json_data.decode("utf-8")
if decoded_line.lower() != '[DONE]'.lower(): if decoded_line.lower() != "[DONE]".lower():
obj = json.loads(json_data) obj = json.loads(json_data)
if obj['choices'][0]['delta'].get('content') is not None: if obj["choices"][0]["delta"].get("content") is not None:
content = obj['choices'][0]['delta']['content'] content = obj["choices"][0]["delta"]["content"]
text += content text += content
yield text yield text
@ -104,4 +104,4 @@ def proxyllm_generate_stream(model, tokenizer, params, device, context_len=2048)
# json_line = json.loads(decoded_line) # json_line = json.loads(decoded_line)
# print(json_line) # print(json_line)
# text += json_line["choices"][0]["message"]["content"] # text += json_line["choices"][0]["message"]["content"]
# yield text # yield text