def __init__(
        self,
        labelled: Optional[DataModule] = None,
        heuristic: "AbstractHeuristic" = BALD(),
        map_dataset_to_labelled: Optional[Callable] = dataset_to_non_labelled_tensor,
        filter_unlabelled_data: Optional[Callable] = filter_unlabelled_data,
        initial_num_labels: Optional[int] = None,
        query_size: int = 1,
        val_split: Optional[float] = None,
    ):
        """The `ActiveLearningDataModule` handles data manipulation for ActiveLearning.

        Args:
            labelled: DataModule containing labelled train data for research use-case.
                The labelled data would be masked.
            heuristic: Sorting algorithm used to rank samples on how likely they can help with model performance.
            map_dataset_to_labelled: Function used to emulate masking on labelled dataset.
            filter_unlabelled_data: Function used to filter the unlabelled data while computing uncertainties.
            initial_num_labels: Number of samples to randomly label to start the training with.
            query_size: Number of samples to be labelled at each Active Learning loop based on the fed heuristic.
            val_split: Float to split train dataset into train and validation set.

        """
        super().__init__(batch_size=1)
        self.labelled = labelled
        self.heuristic = heuristic
        self.map_dataset_to_labelled = map_dataset_to_labelled
        self.filter_unlabelled_data = filter_unlabelled_data
        self.initial_num_labels = initial_num_labels
        self.query_size = query_size
        self.val_split = val_split
        self._dataset: Optional[ActiveLearningDataset] = None

        if not self.labelled:
            raise TypeError("The labelled `datamodule` should be provided.")

        if not self.labelled.num_classes:
            raise TypeError("The labelled dataset should be labelled")

        if self.labelled and (self.labelled._val_input or self.labelled._predict_input):
            raise TypeError("The labelled `datamodule` should have only train data.")

        self._dataset = ActiveLearningDataset(
            self.labelled._train_input, labelled=self.map_dataset_to_labelled(self.labelled._train_input)
        )

        if not self.val_split or not self.has_labelled_data:
            self.val_dataloader = None
        elif self.val_split < 0 or self.val_split > 1:
            raise ValueError("The `val_split` should a float between 0 and 1.")

        if self.labelled._test_input:
            self.test_dataloader = self._test_dataloader

        if hasattr(self.labelled, "on_after_batch_transfer"):
            self.on_after_batch_transfer = self.labelled.on_after_batch_transfer

        if not self.initial_num_labels:
            warnings.warn(
                "No labels provided for the initial step," "the estimated uncertainties are unreliable!", UserWarning
            )
        else:
            self._dataset.label_randomly(self.initial_num_labels)
The `ActiveLearningDataModule` handles data manipulation for ActiveLearning.

        Args:
            labelled: DataModule containing labelled train data for research use-case.
                The labelled data would be masked.
            heuristic: Sorting algorithm used to rank samples on how likely they can help with model performance.
            map_dataset_to_labelled: Function used to emulate masking on labelled dataset.
            filter_unlabelled_data: Function used to filter the unlabelled data while computing uncertainties.
            initial_num_labels: Number of samples to randomly label to start the training with.
            query_size: Number of samples to be labelled at each Active Learning loop based on the fed heuristic.
            val_split: Float to split train dataset into train and validation set.

        
src/flash/image/classification/integrations/baal/data.py