async def test_gen_batch(self, trained_policy, default_domain):
training_trackers = await train_trackers(default_domain,
augmentation_factor=0)
training_data = trained_policy.featurize_for_training(
training_trackers, default_domain)
session_data = trained_policy._create_session_data(
training_data.X, training_data.y)
batch_size = 2
batch_x, batch_y, _ = next(
train_utils.gen_batch(session_data=session_data,
batch_size=batch_size,
label_key="action_ids"))
assert batch_x.shape[0] == batch_size and batch_y.shape[0] == batch_size
assert (
batch_x[0].shape == session_data["dialogue_features"][0][0].shape
and batch_y[0].shape == session_data["bot_features"][0][0].shape)
batch_x, batch_y, _ = next(
train_utils.gen_batch(
session_data=session_data,
batch_size=batch_size,
label_key="action_ids",
batch_strategy="balanced",
shuffle=True,
))
assert batch_x.shape[0] == batch_size and batch_y.shape[0] == batch_size
assert (
batch_x[0].shape == session_data["dialogue_features"][0][0].shape
and batch_y[0].shape == session_data["bot_features"][0][0].shape)
def test_gen_batch(session_data: SessionDataType):
iterator = gen_batch(
session_data, 2, "intent_ids", shuffle=True, batch_strategy="balanced"
)
batch = next(iterator)
assert len(batch) == 7
assert len(batch[0]) == 2
batch = next(iterator)
assert len(batch) == 7
assert len(batch[0]) == 2
batch = next(iterator)
assert len(batch) == 7
assert len(batch[0]) == 1
with pytest.raises(StopIteration):
next(iterator)