Sentence Transformers¶

The Snowflake ML Model Registry supports Sentence Transformer models (sentence_transformers.SentenceTransformer). Sentence Transformers are used to generate dense vector embeddings for text, which can be used for semantic search, clustering, and other natural language processing tasks.

The following additional options can be used in the options dictionary when you call log_model:

Option

Description

target_methods

A list of the names of the methods available on the model object. Sentence Transformer models only support encode as the target method, which is the default.

batch_size

The batch size to use when encoding text. Defaults to 32.

cuda_version

The version of the CUDA runtime to be used when deploying to a platform with GPU; defaults to 11.8. If manually set to None, the model cannot be deployed to a platform having a GPU.

The sample_input_data parameter is optional when logging a Sentence Transformer model. If not provided, the registry will infer the signature automatically based on the model’s encode method.

Note

Sentence Transformer models only support the encode target method. Input must be a single column containing text strings, and output is a single column containing embedding vectors as arrays.

Examples¶

These examples assume reg is an instance of snowflake.ml.registry.Registry.

Basic Sentence Transformer¶

The following example demonstrates the key steps to load a Sentence Transformer model, log it to the Snowflake ML Model Registry, and use the registered model to generate embeddings. The workflow includes:

  • Loads a pre-trained Sentence Transformer model.

  • Logs the model to the Snowflake ML Model Registry.

  • Generates text embeddings.

from sentence_transformers import SentenceTransformer
import pandas as pd

# Load a pre-trained Sentence Transformer model
model = SentenceTransformer('all-MiniLM-L6-v2')

# Sample input data as DataFrame
sample_df = pd.DataFrame({"sentences": [
    "How do I reset my password?",
    "What are your business hours?",
    "I need help with my recent order.",
]})

# Log the model
model_ref = reg.log_model(
    model=model,
    model_name="my_sentence_transformer",
    version_name="v1",
    sample_input_data=sample_df,
)

# Generate embeddings
test_df = pd.DataFrame({"sentences": [
    "I forgot my login credentials.",
    "When is your store open?",
]})
embeddings_df = model_ref.run(test_df, function_name="encode")
Copy

Custom Batch Size¶

You can specify a custom batch size for encoding:

from sentence_transformers import SentenceTransformer
import pandas as pd

# Load model
model = SentenceTransformer('all-mpnet-base-v2')

# Sample input data as DataFrame
sample_df = pd.DataFrame({"text": ["Sample text for embedding."]})

# Log the model with custom batch size
model_ref = reg.log_model(
    model=model,
    model_name="my_mpnet_model",
    version_name="v1",
    sample_input_data=sample_df,
    options={"batch_size": 64},
)

# Generate embeddings
test_df = pd.DataFrame({"text": [
    "First sentence to encode.",
    "Second sentence to encode.",
    "Third sentence to encode.",
]})
embeddings_df = model_ref.run(test_df, function_name="encode")
Copy
Language: English