import pickle import numpy as np from fastNLP.core.vocabulary import Vocabulary from fastNLP.io.base_loader import DataBundle from fastNLP.io.dataset_loader import JsonLoader from fastNLP.core.const import Const from tools.logger import * WORD_PAD = "[PAD]" WORD_UNK = "[UNK]" DOMAIN_UNK = "X" TAG_UNK = "X" class SummarizationLoader(JsonLoader): """ 读取summarization数据集，读取的DataSet包含fields:: text: list(str)，document summary: list(str), summary text_wd: list(list(str))，tokenized document summary_wd: list(list(str)), tokenized summary labels: list(int), flatten_label: list(int), 0 or 1, flatten labels domain: str, optional tag: list(str), optional 数据来源: CNN_DailyMail Newsroom DUC """ def __init__(self): super(SummarizationLoader, self).__init__() def _load(self, path): ds = super(SummarizationLoader, self)._load(path) def _lower_text(text_list): return [text.lower() for text in text_list] def _split_list(text_list): return [text.split() for text in text_list] def _convert_label(label, sent_len): np_label = np.zeros(sent_len, dtype=int) if label != []: np_label[np.array(label)] = 1 return np_label.tolist() ds.apply(lambda x: _lower_text(x['text']), new_field_name='text') ds.apply(lambda x: _lower_text(x['summary']), new_field_name='summary') ds.apply(lambda x:_split_list(x['text']), new_field_name='text_wd') ds.apply(lambda x:_split_list(x['summary']), new_field_name='summary_wd') ds.apply(lambda x:_convert_label(x["label"], len(x["text"])), new_field_name="flatten_label") return ds def process(self, paths, vocab_size, vocab_path, sent_max_len, doc_max_timesteps, domain=False, tag=False, load_vocab=True): """ :param paths: dict path for each dataset :param vocab_size: int max_size for vocab :param vocab_path: str vocab path :param sent_max_len: int max token number of the sentence :param doc_max_timesteps: int max sentence number of the document :param domain: bool build vocab for publication, use 'X' for unknown :param tag: bool build vocab for tag, use 'X' for unknown :param load_vocab: bool build vocab (False) or load vocab (True) :return: DataBundle datasets: dict keys correspond to the paths dict vocabs: dict key: vocab(if "train" in paths), domain(if domain=True), tag(if tag=True) embeddings: optional """ def _pad_sent(text_wd): pad_text_wd = [] for sent_wd in text_wd: if len(sent_wd) < sent_max_len: pad_num = sent_max_len - len(sent_wd) sent_wd.extend([WORD_PAD] * pad_num) else: sent_wd = sent_wd[:sent_max_len] pad_text_wd.append(sent_wd) return pad_text_wd def _token_mask(text_wd): token_mask_list = [] for sent_wd in text_wd: token_num = len(sent_wd) if token_num < sent_max_len: mask = [1] * token_num + [0] * (sent_max_len - token_num) else: mask = [1] * sent_max_len token_mask_list.append(mask) return token_mask_list def _pad_label(label): text_len = len(label) if text_len < doc_max_timesteps: pad_label = label + [0] * (doc_max_timesteps - text_len) else: pad_label = label[:doc_max_timesteps] return pad_label def _pad_doc(text_wd): text_len = len(text_wd) if text_len < doc_max_timesteps: padding = [WORD_PAD] * sent_max_len pad_text = text_wd + [padding] * (doc_max_timesteps - text_len) else: pad_text = text_wd[:doc_max_timesteps] return pad_text def _sent_mask(text_wd): text_len = len(text_wd) if text_len < doc_max_timesteps: sent_mask = [1] * text_len + [0] * (doc_max_timesteps - text_len) else: sent_mask = [1] * doc_max_timesteps return sent_mask datasets = {} train_ds = None for key, value in paths.items(): ds = self.load(value) # pad sent ds.apply(lambda x:_pad_sent(x["text_wd"]), new_field_name="pad_text_wd") ds.apply(lambda x:_token_mask(x["text_wd"]), new_field_name="pad_token_mask") # pad document ds.apply(lambda x:_pad_doc(x["pad_text_wd"]), new_field_name="pad_text") ds.apply(lambda x:_sent_mask(x["pad_text_wd"]), new_field_name="seq_len") ds.apply(lambda x:_pad_label(x["flatten_label"]), new_field_name="pad_label") # rename field ds.rename_field("pad_text", Const.INPUT) ds.rename_field("seq_len", Const.INPUT_LEN) ds.rename_field("pad_label", Const.TARGET) # set input and target ds.set_input(Const.INPUT, Const.INPUT_LEN) ds.set_target(Const.TARGET, Const.INPUT_LEN) datasets[key] = ds if "train" in key: train_ds = datasets[key] vocab_dict = {} if load_vocab == False: logger.info("[INFO] Build new vocab from training dataset!") if train_ds == None: raise ValueError("Lack train file to build vocabulary!") vocabs = Vocabulary(max_size=vocab_size, padding=WORD_PAD, unknown=WORD_UNK) vocabs.from_dataset(train_ds, field_name=["text_wd","summary_wd"]) vocab_dict["vocab"] = vocabs else: logger.info("[INFO] Load existing vocab from %s!" % vocab_path) word_list = [] with open(vocab_path, 'r', encoding='utf8') as vocab_f: cnt = 2 # pad and unk for line in vocab_f: pieces = line.split("\t") word_list.append(pieces[0]) cnt += 1 if cnt > vocab_size: break vocabs = Vocabulary(max_size=vocab_size, padding=WORD_PAD, unknown=WORD_UNK) vocabs.add_word_lst(word_list) vocabs.build_vocab() vocab_dict["vocab"] = vocabs if domain == True: domaindict = Vocabulary(padding=None, unknown=DOMAIN_UNK) domaindict.from_dataset(train_ds, field_name="publication") vocab_dict["domain"] = domaindict if tag == True: tagdict = Vocabulary(padding=None, unknown=TAG_UNK) tagdict.from_dataset(train_ds, field_name="tag") vocab_dict["tag"] = tagdict for ds in datasets.values(): vocab_dict["vocab"].index_dataset(ds, field_name=Const.INPUT, new_field_name=Const.INPUT) return DataBundle(vocabs=vocab_dict, datasets=datasets)