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)