student2222333051 commited on
Commit
486dbd6
·
verified ·
1 Parent(s): 79c228d

Update summarizer.py

Browse files
Files changed (1) hide show
  1. summarizer.py +5 -3
summarizer.py CHANGED
@@ -3,13 +3,15 @@ from transformers import BartTokenizer, BartForConditionalGeneration
3
  tokenizer = BartTokenizer.from_pretrained("facebook/bart-large-cnn")
4
  model = BartForConditionalGeneration.from_pretrained("facebook/bart-large-cnn")
5
 
6
- def generate_summary(text: str, max_length=200, min_length=50) -> str:
7
  inputs = tokenizer([text], max_length=1024, truncation=True, return_tensors="pt")
 
8
  summary_ids = model.generate(
9
  inputs["input_ids"],
10
  num_beams=4,
11
- max_length=max_length,
12
- min_length=min_length,
13
  early_stopping=True
14
  )
 
15
  return tokenizer.decode(summary_ids[0], skip_special_tokens=True)
 
3
  tokenizer = BartTokenizer.from_pretrained("facebook/bart-large-cnn")
4
  model = BartForConditionalGeneration.from_pretrained("facebook/bart-large-cnn")
5
 
6
+ def generate_summary(text: str) -> str:
7
  inputs = tokenizer([text], max_length=1024, truncation=True, return_tensors="pt")
8
+
9
  summary_ids = model.generate(
10
  inputs["input_ids"],
11
  num_beams=4,
12
+ min_length=40,
13
+ max_length=200,
14
  early_stopping=True
15
  )
16
+
17
  return tokenizer.decode(summary_ids[0], skip_special_tokens=True)