How to use the snorkel.augmentation.RandomPolicy function in snorkel

def test_tf_applier_returns_none_generator(self):
        df = self._get_x_df()
        policy = RandomPolicy(
            1, sequence_length=2, n_per_original=2, keep_original=True
        )
        applier = PandasTFApplier([square_returns_none], policy)
        gen = applier.apply_generator(df, batch_size=2)
        df_expected = [
            make_df([1, 1, 1, 2], [0, 0, 0, 1]),
            make_df([3, 81, 81], [2, 2, 2]),
        ]
        for df_batch, df_batch_expected in zip(gen, df_expected):
            pd.testing.assert_frame_equal(df_batch, df_batch_expected)
        pd.testing.assert_frame_equal(df, self._get_x_df())

def test_tf_applier_pandas_generator(self):
        df = self._get_x_df_with_str()
        policy = RandomPolicy(
            1, sequence_length=2, n_per_original=2, keep_original=False
        )
        applier = PandasTFApplier([square], policy)
        gen = applier.apply_generator(df, batch_size=2)
        df_expected = [
            pd.DataFrame(
                {"num": [1, 1, 16, 16], "strs": ["x", "x", "y", "y"]},
                index=[0, 0, 1, 1],
            ),
            pd.DataFrame({"num": [81, 81], "strs": ["z", "z"]}, index=[2, 2]),
        ]
        for df_batch, df_batch_expected in zip(gen, df_expected):
            self.assertEqual(df_batch.num.dtype, "int64")
            pd.testing.assert_frame_equal(df_batch, df_batch_expected)
        pd.testing.assert_frame_equal(df, self._get_x_df_with_str())

def test_tf_applier_keep_original_generator(self) -> None:
        data = self._get_x_namespace()
        policy = RandomPolicy(
            1, sequence_length=2, n_per_original=2, keep_original=True
        )
        applier = TFApplier([square], policy)
        batches_expected = [[1, 1, 1, 2, 16, 16], [3, 81, 81]]
        gen = applier.apply_generator(data, batch_size=2)
        for batch, batch_expected in zip(gen, batches_expected):
            self.assertEqual(batch, self._get_x_namespace(batch_expected))
        self.assertEqual(data, self._get_x_namespace())

def test_tf_applier_returns_none(self) -> None:
        data = self._get_x_namespace()
        policy = RandomPolicy(
            1, sequence_length=2, n_per_original=2, keep_original=True
        )
        applier = TFApplier([square_returns_none], policy)
        data_augmented = applier.apply(data, progress_bar=False)
        vals = [1, 1, 1, 2, 3, 81, 81]
        self.assertEqual(data_augmented, self._get_x_namespace(vals))
        self.assertEqual(data, self._get_x_namespace())

def test_tf_applier_keep_original(self) -> None:
        data = self._get_x_namespace()
        policy = RandomPolicy(
            1, sequence_length=2, n_per_original=2, keep_original=True
        )
        applier = TFApplier([square], policy)
        data_augmented = applier.apply(data, progress_bar=False)
        vals = [1, 1, 1, 2, 16, 16, 3, 81, 81]
        self.assertEqual(data_augmented, self._get_x_namespace(vals))
        self.assertEqual(data, self._get_x_namespace())

def test_tf_applier_returns_none(self):
        df = self._get_x_df()
        policy = RandomPolicy(
            1, sequence_length=2, n_per_original=2, keep_original=True
        )
        applier = PandasTFApplier([square_returns_none], policy)
        df_augmented = applier.apply(df, progress_bar=False)
        df_expected = pd.DataFrame(
            dict(num=[1, 1, 1, 2, 3, 81, 81]), index=[0, 0, 0, 1, 2, 2, 2]
        )
        self.assertEqual(df_augmented.num.dtype, "int64")
        pd.testing.assert_frame_equal(df_augmented, df_expected)
        pd.testing.assert_frame_equal(df, self._get_x_df())

def test_tf_applier(self) -> None:
        data = self._get_x_namespace()
        policy = RandomPolicy(
            1, sequence_length=2, n_per_original=1, keep_original=False
        )
        applier = TFApplier([square], policy)
        data_augmented = applier.apply(data, progress_bar=False)
        self.assertEqual(data_augmented, self._get_x_namespace([1, 16, 81]))
        self.assertEqual(data, self._get_x_namespace())

        data_augmented = applier.apply(data, progress_bar=True)
        self.assertEqual(data_augmented, self._get_x_namespace([1, 16, 81]))
        self.assertEqual(data, self._get_x_namespace())

def test_tf_applier_pandas_keep_original(self):
        df = self._get_x_df()
        policy = RandomPolicy(
            1, sequence_length=2, n_per_original=2, keep_original=True
        )
        applier = PandasTFApplier([square], policy)
        df_augmented = applier.apply(df, progress_bar=False)
        df_expected = pd.DataFrame(
            dict(num=[1, 1, 1, 2, 16, 16, 3, 81, 81]), index=[0, 0, 0, 1, 1, 1, 2, 2, 2]
        )
        self.assertEqual(df_augmented.num.dtype, "int64")
        pd.testing.assert_frame_equal(df_augmented, df_expected)
        pd.testing.assert_frame_equal(df, self._get_x_df())

def test_tf_applier_returns_none_generator(self) -> None:
        data = self._get_x_namespace()
        policy = RandomPolicy(
            1, sequence_length=2, n_per_original=2, keep_original=True
        )
        applier = TFApplier([square_returns_none], policy)
        batches_expected = [[1, 1, 1, 2], [3, 81, 81]]
        gen = applier.apply_generator(data, batch_size=2)
        for batch, batch_expected in zip(gen, batches_expected):
            self.assertEqual(batch, self._get_x_namespace(batch_expected))
        self.assertEqual(data, self._get_x_namespace())

def test_random_policy(self):
        policy = RandomPolicy(2, sequence_length=2)
        n_samples = 100
        samples = [policy.generate() for _ in range(n_samples)]
        a_ct = samples.count([0, 0])
        b_ct = samples.count([0, 1])
        c_ct = samples.count([1, 0])
        d_ct = samples.count([1, 1])
        self.assertGreater(a_ct, 0)
        self.assertGreater(b_ct, 0)
        self.assertGreater(c_ct, 0)
        self.assertGreater(d_ct, 0)
        self.assertEqual(a_ct + b_ct + c_ct + d_ct, n_samples)