Upload 6 files

plms/exceptions.py

@@ -0,0 +1,33 @@
+""" Custom exceptions """
+
+
+class ExceedMaxLengthError(Exception):
+    """ Token exceed max length. """
+
+    def __init__(self, max_length=None):
+        self.message = f'Input sentence exceeds max length of {max_length}'
+        super().__init__(self.message)
+
+
+class HighlightNotFoundError(Exception):
+    """ Highlight is not in the sentence. """
+
+    def __init__(self, highlight: str, input_sentence: str):
+        self.message = f'Highlight `{highlight}` not found in the input sentence `{input_sentence}`'
+        super().__init__(self.message)
+
+
+class AnswerNotFoundError(Exception):
+    """ Answer cannot found in the context. """
+
+    def __init__(self, context: str):
+        self.message = f'Model cannot find any answer candidates in `{context}`'
+        super().__init__(self.message)
+
+
+class APIError(Exception):
+    """ Error from huggingface inference API. """
+
+    def __init__(self, context: str):
+        self.message = f'Huggingface API Error:\n`{context}`'
+        super().__init__(self.message)

plms/language_model.py

@@ -0,0 +1,613 @@
+import os
+import logging
+import pickle
+import re
+import urllib
+from itertools import chain
+from typing import List, Dict
+from multiprocessing import Pool
+import numpy as np
+from tqdm import tqdm
+import torch
+from torch.nn import functional
+import transformers
+from .exceptions import ExceedMaxLengthError, HighlightNotFoundError, AnswerNotFoundError
+from .spacy_module import SpacyPipeline, VALID_METHODS
+
+__all__ = ('TransformersQG', 'ADDITIONAL_SP_TOKENS', 'TASK_PREFIX', 'clean', 'internet_connection')
+
+os.environ["TOKENIZERS_PARALLELISM"] = "false"  # to turn off warning message
+TASK_PREFIX = {
+    "ae": "extract answers",
+    "qg": "generate question",
+    "qag": "generate question and answer",
+    "qa": "answer question"
+}
+CE_IGNORE_INDEX = -100
+ADDITIONAL_SP_TOKENS = {'hl': '<hl>'}
+NUM_WORKERS = int(os.getenv('NUM_WORKERS', '0'))
+PARALLEL_PROCESSING = bool(int(os.getenv('PARALLEL_PROCESSING', '0')))
+DEFAULT_MODELS = {
+    'vi': 'VietAI/vit5-base'
+}
+
+def pickle_save(obj, path: str):
+    with open(path, "wb") as fp:
+        pickle.dump(obj, fp)
+
+
+def pickle_load(path: str):
+    with open(path, "rb") as fp:  # Unpickling
+        return pickle.load(fp)
+
+
+def clean(string):
+    string = re.sub(r'\A\s*', '', string)
+    string = re.sub(r'\s*\Z', '', string)
+    if len(string) > 0:
+        return string
+    return None
+
+
+def internet_connection(host='http://google.com'):
+    try:
+        urllib.request.urlopen(host)
+        return True
+    except:
+        return False
+
+
+def load_language_model(model_name,
+                        cache_dir: str = None,
+                        use_auth_token: bool = False,
+                        torch_dtype=None,
+                        device_map: str = None,
+                        low_cpu_mem_usage: bool = False):
+    """ load language model from huggingface model hub """
+    # tokenizer
+    local_files_only = not internet_connection()
+    tokenizer = transformers.AutoTokenizer.from_pretrained(
+        model_name, cache_dir=cache_dir, local_files_only=local_files_only, use_auth_token=use_auth_token)
+    config = transformers.AutoConfig.from_pretrained(
+        model_name, local_files_only=local_files_only, cache_dir=cache_dir, use_auth_token=use_auth_token)
+    # model
+    if config.model_type == 't5':  # T5 model requires T5ForConditionalGeneration class
+        model_class = transformers.T5ForConditionalGeneration.from_pretrained
+    elif config.model_type == 'mt5':
+        model_class = transformers.MT5ForConditionalGeneration.from_pretrained
+    elif config.model_type == 'bart':
+        model_class = transformers.BartForConditionalGeneration.from_pretrained
+    elif config.model_type == 'mbart':
+        model_class = transformers.MBartForConditionalGeneration.from_pretrained
+    elif config.model_type == 'switch_transformers':
+        model_class = transformers.SwitchTransformersForConditionalGeneration.from_pretrained
+    else:
+        raise ValueError(f'unsupported model type: {config.model_type}')
+
+    param = {'config': config, "local_files_only": local_files_only, "use_auth_token": use_auth_token,
+             "low_cpu_mem_usage": low_cpu_mem_usage, "cache_dir": cache_dir}
+    if torch_dtype is not None:
+        param['torch_dtype'] = torch_dtype
+    if device_map is not None:
+        param['device_map'] = device_map
+    model = model_class(model_name, **param)
+    # add new special tokens to the tokenizer and the model if they don't have it
+    tokenizer.add_special_tokens({'additional_special_tokens': list(ADDITIONAL_SP_TOKENS.values())})
+    model.resize_token_embeddings(len(tokenizer))
+    return tokenizer, model, config
+
+
+def label_smoothed_loss(logits, labels, epsilon):
+    """ https://github.com/huggingface/transformers/blob/55bb4c06f7be141c6d895dbe1f11018dc8580b2d/src/transformers/trainer_pt_utils.py#L430 """
+    log_probs = - functional.log_softmax(logits, dim=-1)
+    if labels.dim() == log_probs.dim() - 1:
+        labels = labels.unsqueeze(-1)
+
+    padding_mask = labels.eq(CE_IGNORE_INDEX)
+    # In case the ignore_index is -100, the gather will fail, so we replace labels by 0. The padding_mask
+    # will ignore them in any case.
+    labels.clamp_min_(0)
+
+    nll_loss = log_probs.gather(dim=-1, index=labels)
+    nll_loss.masked_fill_(padding_mask, 0.0)
+
+    # works for fp16 input tensor too, by internally upcasting it to fp32
+    smoothed_loss = log_probs.sum(dim=-1, keepdim=True, dtype=torch.float32)
+    smoothed_loss.masked_fill_(padding_mask, 0.0)
+
+    # Take the mean over the label dimensions, then divide by the number of active elements (i.e. not-padded):
+    num_active_elements = padding_mask.numel() - padding_mask.long().sum()
+    nll_loss = nll_loss.sum() / num_active_elements
+    smoothed_loss = smoothed_loss.sum() / (num_active_elements * log_probs.shape[-1])
+    return (1 - epsilon) * nll_loss + epsilon * smoothed_loss
+
+
+class Dataset(torch.utils.data.Dataset):
+    """ torch.utils.data.Dataset wrapper converting into tensor """
+    float_tensors = ['attention_mask']
+
+    def __init__(self, data: List):
+        self.data = data
+
+    def __len__(self):
+        return len(self.data)
+
+    def to_tensor(self, name, data):
+        if name in self.float_tensors:
+            return torch.tensor(data, dtype=torch.float32)
+        return torch.tensor(data, dtype=torch.long)
+
+    def __getitem__(self, idx):
+        return {k: self.to_tensor(k, v) for k, v in self.data[idx].items()}
+
+
+class EncodePlus:
+    """ Wrapper of encode_plus for multiprocessing. """
+
+    def __init__(self,
+                 tokenizer,
+                 max_length: int = 512,
+                 max_length_output: int = 34,
+                 drop_overflow_error_text: bool = False,
+                 skip_overflow_error: bool = False,
+                 drop_highlight_error_text: bool = False,
+                 prefix_type: str = None,
+                 padding: bool = True):
+        """ Wrapper of encode_plus for multiprocessing.
+
+        @param tokenizer: transforms.Tokenizer
+        @param max_length: Max text length of input.
+        @param max_length_output: Max text length of output.
+        @param drop_overflow_error_text: If true, return None when the input exceeds the max length.
+        @param skip_overflow_error: If true, raise an error when the input exceeds the max length.
+        @param drop_highlight_error_text: If true, raise an error when a highlight span is not found in the paragraph.
+        @param prefix_type: Either of `qg` or `answer_extraction`, which is to add at the beginning of the text.
+        @param padding: Pad the sequence to the max length.
+        """
+        self.prefix = TASK_PREFIX[prefix_type] if prefix_type is not None else None
+        self.tokenizer = tokenizer
+        self.max_length = max_length
+        self.max_length_output = max_length_output
+        # NOTE: for model training, we should drop the exceeded input but not for the evaluator
+        self.drop_overflow_error_text = drop_overflow_error_text
+        self.skip_overflow_error = skip_overflow_error
+        self.drop_highlight_error_text = drop_highlight_error_text
+        # truncation should be true for the batch process, but not necessary to process single input
+        self.param_in = {'truncation': True, 'max_length': self.max_length}
+        self.param_out = {'truncation': True, 'max_length': self.max_length_output}
+        if padding:
+            self.param_in['padding'] = 'max_length'
+            self.param_out['padding'] = 'max_length'
+
+    def __call__(self, inputs):
+        return self.encode_plus(*inputs)
+
+    def encode_plus(self, input_sequence: str, output_sequence: str = None, input_highlight: str = None):
+        """ encode_plus
+
+        @param input_sequence: Input sequence.
+        @param output_sequence: Output sequence.
+        @param input_highlight: Sub-sequence of `input_sequence` to be surrounded by <hl>.
+        @return: The output of `encode_plus`.
+        """
+        # add highlight to the input
+        if input_highlight is not None:
+            position = input_sequence.find(input_highlight)
+            if position == -1:
+                if self.drop_highlight_error_text:
+                    return None
+                raise HighlightNotFoundError(input_highlight, input_sequence)
+            input_sequence = '{0}{1} {2} {1}{3}'.format(
+                input_sequence[:position], ADDITIONAL_SP_TOKENS['hl'], input_highlight,
+                input_sequence[position+len(input_highlight):])
+        if self.prefix is not None:
+            input_sequence = f'{self.prefix}: {input_sequence}'
+
+        # handling overflow text
+        # drop_overflow_error_text ==> remove the overflow sentence from input
+        # skip_overflow_error ==> keep the overflow sentence
+        # none of them ==> raise error
+        if self.drop_overflow_error_text or not self.skip_overflow_error:
+            if len(self.tokenizer.encode(input_sequence)) > self.max_length:
+                if not self.drop_overflow_error_text:  # raise error for overflow text
+                    raise ExceedMaxLengthError(self.max_length)
+                return None  # remove overflow text
+            if output_sequence is not None:
+                if len(self.tokenizer.encode(output_sequence)) > self.max_length_output:
+                    if not self.drop_overflow_error_text:  # raise error for overflow text
+                        raise ExceedMaxLengthError(self.max_length)
+                    return None  # remove overflow text
+        if type(self.tokenizer) is transformers.models.mbart.tokenization_mbart_fast.MBartTokenizerFast:
+            encode = self.tokenizer(input_sequence, **self.param_in)
+        else:
+            encode = self.tokenizer(text_target=input_sequence, **self.param_in)
+        if output_sequence is not None:
+            encode['labels'] = self.tokenizer.encode(output_sequence, **self.param_out)
+        return encode
+
+
+class TransformersQG:
+    """ Transformers Language Model for Question Generation. """
+
+    def __init__(self,
+                 model: str = None,
+                 max_length: int = 512,
+                 max_length_output: int = 256,
+                 model_ae: str = None,
+                 max_length_ae: int = 512,
+                 max_length_output_ae: int = 64,
+                 cache_dir: str = None,
+                 add_prefix: bool = None,
+                 language: str = 'vi',
+                 label_smoothing: float = None,
+                 skip_overflow_error: bool = False,
+                 drop_overflow_error_text: bool = False,
+                 drop_highlight_error_text: bool = False,
+                 drop_answer_error_text: bool = False,
+                 use_auth_token: bool = False,
+                 torch_dtype=None,
+                 device_map: str = None,
+                 low_cpu_mem_usage: bool = False,
+                 is_qg: bool = None,
+                 is_qag: bool = None,
+                 is_qa: bool = None,
+                 is_ae: bool = None):
+        """ Transformers Language Model for Question Generation.
+
+        @param model: Model alias or path to local model file.
+        @param max_length: Max text length of input.
+        @param max_length_output: Max text length of output.
+        @param cache_dir: Directory to cache transformers model files.
+        @param add_prefix: Whether model uses task-specific prefix (eg. True for T5 but False for BART models).
+        @param language: Language alias for SpaCy language-specific pipelines (sentencizer/keyword extraction).
+        @param label_smoothing: [Fine-tuning parameter] Label smoothing.
+        @param drop_overflow_error_text: If true, return None when the input exceeds the max length.
+        @param skip_overflow_error: If true, raise an error when the input exceeds the max length.
+        @param drop_highlight_error_text: If true, raise an error when a highlight span is not found in the paragraph.
+        @param use_auth_token: [optional] Huggingface transformers argument of `use_auth_token`
+        """
+
+        # take default model given the language
+        if model is None:
+            assert language in DEFAULT_MODELS.keys(),\
+                f"Model with language '{language}' is not available. Please choose language from " \
+                f"'{DEFAULT_MODELS.keys()}' or specify 'model'."
+            model = DEFAULT_MODELS[language]
+
+        # classify model type
+        self.is_qg = 'qg' in model.split('-') if is_qg is None else is_qg
+        self.is_ae = 'ae' in model.split('-') if is_ae is None else is_ae
+        self.is_qa = 'qa' in model.split('-') if is_qa is None else is_qa
+        self.is_qag = 'qag' in model.split('-') if is_qag is None else is_qag
+        # configs
+        self.model_name = model
+        self.max_length = max_length
+        self.max_length_output = max_length_output
+        self.label_smoothing = label_smoothing
+        self.drop_overflow_error_text = drop_overflow_error_text
+        self.skip_overflow_error = skip_overflow_error
+        self.drop_highlight_error_text = drop_highlight_error_text
+        self.drop_answer_error_text = drop_answer_error_text
+        self.model_name_ae = model_ae
+        self.max_length_ae = max_length_ae
+        self.max_length_output_ae = max_length_output_ae
+        # load model
+        self.tokenizer, self.model, config = load_language_model(
+            self.model_name, cache_dir=cache_dir, use_auth_token=use_auth_token, device_map=device_map,
+            torch_dtype=torch_dtype, low_cpu_mem_usage=low_cpu_mem_usage)
+        if 'add_prefix' not in config.to_dict().keys():
+            # this means the model is not fine-tuned
+            # assert add_prefix, '`add_prefix` is required for non-fine-tuned models'
+            self.add_prefix = add_prefix
+        else:
+            self.add_prefix = config.add_prefix
+
+        # set default behaviour for answer extraction
+        if self.model_name_ae is None:
+            self.model_name_ae = self.model_name if self.is_ae else "positionrank"
+        # load answer extraction model
+        self.answer_model_type = None
+        if self.model_name_ae in VALID_METHODS:
+            logging.info(f'use spaCy answer extraction model: {self.model_name_ae}')
+            self.tokenizer_ae = self.model_ae = self.add_prefix_ae = None
+            self.spacy_module = SpacyPipeline(language, self.model_name_ae)
+            self.answer_model_type = 'spacy'
+        else:
+            logging.info(f'use LMQG fine-tuned answer extraction model: {self.model_name_ae}')
+            if self.model_name == self.model_name_ae:
+                logging.info("the same model as QG is used as AE")
+                assert self.is_ae, f"the model ({self.model_name_ae}) is not fine-tuned for AE"
+                self.tokenizer_ae = self.model_ae = self.add_prefix_ae = None
+                self.answer_model_type = 'multitask'
+            else:
+                logging.info(f"loading 2nd model for AE: {self.model_name_ae}")
+                self.tokenizer_ae, self.model_ae, config_ae = load_language_model(model_ae, cache_dir=cache_dir, use_auth_token=use_auth_token)
+                self.add_prefix_ae = config_ae.add_prefix
+                self.answer_model_type = 'pipeline'
+            self.spacy_module = SpacyPipeline(language)
+
+        # GPU setup
+        self.device = 'cuda' if torch.cuda.device_count() > 0 else 'cpu'
+        self.parallel = False
+        if torch.cuda.device_count() > 1:
+            self.parallel = True
+            self.model = torch.nn.DataParallel(self.model)
+            if self.model_ae is not None:
+                self.model_ae = torch.nn.DataParallel(self.model_ae)
+        self.model.to(self.device)
+        if self.model_ae is not None:
+            self.model_ae.to(self.device)
+        logging.info(f'Model `{self.model_name}`')
+        logging.info(f'\t * Num of GPU in use: {torch.cuda.device_count()}')
+        logging.info(f'\t * Prefix: {self.add_prefix}')
+        logging.info(f'\t * Language: {language} (ignore at the training phase)')
+
+    def generate_qa_end2end(self,
+                            list_context: str or List,
+                            batch_size: int = None,
+                            num_beams: int = 4,
+                            cache_path: str = None,
+                            splitting_symbol: str = ' [SEP] ',
+                            question_prefix: str = "question: ",
+                            answer_prefix: str = ", answer: "):
+        """ Generate question from paragraph and answer. Note that `list_answer` is needed unless they are already
+        highlighted in the `list_context`. eg) "I live in <hl> Tokyo <hl>."
+
+        @param list_context: List of input texts.
+        @param batch_size: Batch size.
+        @param num_beams: Number of beam for model generation.
+        @param cache_path: Path to pre-compute features.
+        @return: List of generated sentences.
+        """
+        logging.info(f'running model for `question_answer_pair_generation`')
+        assert self.is_qag, "`generate_qa_end2end` is available for end2end_qag_model"
+        prefix_type = 'qag' if self.add_prefix else None
+        single_input = type(list_context) is str
+        list_context = [list_context] if single_input else list_context
+        output = self.generate_prediction(
+            list_context, prefix_type=prefix_type, cache_path=cache_path, num_beams=num_beams, batch_size=batch_size
+        )
+
+        def format_qa(list_raw_string):
+            tmp = []
+            for raw_string in list_raw_string:
+                if len(raw_string.split(answer_prefix)) != 2 or question_prefix not in raw_string:
+                    logging.info(f"invalid prediction: {raw_string}")
+                else:
+                    q, a = raw_string.split(answer_prefix)
+                    a = re.sub(r'\A\s+', '', a)
+                    a = re.sub(r'\s+\Z', '', a)
+                    q = q.replace(question_prefix, "")
+                    q = re.sub(r'\A\s+', '', q)
+                    q = re.sub(r'\s+\Z', '', q)
+                    tmp.append((q, a))
+            return tmp
+
+        output = [format_qa(o.split(splitting_symbol)) for o in output]
+        return output[0] if single_input else output
+
+    def generate_qa(self,
+                    list_context: str or List,
+                    batch_size: int = None,
+                    num_beams: int = 4,
+                    cache_path: str = None,
+                    num_questions: int = None,
+                    sentence_level: bool = False):
+        """ Generate question given context.
+
+        @param list_context: Input text.
+        @param batch_size: Batch size.
+        @param num_beams: Number of beam for model generation.
+        @param cache_path: Path to pre-compute features.
+        @param num_questions: Max number of questions.
+        @param sentence_level: Run prediction on each sentence of the context independently to reduce complexity.
+        @return: List of generated sentences.
+        """
+        if self.is_qag:
+            return self.generate_qa_end2end(list_context, batch_size, num_beams, cache_path)
+        single_input = type(list_context) is str
+        list_context = [list_context] if single_input else list_context
+        original_input_length = len(list_context)
+
+        logging.info('running model for `ae`')
+        list_answer = self.generate_a(
+            list_context,
+            batch_size=batch_size,
+            num_beams=num_beams,
+            cache_path=cache_path,
+            sentence_level=sentence_level,
+            num_questions=num_questions
+        )
+        valid_context_id = [n for n, a in enumerate(list_answer) if a is not None]
+        list_context = [list_context[n] for n in valid_context_id]
+        list_answer = [list_answer[n] for n in valid_context_id]
+        qg_input, qg_hl, list_length = [], [], [0]
+        for c, a in zip(list_context, list_answer):
+            qg_hl += a
+            qg_input += [c] * len(a)
+            list_length.append(list_length[-1] + len(a))
+        logging.info('running model for `qg`')
+        list_question = self.generate_q(
+            qg_input,
+            list_answer=qg_hl,
+            batch_size=batch_size,
+            cache_path=cache_path,
+            num_beams=num_beams,
+            sentence_level=sentence_level
+        )
+
+        assert len(qg_hl) == len(list_question), f"{len(qg_input)} != {len(list_question)}"
+
+        # return to nested list
+        list_question = [list_question[list_length[n - 1]:list_length[n]] for n in range(1, len(list_length))]
+        list_answer = [qg_hl[list_length[n - 1]:list_length[n]] for n in range(1, len(list_length))]
+        output_list = [None] * original_input_length
+
+        for n, _id in enumerate(valid_context_id):
+            output_list[_id] = [(q, a) for q, a in zip(list_question[n], list_answer[n])]
+        return output_list[0] if single_input else output_list
+
+    def generate_a(self,
+                   context: str or List,
+                   batch_size: int = None,
+                   num_beams: int = 4,
+                   cache_path: str = None,
+                   sentence_level: bool = False,
+                   num_questions: int = None):
+        """ Generate answers from each sentence.
+
+        @param context: Input text.
+        @param batch_size: Batch size.
+        @param num_beams: Number of beam for model generation.
+        @param cache_path: Path to pre-compute features.
+        @param sentence_level: Run prediction on each sentence of the context independently to reduce complexity.
+        @param num_questions: Max number of questions.
+        @return: List of generated answers.
+        """
+        logging.info(f'running model for `answer_extraction`')
+        if self.answer_model_type == 'spacy':
+            num_questions = 10 if num_questions is None else num_questions
+            if type(context) is str:
+                return self.spacy_module.keyword(context, num_questions)
+            else:
+                return [self.spacy_module.keyword(c, num_questions) for c in context]
+        single_input = type(context) is str
+        context = [context] if single_input else context
+        list_sentences = [self.spacy_module.sentence(c) for c in context]  # split into sentence
+        list_inputs = [[c] * len(s) for c, s in zip(context, list_sentences)]
+        list_length = [0] + np.cumsum([len(s) for s in list_sentences]).tolist()
+        if sentence_level:
+            list_inputs = list_sentences
+        # flatten inputs
+        flat_sentences = list(chain(*list_sentences))
+        flat_inputs = list(chain(*list_inputs))
+        if self.answer_model_type == 'multitask':
+            answer = self.generate_prediction(
+                flat_inputs,  # list_input,
+                highlights=flat_sentences,  # highlights=list_sentence,
+                prefix_type='ae' if self.add_prefix else None,
+                cache_path=cache_path,
+                num_beams=num_beams,
+                batch_size=batch_size
+            )
+        elif self.answer_model_type == 'pipeline':
+            answer = self.generate_prediction(
+                flat_inputs,  # list_input,
+                highlights=flat_sentences,  # highlights=list_sentence,
+                prefix_type='ae' if self.add_prefix_ae else None,
+                cache_path=cache_path,
+                num_beams=num_beams,
+                batch_size=batch_size,
+                switch_to_model_ae=True
+            )
+        else:
+            raise ValueError(f"unknown answer model type: {self.answer_model_type}")
+        # return to nested list
+        answer = [clean(a) for a in answer]
+        list_answer = [answer[list_length[n - 1]:list_length[n]] for n in range(1, len(list_length))]
+        list_answer = [[a for a, c in zip(a_sent, c_sent) if a is not None and a in c]
+                       for a_sent, c_sent in zip(list_answer, list_inputs)]
+        list_answer = [None if len(a) == 0 else a for a in list_answer]
+        if not self.drop_answer_error_text:
+            if any(a is None for a in list_answer):
+                raise AnswerNotFoundError([context[n] for n, a in enumerate(list_answer) if a is None][0])
+        return list_answer[0] if single_input else list_answer
+
+    def generate_q(self,
+                   list_context: str or List,
+                   list_answer: List = None,
+                   batch_size: int = None,
+                   num_beams: int = 4,
+                   cache_path: str = None,
+                   sentence_level: bool = False):
+        """ Generate question from paragraph and answer. Note that `list_answer` is needed unless they are already
+        highlighted in the `list_context`. eg) "I live in <hl> Tokyo <hl>."
+
+        @param list_context: List of input texts.
+        @param list_answer: List of answers in the `list_context` that are highlighted by <hl>.
+        @param batch_size: Batch size.
+        @param num_beams: Number of beam for model generation.
+        @param cache_path: Path to pre-compute features.
+        @param sentence_level: Run prediction on each sentence of the context independently to reduce complexity.
+        @return: List of generated sentences.
+        """
+        assert self.is_qg, "model is not fine-tuned for QG"
+        if list_answer is not None:
+            assert type(list_context) is type(list_answer), f"{type(list_context)} != {type(list_answer)}"
+        single_input = False
+        if type(list_context) is str:
+            list_context = [list_context]
+            list_answer = [list_answer] if list_answer is not None else None
+            single_input = True
+        output = self.generate_prediction(
+            list_context,
+            highlights=list_answer,
+            prefix_type='qg' if self.add_prefix else None,
+            cache_path=cache_path,
+            num_beams=num_beams,
+            batch_size=batch_size,
+            sentence_level=sentence_level
+        )
+        if single_input:
+            return output[0]
+        return output
+
+    def generate_prediction(self,
+                            inputs: List,
+                            highlights: List or None = None,
+                            prefix_type: str = None,
+                            num_beams: int = 4,
+                            batch_size: int = None,
+                            cache_path: str = None,
+                            sentence_level: bool = False,
+                            switch_to_model_ae: bool = False):
+        """ General method to generate model prediction
+
+        @param inputs: List of input sequences.
+        @param highlights: List of sub-sequences from list_context to be highlighted by <hl>.
+        @param batch_size: Batch size.
+        @param num_beams: Number of beam for model generation.
+        @param cache_path: Path to pre-compute features.
+        @param prefix_type: Either of `qg` or `answer_extraction`, which is to add at the beginning of the text.
+        @return: List of generated sequences.
+        """
+        self.eval()
+        if switch_to_model_ae:
+            assert self.model_ae is not None and self.tokenizer_ae is not None
+            model = self.model_ae
+            tokenizer = self.tokenizer_ae
+            max_length_output = self.max_length_output_ae
+        else:
+            model = self.model
+            tokenizer = self.tokenizer
+            max_length_output = self.max_length_output
+
+        if sentence_level:
+            assert highlights is not None, '`sentence_level` needs `highlights`.'
+            assert len(highlights) == len(inputs), str([len(highlights), len(inputs)])
+            list_sentence = []
+            for context, answer in zip(inputs, highlights):
+                s = [sentence for sentence in self.spacy_module.sentence(context) if answer in sentence]
+                list_sentence.append(s[0] if len(s) != 0 else context)
+            inputs = list_sentence
+
+        assert type(inputs) is list, inputs
+        encode_list = self.text_to_encode(
+            inputs,
+            highlights=highlights,
+            prefix_type=prefix_type,
+            cache_path=cache_path,
+            switch_to_model_ae=switch_to_model_ae
+        )
+        loader = self.get_data_loader(encode_list, batch_size=batch_size)
+        outputs = []
+        for encode in loader:
+            with torch.no_grad():
+                if 'labels' in encode:
+                    encode.pop('labels')
+                encode = {k: v.to(self.device) for k, v in encode.items()}
+                encode['max_length'] = max_length_output
+                encode['num_beams'] = num_beams
+                tensor = model.module.generate(**encode) if self.parallel else model.generate(**encode)
+                outputs += tokenizer.batch_decode(tensor, skip_special_tokens=True)
+        return outputs

plms/spacy_module.py

@@ -0,0 +1,72 @@
+import spacy
+
+__all__ = 'SpacyPipeline'
+
+MODELS = {
+    "en": "en_core_web_sm",
+    "ja": "ja_core_news_sm",
+    "zh": "zh_core_web_sm",
+    "de": "de_core_news_sm",
+    "es": "es_core_news_sm",
+    "it": "it_core_news_sm",
+    "ko": "ko_core_news_sm",
+    "ru": "ru_core_news_sm",
+    "fr": "fr_core_news_sm",
+    "vi": "vi_core_news_lg"
+}
+VALID_METHODS = ['positionrank', 'textrank', 'biasedtextrank', 'positionrank', 'ner']
+
+
+class SpacyPipeline:
+
+    def __init__(self, language, algorithm: str = None):
+        model = "vi_core_news_lg" if language not in MODELS else MODELS[language]
+
+        self.nlp = spacy.load(model)
+        self.nlp.add_pipe("sentencizer")
+        self.algorithm = algorithm
+        self.library = None
+        if self.algorithm is not None and self.algorithm != 'ner':
+            assert algorithm in VALID_METHODS, f'invalid algorithm {algorithm}\n- valid list: {VALID_METHODS}'
+            if self.algorithm == 'yake':
+                import spacy_ke  # need to load yake
+                self.nlp.add_pipe("yake")
+                self.library = 'spacy_ke'
+            elif self.algorithm in ['textrank', 'biasedtextrank', 'positionrank']:
+                import pytextrank
+                self.nlp.add_pipe(algorithm)
+                self.library = 'pytextrank'
+            else:
+                raise ValueError(f'unknown algorithm: {self.algorithm}')
+
+    def _get_keyword(self, output, original_document=None, n=None):
+        if self.algorithm == 'ner':
+            return [str(i) for i in output.ents]
+        assert original_document is not None
+        assert n is not None
+        if self.library == 'spacy_ke':
+            return [str(term) for term, score in output._.extract_keywords(n) if str(term) in original_document]
+        return [str(i.text) for i in output._.phrases[:n] if str(i.text) in original_document]
+
+    def sentence_keyword(self, string: str, n: int = 10):
+        out = self.nlp(string)
+        sentence = [str(i) for i in out.sents if len(i) > 0]
+        keyword = self._get_keyword(out, string, n)
+        return sentence, keyword
+
+    def sentence(self, string: str):
+        return [str(i) for i in self.nlp(string).sents if len(i) > 0]
+
+    def token(self, string: str):
+        return [str(i) for i in self.nlp.tokenizer(string)]
+
+    def keyword(self, string: str, n: int = 10):
+        return self._get_keyword(self.nlp(string), string, n)
+
+    def ner(self, string: str, n: int = None):
+        keywords = self.nlp(string).ents
+        return keywords[:min(len(keywords), n)] if n is not None else keywords
+
+    @property
+    def language(self):
+        return self.nlp.lang