chore: Add tests cases

tests/test_classifiers_classic_ml.py

@@ -0,0 +1,106 @@
+from unittest.mock import patch
+
+import pytest
+from sklearn.datasets import make_classification
+from sklearn.decomposition import PCA
+from sklearn.ensemble import RandomForestClassifier
+from sklearn.linear_model import LogisticRegression
+from sklearn.model_selection import train_test_split
+
+from src.classifiers_classic_ml import train_and_evaluate_model, visualize_embeddings
+
+####################################################################################################
+################################### Test the Classical ML Models ###################################
+####################################################################################################
+
+
+@pytest.fixture
+def sample_embedding_data():
+    """
+    Fixture to create a mock dataset for testing dimensionality reduction and model training.
+    Returns:
+        X_train, X_test, y_train, y_test: Training and testing data along with labels.
+    """
+    # Create a synthetic dataset with 20 samples, 6 features, and 3 classes
+    X, y = make_classification(
+        n_samples=20, n_features=6, n_classes=3, random_state=42, n_informative=4
+    )
+
+    # Split the dataset into training and test sets (80% train, 20% test)
+    X_train, X_test, y_train, y_test = train_test_split(
+        X, y, test_size=0.2, random_state=42
+    )
+
+    return X_train, X_test, y_train, y_test
+
+
+@pytest.mark.parametrize(
+    "method, plot_type",
+    [
+        ("PCA", "2D"),  # PCA reduction to 2D
+        ("PCA", "3D"),  # PCA reduction to 3D
+    ],
+)
+def test_visualize_embeddings(method, plot_type, sample_embedding_data):
+    """
+    Test the dimensionality reduction and embedding visualization.
+    This ensures that PCA can reduce embeddings correctly and produce visualizations.
+    """
+    X_train, X_test, y_train, y_test = sample_embedding_data
+
+    # Mock the plotly figures to avoid actual plotting in test environment
+    with patch("plotly.graph_objs.Figure.show"):
+        # Test the visualize_embeddings function
+        model = visualize_embeddings(
+            X_train, X_test, y_train, y_test, plot_type=plot_type, method=method
+        )
+
+    # Check if the PCA model is an instance of the correct class and has the expected number of components
+    assert isinstance(model, PCA), "The model should be an instance of PCA"
+    if plot_type == "2D":
+        assert model.n_components_ == 2, "PCA should reduce data to 2 components"
+    elif plot_type == "3D":
+        assert model.n_components_ == 3, "PCA should reduce data to 3 components"
+
+
+def test_train_and_evaluate_model(sample_embedding_data):
+    """
+    Test the training and evaluation of models (Logistic Regression, Random Forest).
+    Ensures that models are correctly trained and returned in the expected format.
+    """
+    X_train, X_test, y_train, y_test = sample_embedding_data
+
+    # Train and evaluate the models
+    trained_models = train_and_evaluate_model(
+        X_train, X_test, y_train, y_test, test=False
+    )
+
+    # Verify that trained_models is a list
+    assert isinstance(trained_models, list), (
+        "The output should be a list of trained models"
+    )
+
+    # Check that at least two models were trained (Logistic Regression, Random Forest)
+    assert len(trained_models) >= 2, "At least two models should be trained"
+
+    # Check that the models have Logistic Regression and Random Forest
+    models_instances = [model for _, model in trained_models]
+    assert any(isinstance(model, LogisticRegression) for model in models_instances), (
+        "Logistic Regression model not found"
+    )
+    assert any(
+        isinstance(model, RandomForestClassifier) for model in models_instances
+    ), "Random Forest model not found"
+
+    # Ensure that the trained models are indeed fitted (trained)
+    for name, model in trained_models:
+        assert hasattr(model, "fit"), f"{name} should have a fit method"
+        assert hasattr(model, "predict"), f"{name} should have a predict method"
+
+        # Check if the model is correctly trained by predicting on the test set
+        y_pred = model.predict(X_test)
+        assert y_pred is not None, f"{name} should have successfully made predictions"
+
+
+if __name__ == "__main__":
+    pytest.main()

tests/test_classifiers_mlp.py

@@ -0,0 +1,626 @@
+# import numpy as np
+# from sklearn.decomposition import PCA
+# from sklearn.manifold import TSNE
+# from src.classifiers_classic_ml import visualize_embeddings, train_and_evaluate_model
+
+import os
+
+import pandas as pd
+import pytest
+from sklearn.datasets import make_classification
+from sklearn.metrics import accuracy_score, f1_score
+from sklearn.model_selection import train_test_split
+from sklearn.preprocessing import LabelEncoder
+from tensorflow.keras.layers import BatchNormalization, Concatenate, Dense, Dropout
+from tensorflow.keras.losses import CategoricalCrossentropy
+from tensorflow.keras.models import Model
+from tensorflow.keras.optimizers import SGD, Adam
+
+from src.classifiers_mlp import MultimodalDataset, create_early_fusion_model, train_mlp
+
+####################################################################################################
+##################################### Test the Keras MLP Models ####################################
+####################################################################################################
+
+
+@pytest.fixture
+def correlated_sample_data():
+    """
+    Fixture to create a correlated synthetic dataset using make_classification for testing.
+    It generates data with 10 text features and 10 image features.
+    Returns:
+        train_df (pd.DataFrame): DataFrame with train data.
+        test_df (pd.DataFrame): DataFrame with test data.
+    """
+    # Create synthetic multi-class data with 8 features (4 text-like, 4 image-like)
+    X, y = make_classification(
+        n_samples=20, n_features=8, n_informative=6, n_classes=3, random_state=42
+    )
+
+    # Rename features to simulate text and image columns
+    feature_names = [f"text_{i}" for i in range(4)] + [
+        f"image_{i}" for i in range(4, 8)
+    ]
+
+    # Create a DataFrame and assign class labels
+    df = pd.DataFrame(X, columns=feature_names)
+    df["class_id"] = y
+
+    # Split into train and test sets
+    train_df, test_df = train_test_split(df, test_size=0.3, random_state=42)
+
+    return train_df, test_df
+
+
+@pytest.fixture
+def label_encoder(correlated_sample_data):
+    """
+    Fixture to create a label encoder based on the training data.
+    """
+    train_df, _ = correlated_sample_data
+    label_encoder = LabelEncoder()
+    label_encoder.fit(train_df["class_id"])
+    return label_encoder
+
+
+def test_multimodal_dataset_image_only(correlated_sample_data, label_encoder):
+    """
+    Test the MultimodalDataset class with only image data.
+    """
+    train_df, test_df = correlated_sample_data
+
+    # Image columns (the second 4 features)
+    image_columns = [f"image_{i}" for i in range(4, 8)]
+    label_column = "class_id"
+
+    # Create the dataset
+    train_dataset = MultimodalDataset(
+        train_df,
+        text_cols=None,
+        image_cols=image_columns,
+        label_col=label_column,
+        encoder=label_encoder,
+    )
+
+    # Check if the dataset is correctly instantiated
+    assert train_dataset.image_data is not None, "Image data should be instantiated"
+    assert train_dataset.text_data is None, "Text data should be None"
+
+    # Fetch a batch of data
+    (batch_inputs, batch_labels) = train_dataset[0]
+
+    assert "image" in batch_inputs, "Batch should contain image data"
+    assert "text" not in batch_inputs, "Batch should not contain text data"
+    assert batch_inputs["image"].shape[1] == len(image_columns), (
+        "Image data shape is incorrect"
+    )
+    assert batch_labels is not None, "Batch should contain labels"
+    assert batch_labels.shape[0] == batch_inputs["image"].shape[0], (
+        "Labels should match the batch size"
+    )
+
+
+def test_multimodal_dataset_text_only(correlated_sample_data, label_encoder):
+    """
+    Test the MultimodalDataset class with only text data.
+    """
+    train_df, test_df = correlated_sample_data
+
+    # Text columns (the first 4 features)
+    text_columns = [f"text_{i}" for i in range(4)]
+    label_column = "class_id"
+
+    # Create the dataset
+    train_dataset = MultimodalDataset(
+        train_df,
+        text_cols=text_columns,
+        image_cols=None,
+        label_col=label_column,
+        encoder=label_encoder,
+    )
+
+    # Check if the dataset is correctly instantiated
+    assert train_dataset.text_data is not None, "Text data should be instantiated"
+    assert train_dataset.image_data is None, "Image data should be None"
+
+    # Fetch a batch of data
+    (batch_inputs, batch_labels) = train_dataset[0]
+
+    assert "text" in batch_inputs, "Batch should contain text data"
+    assert "image" not in batch_inputs, "Batch should not contain image data"
+    assert batch_inputs["text"].shape[1] == len(text_columns), (
+        "Text data shape is incorrect"
+    )
+    assert batch_labels is not None, "Batch should contain labels"
+    assert batch_labels.shape[0] == batch_inputs["text"].shape[0], (
+        "Labels should match the batch size"
+    )
+
+
+def test_multimodal_dataset_multimodal(correlated_sample_data, label_encoder):
+    """
+    Test the MultimodalDataset class with both text and image data.
+    """
+    train_df, test_df = correlated_sample_data
+
+    # Text and image columns
+    text_columns = [f"text_{i}" for i in range(4)]
+    image_columns = [f"image_{i}" for i in range(4, 8)]
+    label_column = "class_id"
+
+    # Create the dataset
+    train_dataset = MultimodalDataset(
+        train_df,
+        text_cols=text_columns,
+        image_cols=image_columns,
+        label_col=label_column,
+        encoder=label_encoder,
+    )
+
+    # Check if the dataset is correctly instantiated
+    assert train_dataset.text_data is not None, "Text data should be instantiated"
+    assert train_dataset.image_data is not None, "Image data should be instantiated"
+
+    # Fetch a batch of data
+    (batch_inputs, batch_labels) = train_dataset[0]
+    assert "text" in batch_inputs, "Batch should contain text data"
+    assert "image" in batch_inputs, "Batch should contain image data"
+    assert batch_inputs["text"].shape[1] == len(text_columns), (
+        "Text data shape is incorrect"
+    )
+    assert batch_inputs["image"].shape[1] == len(image_columns), (
+        "Image data shape is incorrect"
+    )
+    assert batch_labels is not None, "Batch should contain labels"
+    assert (
+        batch_labels.shape[0]
+        == batch_inputs["text"].shape[0]
+        == batch_inputs["image"].shape[0]
+    ), "Labels should match the batch size"
+
+
+def test_create_early_fusion_model_single_modality_image():
+    """
+    Test the model creation with only image input or only text input.
+    Ensure the architecture matches expectations.
+    """
+    text_input_size = None
+    image_input_size = 4
+    output_size = 3
+
+    # Create the model
+    model = create_early_fusion_model(
+        text_input_size, image_input_size, output_size, hidden=[128, 64], p=0.3
+    )
+
+    # Check if the model has the expected number of layers
+    assert isinstance(model, Model), "Model should be a Keras Model instance"
+
+    # Check that the input and output shapes are consistent
+    assert model.input_shape == (None, image_input_size), (
+        "Input shape should match image input size"
+    )
+    assert model.output_shape == (None, output_size), (
+        "Output shape should match number of classes"
+    )
+
+    # Check that there are the correct number of Dense, Dropout, and BatchNormalization layers
+    dense_layers = [layer for layer in model.layers if isinstance(layer, Dense)]
+    dropout_layers = [layer for layer in model.layers if isinstance(layer, Dropout)]
+    batchnorm_layers = [
+        layer for layer in model.layers if isinstance(layer, BatchNormalization)
+    ]
+
+    assert len(dense_layers) == 3, (
+        "There should be 3 Dense layers (2 hidden + 1 output)"
+    )
+    assert len(dropout_layers) > 0, "There should be at least 1 Dropout layers"
+    assert len(batchnorm_layers) > 0, (
+        "There should be at least 1 BatchNormalization layer"
+    )
+
+
+def test_create_early_fusion_model_single_modality_text():
+    """
+    Test the model creation with only image input or only text input.
+    Ensure the architecture matches expectations.
+    """
+    text_input_size = 4
+    image_input_size = None
+    output_size = 3
+
+    # Create the model
+    model = create_early_fusion_model(
+        text_input_size, image_input_size, output_size, hidden=[128, 64], p=0.3
+    )
+
+    # Check if the model has the expected number of layers
+    assert isinstance(model, Model), "Model should be a Keras Model instance"
+
+    # Check that the input and output shapes are consistent
+    assert model.input_shape == (None, text_input_size), (
+        "Input shape should match text input size"
+    )
+    assert model.output_shape == (None, output_size), (
+        "Output shape should match number of classes"
+    )
+
+    # Check that there are the correct number of Dense, Dropout, and BatchNormalization layers
+    dense_layers = [layer for layer in model.layers if isinstance(layer, Dense)]
+    dropout_layers = [layer for layer in model.layers if isinstance(layer, Dropout)]
+    batchnorm_layers = [
+        layer for layer in model.layers if isinstance(layer, BatchNormalization)
+    ]
+
+    assert len(dense_layers) == 3, (
+        "There should be 3 Dense layers (2 hidden + 1 output)"
+    )
+    assert len(dropout_layers) > 0, "There should be at least 1 Dropout layers"
+    assert len(batchnorm_layers) > 0, (
+        "There should be at least 1 BatchNormalization layer"
+    )
+
+
+def test_create_early_fusion_model_multimodal():
+    """
+    Test the model creation with both text and image input.
+    Ensure the architecture matches expectations.
+    """
+    text_input_size = 4
+    image_input_size = 4
+    output_size = 3
+
+    # Create the model
+    model = create_early_fusion_model(
+        text_input_size, image_input_size, output_size, hidden=[128, 64], p=0.3
+    )
+
+    # Check if the model has the expected number of layers
+    assert isinstance(model, Model), "Model should be a Keras Model instance"
+
+    # Check that the input and output shapes are consistent
+    assert model.input_shape == [(None, text_input_size), (None, image_input_size)], (
+        "Input shape should match both text and image input sizes"
+    )
+    assert model.output_shape == (None, output_size), (
+        "Output shape should match number of classes"
+    )
+
+    # Check that the concatenation of text and image inputs is present
+    assert any(isinstance(layer, Concatenate) for layer in model.layers), (
+        "There should be a Concatenate layer for text and image inputs"
+    )
+
+    # Check that there are the correct number of Dense, Dropout, and BatchNormalization layers
+    dense_layers = [layer for layer in model.layers if isinstance(layer, Dense)]
+    dropout_layers = [layer for layer in model.layers if isinstance(layer, Dropout)]
+    batchnorm_layers = [
+        layer for layer in model.layers if isinstance(layer, BatchNormalization)
+    ]
+
+    assert len(dense_layers) == 3, (
+        "There should be 3 Dense layers (2 hidden + 1 output)"
+    )
+    assert len(dropout_layers) > 0, "There should be at least 1 Dropout layers"
+    assert len(batchnorm_layers) > 0, (
+        "There should be at least 1 BatchNormalization layer"
+    )
+
+
+def test_train_mlp_single_modality_image(correlated_sample_data, label_encoder):
+    """
+    Test the MLP training with only image data.
+    Ensure the model trains and evaluates correctly.
+    """
+    train_df, test_df = correlated_sample_data
+
+    # Image columns (the second 10 features)
+    image_columns = [f"image_{i}" for i in range(4, 8)]
+    label_column = "class_id"
+
+    # Create datasets
+    train_dataset = MultimodalDataset(
+        train_df,
+        text_cols=None,
+        image_cols=image_columns,
+        label_col=label_column,
+        encoder=label_encoder,
+    )
+    test_dataset = MultimodalDataset(
+        test_df,
+        text_cols=None,
+        image_cols=image_columns,
+        label_col=label_column,
+        encoder=label_encoder,
+    )
+
+    image_input_size = len(image_columns)
+    output_size = len(label_encoder.classes_)
+
+    # Train the model
+    model, test_accuracy, f1, macro_auc = train_mlp(
+        train_loader=train_dataset,
+        test_loader=test_dataset,
+        text_input_size=None,
+        image_input_size=image_input_size,
+        output_size=output_size,
+        num_epochs=1,
+        set_weights=True,
+        adam=True,
+        patience=10,
+        save_results=False,
+        train_model=False,
+        test_mlp_model=False,
+    )
+
+    # Check model
+    assert model is not None, "Model should not be None after training."
+
+    # Ensure the model is compiled with the correct loss and optimizer
+    assert (
+        isinstance(model.loss, CategoricalCrossentropy)
+        or model.loss == "categorical_crossentropy"
+    ), f"Loss function should be categorical crossentropy, but got {model.loss}"
+
+    # Check model input and output shapes
+    assert model.input_shape == (None, image_input_size), (
+        "Input shape should match image input size"
+    )
+    assert model.output_shape == (None, output_size), (
+        "Output shape should match number of classes"
+    )
+
+    # Check if the model is compiled with the correct optimizer
+    assert isinstance(model.optimizer, Adam) or isinstance(model.optimizer, SGD), (
+        f"Optimizer should be Adam or SGD, but got {model.optimizer}"
+    )
+
+
+def test_train_mlp_single_modality_text(correlated_sample_data, label_encoder):
+    """
+    Test the MLP training with only text data.
+    Ensure the model trains and evaluates correctly.
+    """
+    train_df, test_df = correlated_sample_data
+
+    # Text columns (the first 10 features)
+    text_columns = [f"text_{i}" for i in range(4)]
+    label_column = "class_id"
+
+    # Create datasets
+    train_dataset = MultimodalDataset(
+        train_df,
+        text_cols=text_columns,
+        image_cols=None,
+        label_col=label_column,
+        encoder=label_encoder,
+    )
+    test_dataset = MultimodalDataset(
+        test_df,
+        text_cols=text_columns,
+        image_cols=None,
+        label_col=label_column,
+        encoder=label_encoder,
+    )
+
+    text_input_size = len(text_columns)
+    output_size = len(label_encoder.classes_)
+
+    # Train the model
+    model, test_accuracy, f1, macro_auc = train_mlp(
+        train_loader=train_dataset,
+        test_loader=test_dataset,
+        text_input_size=text_input_size,
+        image_input_size=None,
+        output_size=output_size,
+        num_epochs=1,
+        set_weights=True,
+        adam=True,
+        patience=10,
+        save_results=False,
+        train_model=False,
+        test_mlp_model=False,
+    )
+
+    # Check model
+    assert model is not None, "Model should not be None after training."
+
+    # Ensure the model is compiled with the correct loss and optimizer
+    assert (
+        isinstance(model.loss, CategoricalCrossentropy)
+        or model.loss == "categorical_crossentropy"
+    ), f"Loss function should be categorical crossentropy, but got {model.loss}"
+
+    # Check model input and output shapes
+    assert model.input_shape == (None, text_input_size), (
+        "Input shape should match text input size"
+    )
+    assert model.output_shape == (None, output_size), (
+        "Output shape should match number of classes"
+    )
+
+    # Check if the model is compiled with the correct optimizer
+    assert isinstance(model.optimizer, Adam) or isinstance(model.optimizer, SGD), (
+        f"Optimizer should be Adam or SGD, but got {model.optimizer}"
+    )
+
+
+def test_train_mlp_multimodal(correlated_sample_data, label_encoder):
+    """
+    Test the MLP training with class weights for an imbalanced dataset.
+    Ensure class weights are applied correctly and early stopping works.
+    """
+    train_df, test_df = correlated_sample_data
+
+    # Text and image columns
+    text_columns = [f"text_{i}" for i in range(4)]
+    image_columns = [f"image_{i}" for i in range(4, 8)]
+    label_column = "class_id"
+
+    # Create datasets
+    train_dataset = MultimodalDataset(
+        train_df,
+        text_cols=text_columns,
+        image_cols=image_columns,
+        label_col=label_column,
+        encoder=label_encoder,
+    )
+    test_dataset = MultimodalDataset(
+        test_df,
+        text_cols=text_columns,
+        image_cols=image_columns,
+        label_col=label_column,
+        encoder=label_encoder,
+    )
+
+    text_input_size = len(text_columns)
+    image_input_size = len(image_columns)
+    output_size = len(label_encoder.classes_)
+
+    # Train the model
+    model, test_accuracy, f1, macro_auc = train_mlp(
+        train_loader=train_dataset,
+        test_loader=test_dataset,
+        text_input_size=text_input_size,
+        image_input_size=image_input_size,
+        output_size=output_size,
+        num_epochs=1,
+        set_weights=True,
+        adam=True,
+        patience=10,
+        save_results=False,
+        train_model=False,
+        test_mlp_model=False,
+    )
+
+    # Check model
+    assert model is not None, "Model should not be None after training."
+
+    # Ensure the model is compiled with the correct loss and optimizer
+    assert (
+        isinstance(model.loss, CategoricalCrossentropy)
+        or model.loss == "categorical_crossentropy"
+    ), f"Loss function should be categorical crossentropy, but got {model.loss}"
+
+    # Check model input and output shapes
+    assert model.input_shape == [(None, text_input_size), (None, image_input_size)], (
+        "Input shape should match both text and image input sizes"
+    )
+    assert model.output_shape == (None, output_size), (
+        "Output shape should match number of classes"
+    )
+
+    # Check if the model is compiled with the correct optimizer
+    assert isinstance(model.optimizer, Adam) or isinstance(model.optimizer, SGD), (
+        f"Optimizer should be Adam or SGD, but got {model.optimizer}"
+    )
+
+
+# Check if the result files are correctly saved
+def test_result_files():
+    """
+    Test if the result files are created for each modality and have the correct format.
+    """
+    # Get the absolute path of the directory where this test file is located
+    test_dir = os.path.dirname(os.path.abspath(__file__))
+
+    # Paths for result files relative to the test file location
+    multimodal_results_path = os.path.join(
+        test_dir, "../results/multimodal_results.csv"
+    )
+    text_results_path = os.path.join(test_dir, "../results/text_results.csv")
+    image_results_path = os.path.join(test_dir, "../results/image_results.csv")
+
+    # Check if the files exist
+    assert os.path.exists(multimodal_results_path), "Multimodal result file is missing!"
+    assert os.path.exists(text_results_path), "Text result file is missing!"
+    assert os.path.exists(image_results_path), "Image result file is missing!"
+
+    # Check if the files are not empty and in correct format (CSV)
+    for file_path in [multimodal_results_path, text_results_path, image_results_path]:
+        df = pd.read_csv(file_path)
+        assert not df.empty, f"{file_path} is empty!"
+        assert "Predictions" in df.columns and "True Labels" in df.columns, (
+            f"{file_path} is not in the correct format!"
+        )
+
+
+# Check if the accuracy and F1 scores meet the specified thresholds
+def test_model_performance():
+    """
+    Test if the accuracy and F1 score are above the required thresholds.
+    """
+    # Get the absolute path of the directory where this test file is located
+    test_dir = os.path.dirname(os.path.abspath(__file__))
+
+    # Paths for result files relative to the test file location
+    multimodal_results_path = os.path.join(
+        test_dir, "../results/multimodal_results.csv"
+    )
+    text_results_path = os.path.join(test_dir, "../results/text_results.csv")
+    image_results_path = os.path.join(test_dir, "../results/image_results.csv")
+
+    # Load the result files
+    multimodal_results = pd.read_csv(multimodal_results_path)
+    text_results = pd.read_csv(text_results_path)
+    image_results = pd.read_csv(image_results_path)
+
+    # Define the accuracy and F1-score thresholds
+    multimodal_accuracy_threshold = 0.85
+    multimodal_f1_threshold = 0.80
+    text_accuracy_threshold = 0.85
+    text_f1_threshold = 0.80
+    image_accuracy_threshold = 0.75
+    image_f1_threshold = 0.70
+
+    # Calculate accuracy and F1 score for multimodal results
+    multimodal_accuracy = accuracy_score(
+        multimodal_results["True Labels"], multimodal_results["Predictions"]
+    )
+    multimodal_f1 = f1_score(
+        multimodal_results["True Labels"],
+        multimodal_results["Predictions"],
+        average="macro",
+    )
+
+    # Calculate accuracy and F1 score for text results
+    text_accuracy = accuracy_score(
+        text_results["True Labels"], text_results["Predictions"]
+    )
+    text_f1 = f1_score(
+        text_results["True Labels"], text_results["Predictions"], average="macro"
+    )
+
+    # Calculate accuracy and F1 score for image results
+    image_accuracy = accuracy_score(
+        image_results["True Labels"], image_results["Predictions"]
+    )
+    image_f1 = f1_score(
+        image_results["True Labels"], image_results["Predictions"], average="macro"
+    )
+
+    # Check multimodal performance
+    assert multimodal_accuracy > multimodal_accuracy_threshold, (
+        f"Multimodal accuracy is below {multimodal_accuracy_threshold}"
+    )
+    assert multimodal_f1 > multimodal_f1_threshold, (
+        f"Multimodal F1 score is below {multimodal_f1_threshold}"
+    )
+
+    # Check text performance
+    assert text_accuracy > text_accuracy_threshold, (
+        f"Text accuracy is below {text_accuracy_threshold}"
+    )
+    assert text_f1 > text_f1_threshold, f"Text F1 score is below {text_f1_threshold}"
+
+    # Check image performance
+    assert image_accuracy > image_accuracy_threshold, (
+        f"Image accuracy is below {image_accuracy_threshold}"
+    )
+    assert image_f1 > image_f1_threshold, (
+        f"Image F1 score is below {image_f1_threshold}"
+    )
+
+
+if __name__ == "__main__":
+    pytest.main()

tests/test_nlp_models.py

@@ -0,0 +1,82 @@
+import numpy as np
+import pandas as pd
+import pytest
+from transformers import AutoModel, AutoTokenizer
+
+from src.nlp_models import HuggingFaceEmbeddings
+
+# import torch
+# import os
+
+####################################################################################################
+################################## Test the Text Embeddings Model ##################################
+####################################################################################################
+
+
+@pytest.fixture
+def mock_text_data(tmp_path):
+    """
+    Fixture to create a mock CSV file with text data for testing.
+    """
+    data = {"description": ["Product 1 description", "Product 2 description"]}
+    df = pd.DataFrame(data)
+    file_path = tmp_path / "test_text_data.csv"
+    df.to_csv(file_path, index=False)
+    return str(file_path)
+
+
+@pytest.mark.parametrize(
+    "model_name, expected_hidden_size",
+    [
+        ("sentence-transformers/all-MiniLM-L6-v2", 384),  # MiniLM with 384 hidden units
+        # ('bert-base-uncased', 768),  # BERT base with 768 hidden units
+    ],
+)
+def test_huggingface_embeddings_generic(
+    model_name, expected_hidden_size, mock_text_data
+):
+    """
+    Generic test for loading a Hugging Face model, generating text embeddings, and saving them to a CSV file.
+
+    This test ensures that:
+    - The model and tokenizer are properly loaded from Hugging Face.
+    - Embeddings are correctly generated for text descriptions.
+    - Embeddings are saved in the correct format to a CSV file.
+
+    Parameters:
+    ----------
+    model_name : str
+        The name of the Hugging Face model to test.
+    expected_hidden_size : int
+        The expected hidden size (dimensionality) of the embeddings generated by the model.
+    mock_text_data : str
+        Path to the mock CSV file containing text descriptions.
+    """
+    # Initialize the HuggingFaceEmbeddings model with the provided model name
+    model = HuggingFaceEmbeddings(
+        model_name=model_name, path=mock_text_data, device="cpu"
+    )
+
+    # Check that the tokenizer and model were loaded correctly
+    assert isinstance(
+        model.tokenizer, type(AutoTokenizer.from_pretrained(model_name))
+    ), (
+        f"Tokenizer should be an instance of {type(AutoTokenizer.from_pretrained(model_name))}"
+    )
+    assert isinstance(model.model, type(AutoModel.from_pretrained(model_name))), (
+        f"Model should be an instance of {type(AutoModel.from_pretrained(model_name))}"
+    )
+
+    # Generate embeddings for a sample text
+    sample_text = "This is a test description."
+    embeddings = model.get_embedding(sample_text)
+
+    # Check that the embeddings are a NumPy array with the expected shape
+    assert isinstance(embeddings, np.ndarray), "Embeddings should be a NumPy array"
+    assert embeddings.shape == (expected_hidden_size,), (
+        f"Embeddings shape should be ({expected_hidden_size},), got {embeddings.shape}"
+    )
+
+
+if __name__ == "__main__":
+    pytest.main()

tests/test_utils.py

@@ -0,0 +1,93 @@
+# import numpy as np
+# import os
+
+# from src.utils import preprocess_data
+# from sklearn.model_selection import train_test_split
+
+
+import numpy as np
+import pandas as pd
+import pytest
+
+from src.utils import train_test_split_and_feature_extraction
+
+####################################################################################################
+######################### Test the Train-Test Split and variable selection #########################
+####################################################################################################
+
+
+@pytest.fixture
+def big_fake_data():
+    # Create a fake dataset with 100 rows
+    num_rows = 100
+    num_image_columns = 10
+    num_text_columns = 11
+
+    data = {
+        "id": np.arange(1, num_rows + 1),
+        "image": [f"path/{i}.jpg" for i in range(1, num_rows + 1)],
+    }
+
+    # Add image_0 to image_9 columns
+    for i in range(num_image_columns):
+        data[f"image_{i}"] = np.random.rand(num_rows)
+
+    # Add text_0 to text_10 columns
+    for i in range(num_text_columns):
+        data[f"text_{i}"] = np.random.rand(num_rows)
+
+    # Add a class_id column
+    data["class_id"] = np.random.choice(["label1", "label2", "label3"], size=num_rows)
+
+    return pd.DataFrame(data)
+
+
+def test_train_test_split_and_feature_extraction(big_fake_data):
+    # Split the data and extract features and labels
+    train_df, test_df, text_columns, image_columns, label_columns = (
+        train_test_split_and_feature_extraction(
+            big_fake_data, test_size=0.3, random_state=42
+        )
+    )
+
+    # Check that the correct columns were identified
+    assert text_columns == [f"text_{i}" for i in range(11)], (
+        "The text embedding columns extraction is incorrect"
+    )
+    assert image_columns == [f"image_{i}" for i in range(10)], (
+        "The image embedding columns extraction is incorrect"
+    )
+    assert label_columns == ["class_id"], (
+        "The label column extraction is incorrect, should be 'class_id'"
+    )
+
+    # Check if 'image' is in the columns
+    assert "image" not in image_columns, (
+        "'image' column is not part of the embedding columns"
+    )
+
+    # Check the train-test split sizes (30% of 100 rows should be 70 train, 30 test)
+    assert len(train_df) == 70, f"Train size should be 70%, but got {len(train_df)}%"
+    assert len(test_df) == 30, f"Test size should be 30%, but got {len(test_df)}%"
+
+    # Check random state consistency by ensuring the split results are reproducible
+    expected_train_indices = train_df.index.tolist()
+    expected_test_indices = test_df.index.tolist()
+
+    # Re-run the function to check for consistency in split
+    train_df_recheck, test_df_recheck, _, _, _ = (
+        train_test_split_and_feature_extraction(
+            big_fake_data, test_size=0.3, random_state=42
+        )
+    )
+
+    assert expected_train_indices == train_df_recheck.index.tolist(), (
+        "Train set indices are not consistent with the random state"
+    )
+    assert expected_test_indices == test_df_recheck.index.tolist(), (
+        "Test set indices are not consistent with the random state"
+    )
+
+
+if __name__ == "__main__":
+    pytest.main()

tests/test_vision_embeddings_tf.py

@@ -0,0 +1,110 @@
+# import os
+# import pandas as pd
+
+# from src.vision_embeddings_tf import get_embeddings_df
+
+
+import numpy as np
+import pytest
+from PIL import Image
+from tensorflow.keras.applications import ResNet50
+from transformers import TFConvNextV2Model
+
+from src.vision_embeddings_tf import FoundationalCVModel, load_and_preprocess_image
+
+# Run tests with CPU and not GPU (custom added)
+# os.environ["CUDA_VISIBLE_DEVICES"] = "-1"
+
+
+####################################################################################################
+#################### Test the foundational CV model and image preprocessing ########################
+####################################################################################################
+@pytest.fixture
+def mock_image(tmp_path):
+    """
+    Fixture to create a mock image for testing.
+    """
+    img_path = tmp_path / "test_image.jpg"
+    img = Image.new("RGB", (300, 300), color="red")
+    img.save(img_path)
+    return str(img_path)
+
+
+def test_load_and_preprocess_image(mock_image):
+    """
+    Test loading and preprocessing of an image.
+    """
+    # Test the load_and_preprocess_image function
+    img = load_and_preprocess_image(mock_image, target_size=(224, 224))
+
+    # Check if the output is a numpy array
+    assert isinstance(img, np.ndarray), "Output is not a numpy array"
+
+    # Check if the image has the correct shape
+    assert img.shape == (224, 224, 3), (
+        f"Image shape is {img.shape}, expected (224, 224, 3)"
+    )
+
+    # Check if the pixel values are in the range [0, 1]
+    assert img.min() >= 0 and img.max() <= 1, (
+        "Image pixel values are not in the range [0, 1]"
+    )
+
+
+@pytest.mark.parametrize(
+    "backbone, expected_model_class, expected_output_shape",
+    [
+        ("resnet50", type(ResNet50()), (2048,)),  # Keras ResNet50 with 2048 features
+        (
+            "convnextv2_tiny",
+            TFConvNextV2Model,
+            (768,),
+        ),  # ConvNeXt V2 Tiny from Hugging Face with 768 features
+    ],
+)
+def test_foundational_cv_model_generic(
+    backbone, expected_model_class, expected_output_shape
+):
+    """
+    Generic test for loading a foundational CV model and making predictions.
+
+    This test ensures that:
+    - The correct backbone model is loaded.
+    - The input shape matches the model's requirements (224x224x3).
+    - The output embedding shape matches the expected shape for the backbone.
+
+    Parameters:
+    ----------
+    backbone : str
+        The name of the model backbone to test.
+    expected_model_class : class
+        The expected class of the loaded backbone model (e.g., ResNet50 or TFConvNextV2Model).
+    expected_output_shape : tuple
+        The expected shape of the output embedding vector.
+    """
+    # Initialize the model with the provided backbone
+    model = FoundationalCVModel(backbone=backbone, mode="eval")
+
+    # Check if the model is an instance of the expected model class
+    assert isinstance(model.base_model, expected_model_class), (
+        f"Expected model class {expected_model_class}, got {type(model.model)}"
+    )
+
+    # Create a batch of random images (2 images of shape 224x224x3)
+    batch_images = np.random.rand(2, 224, 224, 3)
+
+    # Ensure that the input shape matches the model's input requirements
+    assert model.model.input_shape == (None, 224, 224, 3), (
+        f"Expected input shape (None, 224, 224, 3), got {model.model.input_shape}"
+    )
+
+    # Ensure that the output shape matches the expected output shape without using the model.predict method
+    output = model.get_output_shape()
+
+    assert output == (None, *expected_output_shape), (
+        f"Expected output shape (None, {expected_output_shape}), got {output}"
+    )
+
+
+if __name__ == "__main__":
+    pytest.main()