|
from transformers import Pipeline |
|
import torch |
|
from .utilities import padToSize |
|
from .summary import select, splitDocument |
|
|
|
|
|
""" |
|
Generates the segments ids for BERT |
|
""" |
|
def generateSegmentIds(doc_ids, tokenizer): |
|
|
|
segments_ids = [0] * len(doc_ids) |
|
curr_segment = 0 |
|
|
|
for i, token in enumerate(doc_ids): |
|
segments_ids[i] = curr_segment |
|
if token == tokenizer.vocab["[SEP]"]: |
|
curr_segment = 1 - curr_segment |
|
|
|
return segments_ids |
|
|
|
|
|
class ExtSummPipeline(Pipeline): |
|
""" |
|
Extractive summarization pipeline |
|
|
|
Inputs |
|
------ |
|
inputs : dict |
|
'sentences' : list[str] |
|
Sentences of the document |
|
|
|
strategy : str |
|
Strategy to summarize the document: |
|
- 'length': summary with a maximum length (strategy_args is the maximum length). |
|
- 'count': summary with the given number of sentences (strategy_args is the number of sentences). |
|
- 'ratio': summary proportional to the length of the document (strategy_args is the ratio [0, 1]). |
|
- 'threshold': summary only with sentences with a score higher than a given value (strategy_args is the minimum score). |
|
|
|
strategy_args : any |
|
Parameters of the strategy. |
|
|
|
Outputs |
|
------- |
|
selected_sents : list[str] |
|
List of the selected sentences |
|
|
|
selected_idxs : list[int] |
|
List of the indexes of the selected sentences in the original input |
|
""" |
|
|
|
|
|
def _sanitize_parameters(self, **kwargs): |
|
postprocess_kwargs = {} |
|
|
|
if ("strategy" in kwargs and "strategy_args" not in kwargs) or ("strategy" not in kwargs and "strategy_args" in kwargs): |
|
raise ValueError("`strategy` and `strategy_args` have to be both set") |
|
if "strategy" in kwargs: |
|
postprocess_kwargs["strategy"] = kwargs["strategy"] |
|
if "strategy_args" in kwargs: |
|
postprocess_kwargs["strategy_args"] = kwargs["strategy_args"] |
|
|
|
return {}, {}, postprocess_kwargs |
|
|
|
|
|
def preprocess(self, inputs): |
|
sentences = inputs["sentences"] |
|
|
|
|
|
doc_tokens = self.tokenizer.tokenize( f"{self.tokenizer.sep_token}{self.tokenizer.cls_token}".join(sentences) ) |
|
doc_tokens = [self.tokenizer.cls_token] + doc_tokens + [self.tokenizer.sep_token] |
|
doc_chunks = splitDocument(doc_tokens, self.tokenizer.cls_token, self.tokenizer.sep_token, self.model.config.input_size) |
|
|
|
|
|
batch = { |
|
"ids": [], |
|
"segments_ids": [], |
|
"clss_mask": [], |
|
"attn_mask": [], |
|
} |
|
for chunk_tokens in doc_chunks: |
|
doc_ids = self.tokenizer.convert_tokens_to_ids(chunk_tokens) |
|
segment_ids = generateSegmentIds(doc_ids, self.tokenizer) |
|
clss_mask = [True if token == self.tokenizer.cls_token_id else False for token in doc_ids] |
|
attn_mask = [1 for _ in range(len(doc_ids))] |
|
|
|
batch["ids"].append( padToSize(doc_ids, self.model.config.input_size, self.tokenizer.pad_token_id) ) |
|
batch["segments_ids"].append( padToSize(segment_ids, self.model.config.input_size, 0) ) |
|
batch["clss_mask"].append( padToSize(clss_mask, self.model.config.input_size, False) ) |
|
batch["attn_mask"].append( padToSize(attn_mask, self.model.config.input_size, 0) ) |
|
|
|
batch["ids"] = torch.as_tensor(batch["ids"]) |
|
batch["segments_ids"] = torch.as_tensor(batch["segments_ids"]) |
|
batch["clss_mask"] = torch.as_tensor(batch["clss_mask"]) |
|
batch["attn_mask"] = torch.as_tensor(batch["attn_mask"]) |
|
return { "inputs": batch, "sentences": sentences } |
|
|
|
|
|
def _forward(self, args): |
|
batch = args["inputs"] |
|
sentences = args["sentences"] |
|
out_predictions = torch.as_tensor([]).to(self.device) |
|
|
|
self.model.eval() |
|
with torch.no_grad(): |
|
batch_preds, _ = self.model(batch) |
|
for i, clss_mask in enumerate(batch["clss_mask"]): |
|
out_predictions = torch.cat((out_predictions, batch_preds[i][:torch.sum(clss_mask == True)])) |
|
|
|
return { "predictions": out_predictions, "sentences": sentences } |
|
|
|
|
|
def postprocess(self, args, strategy: str="count", strategy_args=3): |
|
predictions = args["predictions"] |
|
sentences = args["sentences"] |
|
return select(sentences, predictions, strategy, strategy_args) |
