iBrokeTheCode/Multimodal_Product_Classification

1# import numpy as np2# import os3 4# from src.utils import preprocess_data5# from sklearn.model_selection import train_test_split6 7 8import numpy as np9import pandas as pd10import pytest11 12from src.utils import train_test_split_and_feature_extraction13 14####################################################################################################15######################### Test the Train-Test Split and variable selection #########################16####################################################################################################17 18 19@pytest.fixture20def big_fake_data():21    # Create a fake dataset with 100 rows22    num_rows = 10023    num_image_columns = 1024    num_text_columns = 1125 26    data = {27        "id": np.arange(1, num_rows + 1),28        "image": [f"path/{i}.jpg" for i in range(1, num_rows + 1)],29    }30 31    # Add image_0 to image_9 columns32    for i in range(num_image_columns):33        data[f"image_{i}"] = np.random.rand(num_rows)34 35    # Add text_0 to text_10 columns36    for i in range(num_text_columns):37        data[f"text_{i}"] = np.random.rand(num_rows)38 39    # Add a class_id column40    data["class_id"] = np.random.choice(["label1", "label2", "label3"], size=num_rows)41 42    return pd.DataFrame(data)43 44 45def test_train_test_split_and_feature_extraction(big_fake_data):46    # Split the data and extract features and labels47    train_df, test_df, text_columns, image_columns, label_columns = (48        train_test_split_and_feature_extraction(49            big_fake_data, test_size=0.3, random_state=4250        )51    )52 53    # Check that the correct columns were identified54    assert text_columns == [f"text_{i}" for i in range(11)], (55        "The text embedding columns extraction is incorrect"56    )57    assert image_columns == [f"image_{i}" for i in range(10)], (58        "The image embedding columns extraction is incorrect"59    )60    assert label_columns == ["class_id"], (61        "The label column extraction is incorrect, should be 'class_id'"62    )63 64    # Check if 'image' is in the columns65    assert "image" not in image_columns, (66        "'image' column is not part of the embedding columns"67    )68 69    # Check the train-test split sizes (30% of 100 rows should be 70 train, 30 test)70    assert len(train_df) == 70, f"Train size should be 70%, but got {len(train_df)}%"71    assert len(test_df) == 30, f"Test size should be 30%, but got {len(test_df)}%"72 73    # Check random state consistency by ensuring the split results are reproducible74    expected_train_indices = train_df.index.tolist()75    expected_test_indices = test_df.index.tolist()76 77    # Re-run the function to check for consistency in split78    train_df_recheck, test_df_recheck, _, _, _ = (79        train_test_split_and_feature_extraction(80            big_fake_data, test_size=0.3, random_state=4281        )82    )83 84    assert expected_train_indices == train_df_recheck.index.tolist(), (85        "Train set indices are not consistent with the random state"86    )87    assert expected_test_indices == test_df_recheck.index.tolist(), (88        "Test set indices are not consistent with the random state"89    )90 91 92if __name__ == "__main__":93    pytest.main()94