def predict(
        self,
        test_dataset,
        per_gpu_batch_size=4,
        max_tgt_length=64,
        beam_size=1,
        need_score_traces=False,
        length_penalty=0,
        forbid_duplicate_ngrams=True,
        forbid_ignore_word=".",
        s2s_config=S2SConfig(),
        num_gpus=None,
        gpu_ids=None,
        local_rank=-1,
        fp16=False,
        verbose=True,
    ):
        """
        Method for predicting, i.e. generating summaries.
        Args:
            test_dataset (S2SAbsSumDataset): Testing dataset.
            per_gpu_batch_size (int, optional): Number of testing samples in each
                batch per GPU. Defaults to 4.
            max_tgt_length (int, optional): Maximum number of tokens in output
                sequence. Defaults to 64.
            beam_size (int, optional): Beam size of beam search. Defaults to 1.
            need_score_traces (bool, optional): Whether to return score traces of
                beam search. Defaults to False.
            length_penalty (float, optional): Length penalty for beam search.
                Defaults to 0.
            forbid_duplicate_ngrams (bool, optional): Whether to forbid duplicate
                n-grams when generating output. Size of the n-gram is determined by
                `S2SConfig.ngram_size` which defaults to 3. Defaults to True.
            forbid_ignore_word (str, optional): Words to ignore when forbidding
                duplicate ngrams. Multiple words should be separated by "|", for
                example, ".|[X_SEP]". Defaults to ".".
            s2s_config (S2SConfig, optional): Some default decoding settings that
                the users usually don't need to change. Defaults to S2SConfig().
            num_gpus (int, optional): Number of GPUs to use. Ignored if `gpu_ids` is
                provided. Defaults to None and all available GPUs are used.
            gpu_ids (list, optional): List of GPU IDs ot use. Defaults to None and GPUs
                used are determined by num_gpus.
            local_rank (int, optional): Rank of the device in distributed training.
                Defaults to -1 which means non-distributed training.
            fp16 (bool, optional): Whether to use 16-bit mixed precision through Apex.
                Defaults to False.
            verbose(bool, optional): Whether to output predicting log. Defaults to True.

        Returns:
            List or tuple of lists: List of generated summaries. If `need_score_traces`
                is True, also returns the score traces of beam search.

        """

        if need_score_traces and beam_size <= 1:
            raise ValueError(
                "Score trace is only available for beam search with beam size > 1."
            )
        if max_tgt_length >= self.max_seq_length - 2:
            raise ValueError("Maximum tgt length exceeds max seq length - 2.")

        # preprocessing pipeline
        if self._model_type == "roberta":
            is_roberta = True
            no_segment_embedding = True
            vocab = self.tokenizer.encoder
        else:
            is_roberta = False
            no_segment_embedding = False
            vocab = self.tokenizer.vocab

        if not self._model_name.startswith("unilm1.2"):
            if self._model_name.startswith("unilm-") or self._model_name.startswith(
                "unilm1-"
            ):
                new_segment_ids = True
            else:
                new_segment_ids = False
        else:
            new_segment_ids = False

        cls_token = "<s>" if is_roberta else "[CLS]"
        sep_token = "</s>" if is_roberta else "[SEP]"
        pad_token = "<pad>" if is_roberta else "[PAD]"
        mask_token = "<mask>" if is_roberta else "[MASK]"

        max_src_length = self.max_seq_length - 2 - max_tgt_length
        tokenizer = self.tokenizer
        bi_uni_pipeline = []
        bi_uni_pipeline.append(
            seq2seq_loader.Preprocess4Seq2seqDecoder(
                list(vocab.keys()),
                tokenizer.convert_tokens_to_ids,
                self.max_seq_length,
                max_tgt_length=max_tgt_length,
                pos_shift=s2s_config.pos_shift,
                source_type_id=self.model_config.source_type_id,
                target_type_id=self.model_config.target_type_id,
                cls_token=tokenizer.cls_token,
                sep_token=tokenizer.sep_token,
                pad_token=tokenizer.pad_token,
            )
        )
        mask_word_id, eos_word_ids, sos_word_id = tokenizer.convert_tokens_to_ids(
        [tokenizer.mask_token, tokenizer.sep_token, tokenizer.sep_token])

        def collate_fn(input_batch):
            buf_id = [x[0] for x in input_batch]
            buf = [x[1][:max_src_length] for x in input_batch]
            max_a_len = max([len(x) for x in buf])
            instances = []
            for instance in [(x, max_a_len) for x in buf]:
                for proc in bi_uni_pipeline:
                    instance = proc(instance)
                instances.append(instance)
            batch = seq2seq_loader.batch_list_to_batch_tensors(instances)

            return (batch, buf_id)

        # prepare decoder
        pair_num_relation = 0
        cls_num_labels = 2
        type_vocab_size = (
            6 + (1 if s2s_config.s2s_add_segment else 0) if new_segment_ids else 2
        )
        forbid_ignore_set = None
        if forbid_ignore_word:
            w_list = []
            for w in forbid_ignore_word.split("|"):
                if w.startswith("[") and w.endswith("]"):
                    w_list.append(w.upper())
                else:
                    w_list.append(w)
            forbid_ignore_set = set(self.tokenizer.convert_tokens_to_ids(w_list))

        if hasattr(self.model, "module"):
            state_dict = self.model.module.state_dict()
        else:
            state_dict = self.model.state_dict()
        
        bert_config = None
        if self._model_type == "minilm":
            bert_config = s2s_ft.modeling_decoding.BertConfig(
                    30522,
                    type_vocab_size=type_vocab_size,
                    hidden_size=384,
                    intermediate_size=1536, 
                    max_position_embeddings=self.max_seq_length,)

        else: # self._model_type == "unilm":
            bert_config = s2s_ft.modeling_decoding.BertConfig(
                    len(list(vocab.keys())) if not is_roberta else 50265,
                    type_vocab_size=type_vocab_size,
                    max_position_embeddings=self.max_seq_length,)

       
        model = BertForSeq2SeqDecoder.from_pretrained(
            self._bert_model_name,
            bert_config,
            state_dict=state_dict,
            cache_dir=self.cache_dir,
            mask_word_id=mask_word_id,
            search_beam_size=beam_size,
            length_penalty=length_penalty,
            eos_id=eos_word_ids,
            sos_id=sos_word_id,
            forbid_duplicate_ngrams=forbid_duplicate_ngrams,
            forbid_ignore_set=forbid_ignore_set,
            ngram_size=s2s_config.forbid_ngram_size,
            min_len=s2s_config.min_len,
            mode=s2s_config.mode,
            max_position_embeddings=self.max_seq_length,
            pos_shift=s2s_config.pos_shift,
        )

        del state_dict

        if fp16:
            model.half()
        # get device
        device, num_gpus = get_device(
            num_gpus=num_gpus, gpu_ids=gpu_ids, local_rank=local_rank
        )

        # # move model
        model = move_model_to_device(model=model, device=device)

        batch_size = per_gpu_batch_size * max(1, num_gpus)

        model = parallelize_model(
            model=model,
            device=device,
            num_gpus=num_gpus,
            gpu_ids=gpu_ids,
            local_rank=local_rank,
        )

        # torch.cuda.empty_cache()
        model.eval()
        first_batch = True
        batch_count = 0

        output_lines = [""] * len(test_dataset)
        score_trace_list = [None] * len(test_dataset)

        test_sampler = SequentialSampler(test_dataset)
        test_dataloader = DataLoader(
            test_dataset,
            sampler=test_sampler,
            batch_size=batch_size,
            collate_fn=collate_fn,
        )
        for batch, buf_id in tqdm(
            test_dataloader, desc="Evaluating", disable=not verbose
        ):
            batch_count += 1
            with torch.no_grad():
                batch = [t.to(device) if t is not None else None for t in batch]
                (
                    input_ids,
                    token_type_ids,
                    position_ids,
                    input_mask,
                    mask_qkv,
                    task_idx,
                ) = batch
                traces = model(
                    input_ids,
                    token_type_ids,
                    position_ids,
                    input_mask,
                    task_idx=task_idx,
                    mask_qkv=mask_qkv,
                )
                if beam_size > 1:
                    traces = {k: v.tolist() for k, v in traces.items()}
                    output_ids = traces["pred_seq"]
                else:
                    output_ids = traces.tolist()

                for i in range(len(batch[0])):
                    w_ids = output_ids[i]
                    output_buf = self.tokenizer.convert_ids_to_tokens(w_ids)
                    output_tokens = []
                    for t in output_buf:
                        if t in (sep_token, pad_token):
                            break
                        output_tokens.append(t)
                    if is_roberta:
                        output_sequence = self.tokenizer.convert_tokens_to_string(
                            output_tokens
                        )
                    else:
                        output_sequence = " ".join(detokenize(output_tokens))
                    if "\n" in output_sequence:
                        output_sequence = " [X_SEP] ".join(output_sequence.split("\n"))
                    output_lines[buf_id[i]] = output_sequence
                    if first_batch or batch_count % 50 == 0:
                        logger.info("{} = {}".format(buf_id[i], output_sequence))
                    if need_score_traces:
                        score_trace_list[buf_id[i]] = {
                            "scores": traces["scores"][i],
                            "wids": traces["wids"][i],
                            "ptrs": traces["ptrs"][i],
                        }

            first_batch = False

        del model
        del batch
        torch.cuda.empty_cache()

        if need_score_traces:
            return output_lines, score_trace_list
        else:
            return output_lines

        Method for predicting, i.e. generating summaries.
        Args:
            test_dataset (S2SAbsSumDataset): Testing dataset.
            per_gpu_batch_size (int, optional): Number of testing samples in each
                batch per GPU. Defaults to 4.
            max_tgt_length (int, optional): Maximum number of tokens in output
                sequence. Defaults to 64.
            beam_size (int, optional): Beam size of beam search. Defaults to 1.
            need_score_traces (bool, optional): Whether to return score traces of
                beam search. Defaults to False.
            length_penalty (float, optional): Length penalty for beam search.
                Defaults to 0.
            forbid_duplicate_ngrams (bool, optional): Whether to forbid duplicate
                n-grams when generating output. Size of the n-gram is determined by
                `S2SConfig.ngram_size` which defaults to 3. Defaults to True.
            forbid_ignore_word (str, optional): Words to ignore when forbidding
                duplicate ngrams. Multiple words should be separated by "|", for
                example, ".|[X_SEP]". Defaults to ".".
            s2s_config (S2SConfig, optional): Some default decoding settings that
                the users usually don't need to change. Defaults to S2SConfig().
            num_gpus (int, optional): Number of GPUs to use. Ignored if `gpu_ids` is
                provided. Defaults to None and all available GPUs are used.
            gpu_ids (list, optional): List of GPU IDs ot use. Defaults to None and GPUs
                used are determined by num_gpus.
            local_rank (int, optional): Rank of the device in distributed training.
                Defaults to -1 which means non-distributed training.
            fp16 (bool, optional): Whether to use 16-bit mixed precision through Apex.
                Defaults to False.
            verbose(bool, optional): Whether to output predicting log. Defaults to True.

        Returns:
            List or tuple of lists: List of generated summaries. If `need_score_traces`
                is True, also returns the score traces of beam search.

        
utils_nlp/models/transformers/abstractive_summarization_seq2seq.py