""" Tests that sentence transformer model is loaded and yielding a structurally correct embedding vector.

To use this test, you may need install the SentenceTransformer library as follows:

    -- pip3 install sentence-transformers

"""


from llmware.models import ModelCatalog
from sentence_transformers import SentenceTransformer


def test_sentence_transformer_model_local_load():

    # This model list was generated by here https://www.sbert.net/docs/pretrained_models.html and
    # selecting the "All Models" switch

    sentence_transformer_models = [
        'all-MiniLM-L12-v1',
        'all-MiniLM-L12-v2',
        'all-MiniLM-L6-v1',
        'all-MiniLM-L6-v2',
        'all-distilroberta-v1',
        'all-mpnet-base-v1',
        'all-mpnet-base-v2',
        'all-roberta-large-v1',
        'average_word_embeddings_glove.6B.300d',
        'average_word_embeddings_komninos',
        'gtr-t5-base',
        'gtr-t5-large',
        'gtr-t5-xl',
        'gtr-t5-xxl',
        'msmarco-bert-base-dot-v5',
        'msmarco-distilbert-base-tas-b',
        'msmarco-distilbert-dot-v5',
        'multi-qa-MiniLM-L6-cos-v1',
        'multi-qa-MiniLM-L6-dot-v1',
        'multi-qa-distilbert-cos-v1',
        'multi-qa-distilbert-dot-v1',
        'multi-qa-mpnet-base-cos-v1',
        'multi-qa-mpnet-base-dot-v1',
        'paraphrase-MiniLM-L12-v2',
        'paraphrase-MiniLM-L3-v2',
        'paraphrase-MiniLM-L6-v2',
        'paraphrase-TinyBERT-L6-v2',
        'paraphrase-albert-small-v2',
        'paraphrase-distilroberta-base-v2',
        'paraphrase-mpnet-base-v2',
        'paraphrase-multilingual-MiniLM-L12-v2',
        'paraphrase-multilingual-mpnet-base-v2',
        'sentence-t5-base',
        'sentence-t5-large',
        'sentence-t5-xl',
        'sentence-t5-xxl'
    ]

    test_text = ("This is just a sample text to confirm that the embedding model is loading and correctly "
                 "converting into a structurally accurate embedding vector.")

    for model_name in sentence_transformer_models:

        print(f"\nloading sentence transformer model: {model_name}")

        st_model = SentenceTransformer(model_name)
        model = ModelCatalog().load_sentence_transformer_model(st_model, model_name=model_name)
        embedding_vector = model.embedding([test_text])

        assert embedding_vector is not None

        print(f"created vector successfully with dimensions: ", embedding_vector.shape)

    return 0


test_sentence_transformer_model_local_load()