zR commited on
Commit
995512e
1 Parent(s): 6f8d4ee

token check

Browse files
Files changed (1) hide show
  1. app.py +2 -4
app.py CHANGED
@@ -46,7 +46,7 @@ def predict(history, prompt, max_length, top_p, temperature):
46
 
47
  model_inputs = tokenizer.build_chat_input(query, history=messages, role='user').input_ids.to(
48
  next(model.parameters()).device)
49
- streamer = TextIteratorStreamer(tokenizer, timeout=600, skip_prompt=True, skip_special_tokens=True)
50
  eos_token_id = [tokenizer.eos_token_id, tokenizer.get_command("<|user|>"),
51
  tokenizer.get_command("<|observation|>")]
52
  generate_kwargs = {
@@ -63,9 +63,7 @@ def predict(history, prompt, max_length, top_p, temperature):
63
  t = Thread(target=model.generate, kwargs=generate_kwargs)
64
  t.start()
65
  for new_token in streamer:
66
- if new_token == '<|user|>':
67
- continue
68
- elif new_token:
69
  history[-1][1] += new_token
70
  yield history
71
 
 
46
 
47
  model_inputs = tokenizer.build_chat_input(query, history=messages, role='user').input_ids.to(
48
  next(model.parameters()).device)
49
+ streamer = TextIteratorStreamer(tokenizer, timeout=600, skip_prompt=True)
50
  eos_token_id = [tokenizer.eos_token_id, tokenizer.get_command("<|user|>"),
51
  tokenizer.get_command("<|observation|>")]
52
  generate_kwargs = {
 
63
  t = Thread(target=model.generate, kwargs=generate_kwargs)
64
  t.start()
65
  for new_token in streamer:
66
+ if new_token and '<|user|>' not in new_token:
 
 
67
  history[-1][1] += new_token
68
  yield history
69