wehub-resource-sync
78 lines
2.4 KiB
Python
78 lines
2.4 KiB
Python
|
|
""" 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()
|