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