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)