def __call__(self, features_batch: List[List[InputFeatures]]) -> Tuple[
        List[List[int]], List[List[int]], List[List[float]], List[List[float]], List[int]]:
        """get predictions using features as input

        Args:
            features_batch: batch of InputFeatures instances

        Returns:
            start_pred_batch: answer start positions
            end_pred_batch: answer end positions
            logits_batch: answer logits
            scores_batch: answer confidences
            ind_batch: indices of paragraph pieces where the answer was found

        """
        predictions = {}
        # TODO: refactor batchification
        indices, input_ids, input_masks, input_type_ids = [], [], [], []
        for n, features_list in enumerate(features_batch):
            for f in features_list:
                input_ids.append(f.input_ids)
                input_masks.append(f.attention_mask)
                input_type_ids.append(f.token_type_ids)
                indices.append(n)

        num_batches = len(indices) // self.batch_size + int(len(indices) % self.batch_size > 0)
        for i in range(num_batches):
            b_input_ids = torch.cat(input_ids[i * self.batch_size:(i + 1) * self.batch_size], dim=0).to(self.device)
            b_input_masks = torch.cat(input_masks[i * self.batch_size:(i + 1) * self.batch_size], dim=0).to(self.device)
            b_input_type_ids = torch.cat(input_type_ids[i * self.batch_size:(i + 1) * self.batch_size],
                                         dim=0).to(self.device)
            input_ = {
                'input_ids': b_input_ids,
                'attention_mask': b_input_masks,
                'token_type_ids': b_input_type_ids,
                'return_dict': True
            }

            with torch.no_grad():
                input_ = {arg_name: arg_value for arg_name, arg_value in input_.items()
                          if arg_name in self.accepted_keys}
                # Forward pass, calculate logit predictions
                outputs = self.model(**input_)

                logits_st = outputs.start_logits
                logits_end = outputs.end_logits

                bs = b_input_ids.size()[0]
                seq_len = b_input_ids.size()[-1]
                mask = torch.cat([torch.ones(bs, 1, dtype=torch.int32),
                                  torch.zeros(bs, seq_len - 1, dtype=torch.int32)], dim=-1).to(self.device)
                logit_mask = b_input_type_ids + mask
                logits_st = softmax_mask(logits_st, logit_mask)
                logits_end = softmax_mask(logits_end, logit_mask)

                start_probs = torch.nn.functional.softmax(logits_st, dim=-1)
                end_probs = torch.nn.functional.softmax(logits_end, dim=-1)
                if self.psg_cls:
                    scores = outputs.rank_logits.squeeze(1)
                else:
                    scores = torch.tensor(1) - start_probs[:, 0] * end_probs[:, 0]

                outer = torch.matmul(start_probs.view(*start_probs.size(), 1),
                                     end_probs.view(end_probs.size()[0], 1, end_probs.size()[1]))
                outer_logits = torch.exp(logits_st.view(*logits_st.size(), 1) + logits_end.view(
                    logits_end.size()[0], 1, logits_end.size()[1]))

                context_max_len = torch.max(torch.sum(b_input_type_ids, dim=1)).to(torch.int64)

                max_ans_length = torch.min(torch.tensor(20).to(self.device), context_max_len).to(torch.int64).item()

                outer = torch.triu(outer, diagonal=0) - torch.triu(outer, diagonal=outer.size()[1] - max_ans_length)
                outer_logits = torch.triu(outer_logits, diagonal=0) - torch.triu(
                    outer_logits, diagonal=outer_logits.size()[1] - max_ans_length)

                start_pred = torch.argmax(torch.max(outer, dim=2)[0], dim=1)
                end_pred = torch.argmax(torch.max(outer, dim=1)[0], dim=1)
                logits = torch.max(torch.max(outer_logits, dim=2)[0], dim=1)[0]

            # Move logits and labels to CPU and to numpy arrays
            start_pred = start_pred.detach().cpu().numpy()
            end_pred = end_pred.detach().cpu().numpy()
            logits = logits.detach().cpu().numpy().tolist()
            scores = scores.detach().cpu().numpy().tolist()

            for j, (start_pred_elem, end_pred_elem, logits_elem, scores_elem) in \
                    enumerate(zip(start_pred, end_pred, logits, scores)):
                ind = indices[i * self.batch_size + j]
                if ind in predictions:
                    predictions[ind] += [(start_pred_elem, end_pred_elem, logits_elem, scores_elem)]
                else:
                    predictions[ind] = [(start_pred_elem, end_pred_elem, logits_elem, scores_elem)]

        start_pred_batch, end_pred_batch, logits_batch, scores_batch, ind_batch = [], [], [], [], []
        for ind in sorted(predictions.keys()):
            prediction = predictions[ind]
            max_ind = np.argmax([pred[2] for pred in prediction])
            start_pred_batch.append(prediction[max_ind][0])
            end_pred_batch.append(prediction[max_ind][1])
            logits_batch.append(prediction[max_ind][2])
            scores_batch.append(prediction[max_ind][3])
            ind_batch.append(max_ind)

        return start_pred_batch, end_pred_batch, logits_batch, scores_batch, ind_batch
get predictions using features as input

        Args:
            features_batch: batch of InputFeatures instances

        Returns:
            start_pred_batch: answer start positions
            end_pred_batch: answer end positions
            logits_batch: answer logits
            scores_batch: answer confidences
            ind_batch: indices of paragraph pieces where the answer was found

        
deeppavlov/models/torch_bert/torch_transformers_squad.py