EE21 commited on
Commit
e09e77e
1 Parent(s): 7a228d7

Update abstractive_summarization.py

Browse files
Files changed (1) hide show
  1. abstractive_summarization.py +8 -3
abstractive_summarization.py CHANGED
@@ -1,10 +1,15 @@
1
  from transformers import AutoTokenizer, AutoModelForSeq2SeqLM, pipeline
2
 
 
 
 
 
3
  # Function to summarize using the fine-tuned BART model
4
  def summarize_with_bart_ft(input_text):
5
- pipe_bart_ft = pipeline("summarization", model="EE21/BART-ToSSimplify")
6
- summary = pipe_bart_ft(input_text, max_length=300, min_length=100, num_beams=1, early_stopping=False, length_penalty=1)
7
- return summary[0]['summary_text']
 
8
 
9
  # Function to summarize using BART-large-cnn
10
  def summarize_with_bart_cnn(input_text):
 
1
  from transformers import AutoTokenizer, AutoModelForSeq2SeqLM, pipeline
2
 
3
+ # Load the fine-tuned BART tokenizer and model
4
+ tokenizer = AutoTokenizer.from_pretrained("EE21/BART-ToSSimplify")
5
+ model = AutoModelForSeq2SeqLM.from_pretrained("EE21/BART-ToSSimplify")
6
+
7
  # Function to summarize using the fine-tuned BART model
8
  def summarize_with_bart_ft(input_text):
9
+ inputs = tokenizer.encode("summarize: " + input_text, return_tensors="pt", max_length=1024, truncation=True)
10
+ summary_ids = model.generate(inputs, max_length=300, min_length=100, num_beams=1, early_stopping=False, length_penalty=1)
11
+ summary = tokenizer.decode(summary_ids[0], skip_special_tokens=False)
12
+ return summary
13
 
14
  # Function to summarize using BART-large-cnn
15
  def summarize_with_bart_cnn(input_text):