BeveledCube commited on
Commit
30e32ac
1 Parent(s): 06b0b4a

Added it back to gpt2

Browse files
Files changed (2) hide show
  1. models/blenderbot.py +1 -1
  2. models/gpt2.py +1 -1
models/blenderbot.py CHANGED
@@ -23,6 +23,6 @@ def generate(input_text):
23
  input_ids = tokenizer.encode(input_text, return_tensors="pt")
24
 
25
  # Generate output using the model
26
- output_ids = model.generate(input_ids, num_beams=5, no_repeat_ngram_size=2)
27
 
28
  return tokenizer.decode(output_ids[0], skip_special_tokens=True)
 
23
  input_ids = tokenizer.encode(input_text, return_tensors="pt")
24
 
25
  # Generate output using the model
26
+ output_ids = model.generate(input_ids, no_repeat_ngram_size=2, max_new_tokens=100)
27
 
28
  return tokenizer.decode(output_ids[0], skip_special_tokens=True)
models/gpt2.py CHANGED
@@ -16,6 +16,6 @@ def generate(input_text):
16
  attention_mask = tf.ones_like(input_ids)
17
 
18
  # Generate output using the model
19
- output_ids = model.generate(input_ids, no_repeat_ngram_size=2, max_new_tokens=100)
20
 
21
  return tokenizer.decode(output_ids[0], skip_special_tokens=True)
 
16
  attention_mask = tf.ones_like(input_ids)
17
 
18
  # Generate output using the model
19
+ output_ids = model.generate(input_ids, num_beams=5, no_repeat_ngram_size=2, max_new_tokens=100)
20
 
21
  return tokenizer.decode(output_ids[0], skip_special_tokens=True)