Method predict_batch

fast_bert/learner_cls copy.py:637–692

(self, texts=None, verbose=True)

    ### Return Predictions ###
    def predict_batch(self, texts=None, verbose=True):

        if verbose:
            if self.logger is None:
                self.logger = logging.getLogger(__name__)
        if texts:
            if verbose:
                self.logger.info("---PROGRESS-STATUS---: Tokenizing input texts...")
            dl = self.data.get_dl_from_texts(texts)
            if verbose:
                self.logger.info("---PROGRESS-STATUS---: Tokenizing input texts...DONE")
        elif self.data.test_dl:
            dl = self.data.test_dl
        else:
            dl = self.data.val_dl

        all_logits = None

        self.model.eval()
        for step, batch in enumerate(dl):
            if verbose:
                self.logger.info(
                    "---PROGRESS-STATUS---: Predicting batch {}/{}".format(
                        step + 1, len(dl)
                    )
                )
            batch = tuple(t.to(self.device) for t in batch)

            inputs = {"input_ids": batch[0], "attention_mask": batch[1], "labels": None}

            if self.model_type in ["bert", "xlnet"]:
                inputs["token_type_ids"] = batch[2]

            with torch.no_grad():
                outputs = self.model(**inputs)
                logits = outputs[0]
                if self.multi_label:
                    logits = logits.sigmoid()
                # elif len(self.data.labels) == 2:
                #     logits = logits.sigmoid()
                else:
                    logits = logits.softmax(dim=1)

            if all_logits is None:
                all_logits = logits.detach().cpu().numpy()
            else:
                all_logits = np.concatenate(
                    (all_logits, logits.detach().cpu().numpy()), axis=0
                )

        result_df = pd.DataFrame(all_logits, columns=self.data.labels)
        results = result_df.to_dict(orient="records")

        if verbose:
            self.logger.info("---PROGRESS-STATUS---: Predicting batch...DONE")
        return [sorted(x.items(), key=lambda kv: kv[1], reverse=True) for x in results]

    # Begin code for LR Finder