YongganFu commited on
Commit
b9d4dfe
·
verified ·
1 Parent(s): e59a204

Update README.md

Browse files
Files changed (1) hide show
  1. README.md +1 -1
README.md CHANGED
@@ -36,7 +36,7 @@ history.append({"role": "user", "content": user_input})
36
 
37
  prompt = tokenizer.apply_chat_template(history, tokenize=False, add_generation_prompt=True)
38
  prompt_ids = tokenizer(prompt, return_tensors='pt').input_ids.to(device='cuda')
39
- out_ids, nfe = model.generate(prompt_ids, max_new_tokens=128, steps=128, block_length=32, shift_logits=False, causal_context=True, threshold=0.9)
40
 
41
  tokenized_out = tokenizer.batch_decode(out_ids[:, prompt_ids.shape[1]:], skip_special_tokens=True)[0]
42
  print(f"Model: {tokenized_out}")
 
36
 
37
  prompt = tokenizer.apply_chat_template(history, tokenize=False, add_generation_prompt=True)
38
  prompt_ids = tokenizer(prompt, return_tensors='pt').input_ids.to(device='cuda')
39
+ out_ids, nfe = model.generate(prompt_ids, max_new_tokens=128, steps=128, block_length=32, shift_logits=False, causal_context=True, threshold=0.9, eos_token_id=tokenizer.eos_token_id)
40
 
41
  tokenized_out = tokenizer.batch_decode(out_ids[:, prompt_ids.shape[1]:], skip_special_tokens=True)[0]
42
  print(f"Model: {tokenized_out}")