nxphi47 commited on
Commit
226d418
1 Parent(s): 57dd736

Update multipurpose_chatbot/engines/transformers_engine.py

Browse files
multipurpose_chatbot/engines/transformers_engine.py CHANGED
@@ -212,6 +212,7 @@ class NewGenerationMixin(GenerationMixin):
212
  ['Today is a beautiful day, and we must do everything possible to make it a day of celebration.']
213
  ```"""
214
  # init values
 
215
  from transformers.generation.utils import (
216
  validate_stopping_criteria, GenerateEncoderDecoderOutput, GenerateDecoderOnlyOutput
217
  )
@@ -443,11 +444,11 @@ class TransformersEngine(BaseEngine):
443
  out_tokens = []
444
  response = None
445
  for token in generator:
446
- out_tokens.append(token.item())
447
  response = self.tokenizer.decode(out_tokens)
448
  num_tokens += 1
449
- # print(f"{num_tokens=}", end='\r')
450
- # sys.stdout.flush()
451
  yield response, num_tokens
452
 
453
  if response is not None:
 
212
  ['Today is a beautiful day, and we must do everything possible to make it a day of celebration.']
213
  ```"""
214
  # init values
215
+ print(f'Streaming tokens...')
216
  from transformers.generation.utils import (
217
  validate_stopping_criteria, GenerateEncoderDecoderOutput, GenerateDecoderOnlyOutput
218
  )
 
444
  out_tokens = []
445
  response = None
446
  for token in generator:
447
+ out_tokens.extend(token.tolist())
448
  response = self.tokenizer.decode(out_tokens)
449
  num_tokens += 1
450
+ print(f"{response}", end='\r')
451
+ sys.stdout.flush()
452
  yield response, num_tokens
453
 
454
  if response is not None: