ka1kuk commited on
Commit
724ec44
·
1 Parent(s): 807a492

Update networks/message_streamer.py

Browse files
Files changed (1) hide show
  1. networks/message_streamer.py +14 -11
networks/message_streamer.py CHANGED
@@ -106,7 +106,6 @@ class MessageStreamer:
106
  return stream_response
107
 
108
  def chat_return_dict(self, stream_response):
109
- # https://platform.openai.com/docs/guides/text-generation/chat-completions-response-format
110
  final_output = self.message_outputer.default_data.copy()
111
  final_output["choices"] = [
112
  {
@@ -118,24 +117,28 @@ class MessageStreamer:
118
  },
119
  }
120
  ]
 
121
  logger.back(final_output)
122
-
123
  final_content = ""
124
  for line in stream_response.iter_lines():
125
  if not line:
126
  continue
 
127
  content = self.parse_line(line)
128
-
129
- if content.strip() == self.stop_sequences:
130
- logger.success("\n[Finished]")
131
- break
132
- else:
133
- logger.back(content, end="")
134
- final_content += content
135
-
 
 
136
  if self.model in self.STOP_SEQUENCES_MAP.keys():
137
  final_content = final_content.replace(self.stop_sequences, "")
138
-
139
  final_output["choices"][0]["message"]["content"] = final_content
140
  return final_output
141
 
 
106
  return stream_response
107
 
108
  def chat_return_dict(self, stream_response):
 
109
  final_output = self.message_outputer.default_data.copy()
110
  final_output["choices"] = [
111
  {
 
117
  },
118
  }
119
  ]
120
+
121
  logger.back(final_output)
122
+
123
  final_content = ""
124
  for line in stream_response.iter_lines():
125
  if not line:
126
  continue
127
+
128
  content = self.parse_line(line)
129
+
130
+ # Check if content is None before calling strip()
131
+ if content is not None:
132
+ if content.strip() == self.stop_sequences:
133
+ logger.success("\n[Finished]")
134
+ break
135
+ else:
136
+ logger.back(content, end="")
137
+ final_content += content
138
+
139
  if self.model in self.STOP_SEQUENCES_MAP.keys():
140
  final_content = final_content.replace(self.stop_sequences, "")
141
+
142
  final_output["choices"][0]["message"]["content"] = final_content
143
  return final_output
144