Paula Leonova commited on
Commit
2b16dfe
1 Parent(s): ae7b1b1

Add early stopping to summarizer

Browse files
Files changed (1) hide show
  1. models.py +1 -1
models.py CHANGED
@@ -44,7 +44,7 @@ def load_summary_model():
44
  # return summarizer
45
 
46
  def summarizer_gen(summarizer, sequence:str, maximum_tokens:int, minimum_tokens:int):
47
- output = summarizer(sequence, num_beams=4, max_length=maximum_tokens, min_length=minimum_tokens, do_sample=False)
48
  return output[0].get('summary_text')
49
 
50
 
 
44
  # return summarizer
45
 
46
  def summarizer_gen(summarizer, sequence:str, maximum_tokens:int, minimum_tokens:int):
47
+ output = summarizer(sequence, num_beams=4, max_length=maximum_tokens, min_length=minimum_tokens, do_sample=False, early_stopping = True)
48
  return output[0].get('summary_text')
49
 
50