File size: 1,372 Bytes
da27d7f 43c8935 da27d7f 43c8935 da27d7f 43c8935 5e6dc54 43c8935 5e6dc54 da27d7f 5e6dc54 da27d7f 5e6dc54 43c8935 da27d7f 5e6dc54 da27d7f 5e6dc54 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 |
import json
class OpenaiStreamOutputer:
"""
Create chat completion - OpenAI API Documentation
* https://platform.openai.com/docs/api-reference/chat/create
"""
def __init__(self):
self.default_data = {
"created": 1700000000,
"id": "chatcmpl-hugginface",
"object": "chat.completion",
"model": "hugginface",
"system_fingerprint": "fp_44709d6fcb",
"choices": [],
"usage": {
"prompt_tokens": 0,
"completion_tokens": 0,
"total_tokens": 0
}
}
def data_to_string(self, data={}):
data_str = f"{json.dumps(data)}"
return data_str
def output(self, content=None, role="assistant", content_type="Completions", tokens_info=None) -> str:
data = self.default_data.copy()
message = {"role": role, "content": content}
if tokens_info is None:
tokens_info = {"prompt_tokens": 9, "completion_tokens": 12, "total_tokens": 21}
data["choices"] = [
{
"index": 0,
"message": message,
"logprobs": None,
"finish_reason": "stop" if content_type == "Finished" else None
}
]
data["usage"] = tokens_info
return self.data_to_string(data)
|