ka1kuk commited on
Commit
9b03fdf
1 Parent(s): 6f83d4a

Update messagers/message_outputer.py

Browse files
Files changed (1) hide show
  1. messagers/message_outputer.py +17 -28
messagers/message_outputer.py CHANGED
@@ -1,17 +1,22 @@
1
  import json
 
 
2
 
3
  class OpenaiStreamOutputer:
4
  """
5
  Create chat completion - OpenAI API Documentation
6
  * https://platform.openai.com/docs/api-reference/chat/create
7
  """
 
8
 
9
  def __init__(self):
 
10
  self.default_data = {
11
- "id": "chatcmpl-123",
12
- "object": "chat.completion",
13
- "created": 1677652288,
14
- "model": "gpt-3.5-turbo-0613",
 
15
  "system_fingerprint": "fp_44709d6fcb",
16
  "choices": [],
17
  "usage": {
@@ -25,18 +30,14 @@ class OpenaiStreamOutputer:
25
  data_str = f"{json.dumps(data)}"
26
  return data_str
27
 
28
- def output(self, content=None, content_type="Completions", tokens_count=0) -> str:
29
  data = self.default_data.copy()
30
  if content_type == "Role":
31
  data["choices"] = [
32
  {
33
  "index": 0,
34
- "message": {
35
- "role": "assistant",
36
- "content": content,
37
- },
38
- "logprobs": None,
39
- "finish_reason": "stop"
40
  }
41
  ]
42
  elif content_type in [
@@ -50,11 +51,7 @@ class OpenaiStreamOutputer:
50
  data["choices"] = [
51
  {
52
  "index": 0,
53
- "message": {
54
- "role": "user",
55
- "content": content,
56
- },
57
- "logprobs": None,
58
  "finish_reason": None,
59
  }
60
  ]
@@ -62,11 +59,7 @@ class OpenaiStreamOutputer:
62
  data["choices"] = [
63
  {
64
  "index": 0,
65
- "message": {
66
- "role": "assistant",
67
- "content": content,
68
- },
69
- "logprobs": None,
70
  "finish_reason": "stop",
71
  }
72
  ]
@@ -74,18 +67,14 @@ class OpenaiStreamOutputer:
74
  data["choices"] = [
75
  {
76
  "index": 0,
77
- "message": {
78
- "role": "assistant",
79
- "content": content,
80
- },
81
- "logprobs": None,
82
  "finish_reason": None,
83
  }
84
  ]
85
 
86
- # Update token counts
87
  data["usage"]["prompt_tokens"] += tokens_count
88
  data["usage"]["completion_tokens"] += len(content.split())
89
  data["usage"]["total_tokens"] = data["usage"]["prompt_tokens"] + data["usage"]["completion_tokens"]
90
-
91
  return self.data_to_string(data, content_type)
 
 
1
  import json
2
+ import time
3
+
4
 
5
  class OpenaiStreamOutputer:
6
  """
7
  Create chat completion - OpenAI API Documentation
8
  * https://platform.openai.com/docs/api-reference/chat/create
9
  """
10
+
11
 
12
  def __init__(self):
13
+ current_time = int(time.time())
14
  self.default_data = {
15
+ "id": "chatcmpl-hugginface",
16
+ "object": "chat.completion.chunk",
17
+ "created": current_time,
18
+ # "content_type": "Completions",
19
+ "model": "hugginface",
20
  "system_fingerprint": "fp_44709d6fcb",
21
  "choices": [],
22
  "usage": {
 
30
  data_str = f"{json.dumps(data)}"
31
  return data_str
32
 
33
+ def output(self, content=None, content_type="Completions") -> str:
34
  data = self.default_data.copy()
35
  if content_type == "Role":
36
  data["choices"] = [
37
  {
38
  "index": 0,
39
+ "delta": {"role": "assistant"},
40
+ "finish_reason": None,
 
 
 
 
41
  }
42
  ]
43
  elif content_type in [
 
51
  data["choices"] = [
52
  {
53
  "index": 0,
54
+ "delta": {"content": content},
 
 
 
 
55
  "finish_reason": None,
56
  }
57
  ]
 
59
  data["choices"] = [
60
  {
61
  "index": 0,
62
+ "delta": {},
 
 
 
 
63
  "finish_reason": "stop",
64
  }
65
  ]
 
67
  data["choices"] = [
68
  {
69
  "index": 0,
70
+ "delta": {},
 
 
 
 
71
  "finish_reason": None,
72
  }
73
  ]
74
 
 
75
  data["usage"]["prompt_tokens"] += tokens_count
76
  data["usage"]["completion_tokens"] += len(content.split())
77
  data["usage"]["total_tokens"] = data["usage"]["prompt_tokens"] + data["usage"]["completion_tokens"]
78
+
79
  return self.data_to_string(data, content_type)
80
+