Method predict

LanguageNetwork/BERT/models/trainer.py:307–418

Validate models. valid_iter: validate data iterator Returns: :obj:`nmt.Statistics`: validation loss statistics

(self, test_iter, step, cal_lead=False, cal_oracle=False)

        return stats

    def predict(self, test_iter, step, cal_lead=False, cal_oracle=False):
        """ Validate models.
            valid_iter: validate data iterator
        Returns:
            :obj:`nmt.Statistics`: validation loss statistics
        """

        # Set models in validating mode.
        def _get_ngrams(n, text):
            ngram_set = set()
            text_length = len(text)
            max_index_ngram_start = text_length - n
            for i in range(max_index_ngram_start + 1):
                ngram_set.add(tuple(text[i:i + n]))
            return ngram_set

        def _block_tri(c, p):
            tri_c = _get_ngrams(3, c.split())
            for s in p:
                tri_s = _get_ngrams(3, s.split())
                if len(tri_c.intersection(tri_s)) > 0:
                    return True
            return False

        if not cal_lead and not cal_oracle:
            self.model.eval()
        stats = Statistics()

        can_path = '%s_step%d.candidate' % (self.args.result_path + self.args.data_name, step)
        gold_path = '%s_step%d.gold' % (self.args.result_path + self.args.data_name, step)
        origin_path = '%s_step%d.origin' % (self.args.result_path + self.args.data_name, step)
        with open(can_path, 'w', encoding='utf-8') as save_pred:
            with open(gold_path, 'w', encoding='utf-8') as save_gold:
                with torch.no_grad():
                    origin = []
                    for batch in test_iter:

                        src = batch.src  # 7 sentences
                        # logger.info('origin sent: %s' % len(batch.src_str))  # 7 sentences

                        labels = batch.labels
                        segs = batch.segs
                        clss = batch.clss
                        mask = batch.mask
                        mask_cls = batch.mask_cls

                        gold = []
                        pred = []

                        if cal_lead:
                            selected_ids = [list(range(batch.clss.size(1)))] * batch.batch_size
                        elif cal_oracle:
                            selected_ids = [[j for j in range(batch.clss.size(1)) if labels[i][j] == 1] for i in
                                            range(batch.batch_size)]
                        else:
                            sent_scores, mask = self.model(src, segs, clss, mask, mask_cls)

                            loss = self.loss(sent_scores, labels.float())