Update Summarizer/Extractive.py

Summarizer/Extractive.py

@@ -1,4 +1,5 @@
 import nltk
+import torch 
 from summarizer import Summarizer
 from sumy.nlp.tokenizers import Tokenizer
 from sumy.summarizers.lsa import LsaSummarizer
@@ -37,6 +38,31 @@ def summarize(file, model):
                                   skip_special_tokens=True, 
                                   clean_up_tokenization_spaces=False)    
     summary =  summary[0]
+  elif model == "LEDBill":
+    tokenizer = AutoTokenizer.from_pretrained("d0r1h/LEDBill")
+    model = AutoModelForSeq2SeqLM.from_pretrained("d0r1h/LEDBill", return_dict_in_generate=True)
+    
+    input_ids = tokenizer(doc, return_tensors="pt").input_ids
+    global_attention_mask = torch.zeros_like(input_ids)
+    global_attention_mask[:, 0] = 1
+    
+    sequences = model.generate(input_ids, global_attention_mask=global_attention_mask).sequences
+    summary = tokenizer.batch_decode(sequences, skip_special_tokens=True)
+    
+    summary = summary[0]
+    
+  elif model == "ILC":
+    tokenizer = AutoTokenizer.from_pretrained("d0r1h/led-base-ilc")
+    model = AutoModelForSeq2SeqLM.from_pretrained("d0r1h/led-base-ilc", return_dict_in_generate=True)
+    
+    input_ids = tokenizer(doc, return_tensors="pt").input_ids
+    global_attention_mask = torch.zeros_like(input_ids)
+    global_attention_mask[:, 0] = 1
+    
+    sequences = model.generate(input_ids, global_attention_mask=global_attention_mask).sequences
+    summary = tokenizer.batch_decode(sequences, skip_special_tokens=True)
+    
+    summary = summary[0]
 
   elif model == "TextRank":
     summary = extractive(LexRankSummarizer(), doc)