Update Summarizer/Extractive.py

Summarizer/Extractive.py

@@ -63,6 +63,20 @@ def summarize(file, model):
     summary = tokenizer.batch_decode(sequences, skip_special_tokens=True)
     
     summary = summary[0]
+  elif model == "Distill":
+    checkpoint = "sshleifer/distill-pegasus-cnn-16-4"
+    tokenizer = AutoTokenizer.from_pretrained(checkpoint)
+    model = AutoModelForSeq2SeqLM.from_pretrained(checkpoint)
+    inputs = tokenizer(doc, 
+                    max_length=1024, 
+                    truncation=True,
+                    return_tensors="pt")
+    
+    summary_ids = model.generate(inputs["input_ids"])
+    summary = tokenizer.batch_decode(summary_ids, 
+                                  skip_special_tokens=True, 
+                                  clean_up_tokenization_spaces=False)    
+    summary =  summary[0]    
 
   elif model == "TextRank":
     summary = extractive(LexRankSummarizer(), doc)