def _get_probs(self, data_type):
        """Returns sampling probabilities for different sampling modes - plain, uniform or anneal"""
        if self.sampling_mode == 'uniform':
            sizes = [1 for _ in self._get_data_size(data_type)]
            # as we sample uniformly
            s = sum(sizes)
            probs = [p / s for p in sizes]
        elif self.sampling_mode == 'plain':
            sizes = self._get_data_size(data_type)
            n_samples = sum(sizes)
            probs = [p / n_samples for p in sizes]
        elif self.sampling_mode == 'anneal':
            alpha = 1.0 - 0.8 * (self.epochs_done / self.num_train_epochs)
            annealed_sizes = [p ** alpha for p in self._get_data_size(data_type)]
            n_samples = sum(annealed_sizes)
            probs = [p / n_samples for p in annealed_sizes]
        else:
            raise ValueError(f'Unsupported sampling mode {self.sampling_mode}')
        return probs
Returns sampling probabilities for different sampling modes - plain, uniform or anneal
deeppavlov/dataset_iterators/multitask_iterator.py