File size: 1,976 Bytes
eb3e513
 
 
 
 
 
 
 
 
cd6b52a
d2b20f2
eb3e513
cd6b52a
eb3e513
d2b20f2
cd6b52a
eb3e513
 
d2b20f2
 
 
 
 
 
 
eb3e513
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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
import json


class OpenaiStreamOutputer:
    """
    Create chat completion - OpenAI API Documentation
    * https://platform.openai.com/docs/api-reference/chat/create
    """

    def __init__(self, owned_by="huggingface", model="mixtral-8x7b"):
        self.default_data = {
            "created": 1700000000,
            "id": f"chatcmpl-{owned_by}",
            "object": "chat.completion.chunk",
            # "content_type": "Completions",
            "model": model,
            "choices": [],
        }

    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,
                }
            ]
        return self.data_to_string(data, content_type)