gmonsoon commited on
Commit
841b74f
1 Parent(s): add2c57

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +5 -0
app.py CHANGED
@@ -28,6 +28,10 @@ model = AutoModelForCausalLM.from_pretrained(
28
  device_map="auto",
29
  torch_dtype=torch.bfloat16,
30
  )
 
 
 
 
31
  model.config.sliding_window = 4096
32
  model.eval()
33
 
@@ -61,6 +65,7 @@ def generate(
61
  temperature=temperature,
62
  num_beams=1,
63
  repetition_penalty=repetition_penalty,
 
64
  )
65
  t = Thread(target=model.generate, kwargs=generate_kwargs)
66
  t.start()
 
28
  device_map="auto",
29
  torch_dtype=torch.bfloat16,
30
  )
31
+ terminators = [
32
+ tokenizer.eos_token_id,
33
+ tokenizer.convert_tokens_to_ids("<|eot_id|>")
34
+ ]
35
  model.config.sliding_window = 4096
36
  model.eval()
37
 
 
65
  temperature=temperature,
66
  num_beams=1,
67
  repetition_penalty=repetition_penalty,
68
+ eos_token_id=terminators,
69
  )
70
  t = Thread(target=model.generate, kwargs=generate_kwargs)
71
  t.start()