fix(kdp): fixing tests fromatting

piotrlaczkowski · piotrlaczkowski · commit 955ed0885a15 · 2025-01-10T22:51:44.000+01:00
diff --git a/test/test_custom_layers.py b/test/test_custom_layers.py
@@ -4,7 +4,7 @@
 import pytest
 import tensorflow as tf
 
-from kdp.custom_layers import TabularAttention, MultiResolutionTabularAttention
+from kdp.custom_layers import MultiResolutionTabularAttention, TabularAttention
 from kdp.layers_factory import PreprocessorLayerFactory
 
 
@@ -18,16 +18,11 @@ def test_tabular_attention_layer_init():
 
 def test_tabular_attention_layer_config():
     """Test get_config and from_config methods."""
-    original_layer = TabularAttention(
-        num_heads=4,
-        d_model=64,
-        dropout_rate=0.2,
-        name="test_attention"
-    )
-    
+    original_layer = TabularAttention(num_heads=4, d_model=64, dropout_rate=0.2, name="test_attention")
+
     config = original_layer.get_config()
     restored_layer = TabularAttention.from_config(config)
-    
+
     assert restored_layer.num_heads == original_layer.num_heads
     assert restored_layer.d_model == original_layer.d_model
     assert restored_layer.dropout_rate == original_layer.dropout_rate
@@ -40,29 +35,26 @@ def test_tabular_attention_computation():
     num_samples = 10
     num_features = 8
     d_model = 16
-    
+
     # Create a layer instance
     layer = TabularAttention(num_heads=2, d_model=d_model)
-    
+
     # Create input data
     inputs = tf.random.normal((batch_size, num_samples, num_features))
-    
+
     # Call the layer
     outputs = layer(inputs, training=True)
-    
+
     # Check output shape - output will have d_model dimension
     assert outputs.shape == (batch_size, num_samples, d_model)
 
 
 def test_tabular_attention_factory():
     """Test creation of TabularAttention layer through PreprocessorLayerFactory."""
     layer = PreprocessorLayerFactory.tabular_attention_layer(
-        num_heads=4,
-        d_model=64,
-        name="test_attention",
-        dropout_rate=0.2
+        num_heads=4, d_model=64, name="test_attention", dropout_rate=0.2
     )
-    
+
     assert isinstance(layer, TabularAttention)
     assert layer.num_heads == 4
     assert layer.d_model == 64
@@ -75,30 +67,30 @@ def test_tabular_attention_training():
     batch_size = 16
     num_samples = 5
     num_features = 4
-    
+
     layer = TabularAttention(num_heads=2, d_model=8, dropout_rate=0.5)
     inputs = tf.random.normal((batch_size, num_samples, num_features))
-    
+
     # Test in training mode
     outputs_training = layer(inputs, training=True)
-    
+
     # Test in inference mode
     outputs_inference = layer(inputs, training=False)
-    
+
     # The outputs should be different due to dropout
     assert not np.allclose(outputs_training.numpy(), outputs_inference.numpy())
 
 
 def test_tabular_attention_invalid_inputs():
     """Test TabularAttention layer with invalid inputs."""
     layer = TabularAttention(num_heads=2, d_model=8)
-    
+
     # Test with wrong input shape
     with pytest.raises(ValueError, match="Input tensor must be 3-dimensional"):
         # Missing batch dimension
         inputs = tf.random.normal((5, 4))
         layer(inputs)
-    
+
     with pytest.raises(ValueError):
         # Wrong rank
         inputs = tf.random.normal((16, 5, 4, 2))
@@ -110,35 +102,31 @@ def test_tabular_attention_end_to_end():
     batch_size = 16
     num_samples = 5
     num_features = 4
-    
+
     # Create a simple model with TabularAttention
     inputs = tf.keras.Input(shape=(num_samples, num_features))
     x = TabularAttention(num_heads=2, d_model=8)(inputs)
     outputs = tf.keras.layers.Dense(1)(x)
     model = tf.keras.Model(inputs=inputs, outputs=outputs)
-    
+
     # Compile the model
-    model.compile(optimizer='adam', loss='mse')
-    
+    model.compile(optimizer="adam", loss="mse")
+
     # Create some dummy data
     X = tf.random.normal((batch_size, num_samples, num_features))
     y = tf.random.normal((batch_size, num_samples, 1))
-    
+
     # Train for one epoch
     history = model.fit(X, y, epochs=1, verbose=0)
-    
+
     # Check if loss was computed
-    assert 'loss' in history.history
-    assert len(history.history['loss']) == 1
+    assert "loss" in history.history
+    assert len(history.history["loss"]) == 1
 
 
 def test_multi_resolution_attention_layer_init():
     """Test initialization of MultiResolutionTabularAttention layer."""
-    layer = MultiResolutionTabularAttention(
-        num_heads=4,
-        d_model=64,
-        embedding_dim=32
-    )
+    layer = MultiResolutionTabularAttention(num_heads=4, d_model=64, embedding_dim=32)
     assert layer.num_heads == 4
     assert layer.d_model == 64
     assert layer.embedding_dim == 32
@@ -148,16 +136,12 @@ def test_multi_resolution_attention_layer_init():
 def test_multi_resolution_attention_layer_config():
     """Test get_config and from_config methods for MultiResolutionTabularAttention."""
     original_layer = MultiResolutionTabularAttention(
-        num_heads=4,
-        d_model=64,
-        embedding_dim=32,
-        dropout_rate=0.2,
-        name="test_multi_attention"
+        num_heads=4, d_model=64, embedding_dim=32, dropout_rate=0.2, name="test_multi_attention"
     )
-    
+
     config = original_layer.get_config()
     restored_layer = MultiResolutionTabularAttention.from_config(config)
-    
+
     assert restored_layer.num_heads == original_layer.num_heads
     assert restored_layer.d_model == original_layer.d_model
     assert restored_layer.embedding_dim == original_layer.embedding_dim
@@ -172,29 +156,21 @@ def test_multi_resolution_attention_computation():
     num_categorical = 5
     numerical_dim = 16
     categorical_dim = 8
-    
+
     # Create a layer instance
-    layer = MultiResolutionTabularAttention(
-        num_heads=2,
-        d_model=numerical_dim,
-        embedding_dim=categorical_dim
-    )
-    
+    layer = MultiResolutionTabularAttention(num_heads=2, d_model=numerical_dim, embedding_dim=categorical_dim)
+
     # Create input data
     numerical_features = tf.random.normal((batch_size, num_numerical, numerical_dim))
     categorical_features = tf.random.normal((batch_size, num_categorical, categorical_dim))
-    
+
     # Call the layer
-    numerical_output, categorical_output = layer(
-        numerical_features,
-        categorical_features,
-        training=True
-    )
-    
+    numerical_output, categorical_output = layer(numerical_features, categorical_features, training=True)
+
     # Check output shapes
     assert numerical_output.shape == (batch_size, num_numerical, numerical_dim)
     assert categorical_output.shape == (batch_size, num_categorical, numerical_dim)
-    
+
     # Test with different batch sizes
     numerical_features_2 = tf.random.normal((64, num_numerical, numerical_dim))
     categorical_features_2 = tf.random.normal((64, num_categorical, categorical_dim))
@@ -210,31 +186,20 @@ def test_multi_resolution_attention_training():
     num_categorical = 3
     numerical_dim = 8
     categorical_dim = 4
-    
+
     layer = MultiResolutionTabularAttention(
-        num_heads=2,
-        d_model=numerical_dim,
-        embedding_dim=categorical_dim,
-        dropout_rate=0.5
+        num_heads=2, d_model=numerical_dim, embedding_dim=categorical_dim, dropout_rate=0.5
     )
-    
+
     numerical_features = tf.random.normal((batch_size, num_numerical, numerical_dim))
     categorical_features = tf.random.normal((batch_size, num_categorical, categorical_dim))
-    
+
     # Test in training mode
-    num_train, cat_train = layer(
-        numerical_features,
-        categorical_features,
-        training=True
-    )
-    
+    num_train, cat_train = layer(numerical_features, categorical_features, training=True)
+
     # Test in inference mode
-    num_infer, cat_infer = layer(
-        numerical_features,
-        categorical_features,
-        training=False
-    )
-    
+    num_infer, cat_infer = layer(numerical_features, categorical_features, training=False)
+
     # The outputs should be different due to dropout
     assert not np.allclose(num_train.numpy(), num_infer.numpy())
     assert not np.allclose(cat_train.numpy(), cat_infer.numpy())
@@ -243,13 +208,9 @@ def test_multi_resolution_attention_training():
 def test_multi_resolution_attention_factory():
     """Test creation of MultiResolutionTabularAttention layer through PreprocessorLayerFactory."""
     layer = PreprocessorLayerFactory.multi_resolution_attention_layer(
-        num_heads=4,
-        d_model=64,
-        embedding_dim=32,
-        name="test_multi_attention",
-        dropout_rate=0.2
+        num_heads=4, d_model=64, embedding_dim=32, name="test_multi_attention", dropout_rate=0.2
     )
-    
+
     assert isinstance(layer, MultiResolutionTabularAttention)
     assert layer.num_heads == 4
     assert layer.d_model == 64
@@ -266,44 +227,34 @@ def test_multi_resolution_attention_end_to_end():
     numerical_dim = 8
     categorical_dim = 4
     output_dim = 1
-    
+
     # Create inputs
     numerical_inputs = tf.keras.Input(shape=(num_numerical, numerical_dim))
     categorical_inputs = tf.keras.Input(shape=(num_categorical, categorical_dim))
-    
+
     # Apply multi-resolution attention
     num_attended, cat_attended = MultiResolutionTabularAttention(
-        num_heads=2,
-        d_model=numerical_dim,
-        embedding_dim=categorical_dim
+        num_heads=2, d_model=numerical_dim, embedding_dim=categorical_dim
     )(numerical_inputs, categorical_inputs)
-    
+
     # Combine outputs
     combined = tf.keras.layers.Concatenate(axis=1)([num_attended, cat_attended])
     outputs = tf.keras.layers.Dense(output_dim)(combined)
-    
+
     # Create model
-    model = tf.keras.Model(
-        inputs=[numerical_inputs, categorical_inputs],
-        outputs=outputs
-    )
-    
+    model = tf.keras.Model(inputs=[numerical_inputs, categorical_inputs], outputs=outputs)
+
     # Compile the model
-    model.compile(optimizer='adam', loss='mse')
-    
+    model.compile(optimizer="adam", loss="mse")
+
     # Create dummy data
     X_num = tf.random.normal((batch_size, num_numerical, numerical_dim))
     X_cat = tf.random.normal((batch_size, num_categorical, categorical_dim))
     y = tf.random.normal((batch_size, num_numerical + num_categorical, output_dim))
-    
+
     # Train for one epoch
-    history = model.fit(
-        [X_num, X_cat],
-        y,
-        epochs=1,
-        verbose=0
-    )
-    
+    history = model.fit([X_num, X_cat], y, epochs=1, verbose=0)
+
     # Check if loss was computed
-    assert 'loss' in history.history
-    assert len(history.history['loss']) == 1
+    assert "loss" in history.history
+    assert len(history.history["loss"]) == 1
