iBrokeTheCode/Multimodal_Product_Classification

1import numpy as np2import pandas as pd3import pytest4from transformers import AutoModel, AutoTokenizer5 6from src.nlp_models import HuggingFaceEmbeddings7 8# import torch9# import os10 11####################################################################################################12################################## Test the Text Embeddings Model ##################################13####################################################################################################14 15 16@pytest.fixture17def mock_text_data(tmp_path):18    """19    Fixture to create a mock CSV file with text data for testing.20    """21    data = {"description": ["Product 1 description", "Product 2 description"]}22    df = pd.DataFrame(data)23    file_path = tmp_path / "test_text_data.csv"24    df.to_csv(file_path, index=False)25    return str(file_path)26 27 28@pytest.mark.parametrize(29    "model_name, expected_hidden_size",30    [31        ("sentence-transformers/all-MiniLM-L6-v2", 384),  # MiniLM with 384 hidden units32        # ('bert-base-uncased', 768),  # BERT base with 768 hidden units33    ],34)35def test_huggingface_embeddings_generic(36    model_name, expected_hidden_size, mock_text_data37):38    """39    Generic test for loading a Hugging Face model, generating text embeddings, and saving them to a CSV file.40 41    This test ensures that:42    - The model and tokenizer are properly loaded from Hugging Face.43    - Embeddings are correctly generated for text descriptions.44    - Embeddings are saved in the correct format to a CSV file.45 46    Parameters:47    ----------48    model_name : str49        The name of the Hugging Face model to test.50    expected_hidden_size : int51        The expected hidden size (dimensionality) of the embeddings generated by the model.52    mock_text_data : str53        Path to the mock CSV file containing text descriptions.54    """55    # Initialize the HuggingFaceEmbeddings model with the provided model name56    model = HuggingFaceEmbeddings(57        model_name=model_name, path=mock_text_data, device="cpu"58    )59 60    # Check that the tokenizer and model were loaded correctly61    assert isinstance(62        model.tokenizer, type(AutoTokenizer.from_pretrained(model_name))63    ), (64        f"Tokenizer should be an instance of {type(AutoTokenizer.from_pretrained(model_name))}"65    )66    assert isinstance(model.model, type(AutoModel.from_pretrained(model_name))), (67        f"Model should be an instance of {type(AutoModel.from_pretrained(model_name))}"68    )69 70    # Generate embeddings for a sample text71    sample_text = "This is a test description."72    embeddings = model.get_embedding(sample_text)73 74    # Check that the embeddings are a NumPy array with the expected shape75    assert isinstance(embeddings, np.ndarray), "Embeddings should be a NumPy array"76    assert embeddings.shape == (expected_hidden_size,), (77        f"Embeddings shape should be ({expected_hidden_size},), got {embeddings.shape}"78    )79 80 81if __name__ == "__main__":82    pytest.main()83