def test_not_balance_model_data(model_data: RasaModelData):
test_model_data = RasaModelData(label_key="entities",
label_sub_key="tag_ids",
data=model_data.data)
data = test_model_data._balanced_data(test_model_data.data, 2, False)
assert np.all(data["entities"]["tag_ids"] == test_model_data.get(
"entities", "tag_ids"))
def test_balance_model_data(model_data: RasaModelData):
data = model_data._balanced_data(model_data.data, 2, False)
assert np.all(data["intent"]["ids"][0] == np.array([0, 1, 1, 0, 1]))
def test_not_balance_model_data(model_data: RasaModelData):
test_model_data = RasaModelData(label_key="tag_ids", data=model_data.data)
data = test_model_data._balanced_data(test_model_data.data, 2, False)
assert np.all(data.get("tag_ids") == test_model_data.get("tag_ids"))
