text: import copy
  
  import pytest
  import numpy as np
  
  from rasa.utils.tensorflow.model_data import RasaModelData
  
  
  def test_shuffle_session_data(model_data: RasaModelData):
      before = copy.copy(model_data)
  
      # precondition
      assert np.all(
          np.array(list(before.values())) == np.array(list(model_data.values()))
      )
  
      data = model_data.shuffled_data(model_data.data)
  
      # check that original data didn't change
      assert np.all(
          np.array(list(before.values())) == np.array(list(model_data.values()))
      )
      # check that new data is different
      assert np.all(np.array(model_data.values()) != np.array(data.values()))
  
  
  def test_split_data_by_label(model_data: RasaModelData):
      split_model_data = model_data._split_by_label_ids(
          model_data.data, model_data.get("label", "ids")[0], np.array([0, 1])
      )
  
      assert len(split_model_data) == 2
      for s in split_model_data:
          assert len(set(s.get("label", "ids")[0])) == 1
  
      for key, attribute_data in split_model_data[0].items():
          for sub_key, features in attribute_data.items():
              assert len(features) == len(model_data.data[key][sub_key])
              assert len(features[0]) == 2
  
  
  def test_split_data_by_none_label(model_data: RasaModelData):
      model_data.label_key = None
      model_data.label_sub_key = None
  
      split_model_data = model_data.split(2, 42)
  
      assert len(split_model_data) == 2
  
      train_data = split_model_data[0]
      test_data = split_model_data[1]
  
      # train data should have 3 examples
      assert len(train_data.get("label", "ids")[0]) == 3
      # test data should have 2 examples
      assert len(test_data.get("label", "ids")[0]) == 2
  
  
  def test_train_val_split(model_data: RasaModelData):
      train_model_data, test_model_data = model_data.split(2, 42)
  
      for key, values in model_data.items():
          assert len(values) == len(train_model_data.get(key))
          assert len(values) == len(test_model_data.get(key))
          for sub_key, data in values.items():
              assert len(data) == len(train_model_data.get(key, sub_key))
              assert len(data) == len(test_model_data.get(key, sub_key))
              for i, v in enumerate(data):
                  if isinstance(v[0], list):
                      assert (
                          v[0][0].dtype
                          == train_model_data.get(key, sub_key)[i][0][0].dtype
                      )
                  else:
                      assert v[0].dtype == train_model_data.get(key, sub_key)[i][0].dtype
  
      for values in train_model_data.values():
          for data in values.values():
              for v in data:
                  assert np.array(v).shape[0] == 3
  
      for values in test_model_data.values():
          for data in values.values():
              for v in data:
                  assert np.array(v).shape[0] == 2
  
  
  @pytest.mark.parametrize("size", [0, 1, 5])
  def test_train_val_split_incorrect_size(model_data: RasaModelData, size: int):
      with pytest.raises(ValueError):
          model_data.split(size, 42)
  
  
  def test_session_data_for_ids(model_data: RasaModelData):
      filtered_data = model_data._data_for_ids(model_data.data, np.array([0, 1]))
  
      for values in filtered_data.values():
          for data in values.values():
              for v in data:
                  assert np.array(v).shape[0] == 2
  
      key = model_data.keys()[0]
      sub_key = model_data.keys(key)[0]
  
      assert np.all(
          np.array(filtered_data[key][sub_key][0][0])
          == np.array(model_data.get(key, sub_key)[0][0])
      )
      assert np.all(
          np.array(filtered_data[key][sub_key][0][1])
          == np.array(model_data.get(key, sub_key)[0][1])
      )
  
  
  def test_get_number_of_examples(model_data: RasaModelData):
      assert model_data.number_of_examples() == 5
  
  
  def test_get_number_of_examples_raises_value_error(model_data: RasaModelData):
      model_data.data["dense"] = {}
      model_data.data["dense"]["data"] = [np.random.randint(5, size=(2, 10))]
      with pytest.raises(ValueError):
          model_data.number_of_examples()
  
  
  def test_is_in_4d_format(model_data: RasaModelData):
      assert model_data.data["action_text"]["sequence"][0].number_of_dimensions == 4
      assert model_data.data["text"]["sentence"][0].number_of_dimensions == 3
  
  
  def test_balance_model_data(model_data: RasaModelData):
      data = model_data.balanced_data(model_data.data, 2, False)
  
      assert np.all(np.array(data["label"]["ids"][0]) == np.array([0, 1, 1, 0, 1]))
  
  
  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_get_num_of_features(model_data: RasaModelData):
      num_features = model_data.number_of_units("text", "sentence")
  
      assert num_features == 24
  
  
  def test_sort(model_data: RasaModelData):
      assert list(model_data.data.keys()) == [
          "text",
          "action_text",
          "dialogue",
          "label",
          "entities",
      ]
  
      model_data.sort()
  
      assert list(model_data.data.keys()) == [
          "action_text",
          "dialogue",
          "entities",
          "label",
          "text",
      ]
  
  
  def test_update_key(model_data: RasaModelData):
      assert model_data.does_feature_exist("label", "ids")
  
      model_data.update_key("label", "ids", "intent", "ids")
  
      assert not model_data.does_feature_exist("label", "ids")
      assert model_data.does_feature_exist("intent", "ids")
      assert "label" not in model_data.data
  