wehub-resource-sync
34 lines
1018 B
Python
34 lines
1018 B
Python
|
|
""" Tests that embedding model is loaded and yielding a structurally correct embedding vector. """
|
|
|
|
|
|
from llmware.models import ModelCatalog
|
|
|
|
|
|
def test_embedding_model_local_load():
|
|
|
|
emb_models = ModelCatalog().list_embedding_models()
|
|
|
|
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_card in emb_models:
|
|
|
|
if model_card["model_family"] in ["HFEmbeddingModel"]:
|
|
|
|
print(f"\nloading model - {model_card['model_name']} - embedding dims - {model_card['embedding_dims']}")
|
|
|
|
model = ModelCatalog().load_model(model_card["model_name"])
|
|
|
|
embedding_vector = model.embedding(test_text)
|
|
|
|
assert embedding_vector is not None
|
|
|
|
print(f"created vector successfully with dimensions: ", embedding_vector[0].shape)
|
|
|
|
assert embedding_vector[0].shape[0] == model_card['embedding_dims']
|
|
|
|
return 0
|
|
|
|
|