Merge pull request ggerganov#1 from sumo43/master

Fix streaming
rooprob · Jul 23, 2023 · deb3818 · deb3818
2 parents ad67d5e + 0bddcd9
commit deb3818
Showing 1 changed file with 5 additions and 4 deletions.
diff --git a/run_wrap.py b/run_wrap.py
@@ -17,15 +17,16 @@
 
 t0 = time.time()
 tokens = []
+last = ''
 for line in proc.stdout:
     token = int(line.decode('utf-8').strip())
-    dec = enc.decode([token])
-    print(dec, end=" ", flush=True)
+    dec = enc.decode(tokens + [token])
+    chunk = dec[len(last):]
+    print(chunk, end='',flush=True)
     tokens.append(token)
+    last = dec
 t1 = time.time()
 
-print('\n---\n')
-print("Sorry I'm not sure why sentencepiece can't stream tokens properly, I'll solve it later. Here is the whole thing properly:")
 print('\n---\n')
 print(enc.decode(tokens))