File tree Expand file tree Collapse file tree 1 file changed +8
-1
lines changed
private_gpt/components/llm/custom Expand file tree Collapse file tree 1 file changed +8
-1
lines changed Original file line number Diff line number Diff line change @@ -243,12 +243,19 @@ def get_stream():
243
243
event_stream = resp ["Body" ]
244
244
start_json = b"{"
245
245
stop_token = "<|endoftext|>"
246
+ first_token = True
246
247
247
248
for line in LineIterator (event_stream ):
248
249
if line != b"" and start_json in line :
249
250
data = json .loads (line [line .find (start_json ) :].decode ("utf-8" ))
250
- if data ["token" ]["text" ] != stop_token :
251
+ special = data ["token" ]["special" ]
252
+ stop = data ["token" ]["text" ] == stop_token
253
+ if not special and not stop :
251
254
delta = data ["token" ]["text" ]
255
+ # trim the leading space for the first token if present
256
+ if first_token :
257
+ delta = delta .lstrip ()
258
+ first_token = False
252
259
text += delta
253
260
yield CompletionResponse (delta = delta , text = text , raw = data )
254
261
You can’t perform that action at this time.
0 commit comments