Arturo Jiménez de los Galanes Reguillos commited on
Commit
73c4071
·
1 Parent(s): 43c389f

Add streamer kwargs

Browse files
Files changed (1) hide show
  1. app.py +3 -1
app.py CHANGED
@@ -36,7 +36,9 @@ def apply_chat_template(messages):
36
  tokenizer = AutoTokenizer.from_pretrained(MODEL)
37
  model = AutoModelForCausalLM.from_pretrained(MODEL, torch_dtype=torch.bfloat16, device_map="auto")
38
  model.eval()
39
- streamer = TextIteratorStreamer(tokenizer)
 
 
40
 
41
  cplusplus = None
42
  def translate(python):
 
36
  tokenizer = AutoTokenizer.from_pretrained(MODEL)
37
  model = AutoModelForCausalLM.from_pretrained(MODEL, torch_dtype=torch.bfloat16, device_map="auto")
38
  model.eval()
39
+
40
+ decode_kwargs = dict(skip_special_tokens=True)
41
+ streamer = TextIteratorStreamer(tokenizer, decode_kwargs=decode_kwargs)
42
 
43
  cplusplus = None
44
  def translate(python):