File size: 2,434 Bytes
da27d7f 9b03fdf 43c8935 da27d7f 43c8935 9b03fdf 43c8935 da27d7f 9b03fdf da27d7f 9b03fdf 5e6dc54 da27d7f 5e6dc54 da27d7f 6f83d4a 43c8935 da27d7f 9b03fdf da27d7f 6f83d4a 9b03fdf 6f83d4a 9b03fdf 6f83d4a 9b03fdf 6f83d4a 9b03fdf 6f83d4a 5e6dc54 6f83d4a 9b03fdf 6f83d4a 9b03fdf |
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 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 |
import json
import time
class OpenaiStreamOutputer:
"""
Create chat completion - OpenAI API Documentation
* https://platform.openai.com/docs/api-reference/chat/create
"""
def __init__(self):
current_time = int(time.time())
self.default_data = {
"id": "chatcmpl-hugginface",
"object": "chat.completion.chunk",
"created": current_time,
# "content_type": "Completions",
"model": "hugginface",
"system_fingerprint": "fp_44709d6fcb",
"choices": [],
"usage": {
"prompt_tokens": 0,
"completion_tokens": 0,
"total_tokens": 0
}
}
def data_to_string(self, data={}, content_type=""):
data_str = f"{json.dumps(data)}"
return data_str
def output(self, content=None, content_type="Completions") -> str:
data = self.default_data.copy()
if content_type == "Role":
data["choices"] = [
{
"index": 0,
"delta": {"role": "assistant"},
"finish_reason": None,
}
]
elif content_type in [
"Completions",
"InternalSearchQuery",
"InternalSearchResult",
"SuggestedResponses",
]:
if content_type in ["InternalSearchQuery", "InternalSearchResult"]:
content += "\n"
data["choices"] = [
{
"index": 0,
"delta": {"content": content},
"finish_reason": None,
}
]
elif content_type == "Finished":
data["choices"] = [
{
"index": 0,
"delta": {},
"finish_reason": "stop",
}
]
else:
data["choices"] = [
{
"index": 0,
"delta": {},
"finish_reason": None,
}
]
data["usage"]["prompt_tokens"] += tokens_count
data["usage"]["completion_tokens"] += len(content.split())
data["usage"]["total_tokens"] = data["usage"]["prompt_tokens"] + data["usage"]["completion_tokens"]
return self.data_to_string(data, content_type)
|