mlflow.sentence_transformers

mlflow.sentence_transformers.get_default_conda_env()[source]

Returns: The default Conda environment for MLflow Models produced with the sentence_transformers flavor.

mlflow.sentence_transformers.get_default_pip_requirements() → list[str][source]

Retrieves the set of minimal dependencies for the sentence_transformers flavor.

Returns: A list of default pip requirements for MLflow Models that have been produced with the sentence-transformers flavor. Calls to save_model() and log_model() produce a pip environment that contain these requirements at a minimum.

mlflow.sentence_transformers.load_model(model_uri: str, dst_path: Optional[str] = None)[source]

Note

The ‘sentence_transformers’ MLflow Models integration is known to be compatible with 2.3.1 <= sentence-transformers <= 5.0.0. MLflow Models integrations with sentence_transformers may not succeed when used with package versions outside of this range.

Load a sentence_transformers object from a local file or a run.

Parameters

model_uri –
The location, in URI format, of the MLflow model. For example:
- /Users/me/path/to/local/model
- relative/path/to/local/model
- s3://my_bucket/path/to/model
- runs:/<mlflow_run_id>/run-relative/path/to/model
- mlflow-artifacts:/path/to/model
For more information about supported URI schemes, see Referencing Artifacts.
dst_path – The local filesystem path to utilize for downloading the model artifact. This directory must already exist if provided. If unspecified, a local output path will be created.

Returns

A sentence_transformers model instance

mlflow.sentence_transformers.log_model(model, artifact_path: Optional[str] = None, task: Optional[str] = None, inference_config: Optional[dict[str, typing.Any]] = None, code_paths: Optional[list[str]] = None, registered_model_name: Optional[str] = None, signature: Optional[mlflow.models.signature.ModelSignature] = None, input_example: Optional[Union[pandas.core.frame.DataFrame, numpy.ndarray, dict, list, csr_matrix, csc_matrix, str, bytes, tuple]] = None, await_registration_for=300, pip_requirements: Optional[Union[list[str], str]] = None, extra_pip_requirements: Optional[Union[list[str], str]] = None, conda_env=None, metadata: Optional[dict[str, typing.Any]] = None, prompts: Optional[list[typing.Union[str, Prompt]]] = None, name: Optional[str] = None, params: Optional[dict[str, typing.Any]] = None, tags: Optional[dict[str, typing.Any]] = None, model_type: Optional[str] = None, step: int = 0, model_id: Optional[str] = None)[source]

Note

The ‘sentence_transformers’ MLflow Models integration is known to be compatible with 2.3.1 <= sentence-transformers <= 5.0.0. MLflow Models integrations with sentence_transformers may not succeed when used with package versions outside of this range.

Note

Logging Sentence Transformers models with custom code (i.e. models that require trust_remote_code=True) is supported in MLflow 2.12.0 and above.

Log a sentence_transformers model as an MLflow artifact for the current run.

# An example of using log_model for a sentence-transformers model and architecture:

from sentence_transformers import SentenceTransformer
import mlflow

model = SentenceTransformer("all-MiniLM-L6-v2")
data = "MLflow is awesome!"
signature = mlflow.models.infer_signature(
    model_input=data,
    model_output=model.encode(data),
)

with mlflow.start_run():
    mlflow.sentence_transformers.log_model(
        model=model,
        name="sbert_model",
        signature=signature,
        input_example=data,
    )

Parameters

model – A trained sentence-transformers model.
artifact_path – Deprecated. Use name instead.
task – MLflow inference task type for sentence-transformers model. Candidate task type is llm/v1/embeddings.
inference_config – A dict of valid overrides that can be applied to a sentence-transformer model instance during inference. These arguments are used exclusively for the case of loading the model as a pyfunc Model or for use in Spark. These values are not applied to a returned model from a call to mlflow.sentence_transformers.load_model()
code_paths –
A list of local filesystem paths to Python file dependencies (or directories containing file dependencies). These files are prepended to the system path when the model is loaded. Files declared as dependencies for a given model should have relative imports declared from a common root path if multiple files are defined with import dependencies between them to avoid import errors when loading the model.

For a detailed explanation of code_paths functionality, recommended usage patterns and limitations, see the code_paths usage guide.
registered_model_name – If given, create a model version under registered_model_name, also creating a registered model if one with the given name does not exist.
signature – an instance of the ModelSignature class that describes the model’s inputs and outputs. If not specified but an input_example is supplied, a signature will be automatically inferred based on the supplied input example and model. If both signature and input_example are not specified or the automatic signature inference fails, a default signature will be adopted. To prevent a signature from being adopted, set signature to False. To manually infer a model signature, call infer_signature() on datasets with valid model inputs and valid model outputs.
input_example – one or several instances of valid model input. The input example is used as a hint of what data to feed the model. It will be converted to a Pandas DataFrame and then serialized to json using the Pandas split-oriented format, or a numpy array where the example will be serialized to json by converting it to a list. Bytes are base64-encoded. When the signature parameter is None, the input example is used to infer a model signature.
await_registration_for – Number of seconds to wait for the model version to finish being created and is in READY status. By default, the function waits for five minutes. Specify 0 or None to skip waiting.
pip_requirements – Either an iterable of pip requirement strings (e.g. ["sentence_transformers", "-r requirements.txt", "-c constraints.txt"]) or the string path to a pip requirements file on the local filesystem (e.g. "requirements.txt"). If provided, this describes the environment this model should be run in. If None, a default list of requirements is inferred by mlflow.models.infer_pip_requirements() from the current software environment. If the requirement inference fails, it falls back to using get_default_pip_requirements. Both requirements and constraints are automatically parsed and written to requirements.txt and constraints.txt files, respectively, and stored as part of the model. Requirements are also written to the pip section of the model’s conda environment (conda.yaml) file.
extra_pip_requirements –
Either an iterable of pip requirement strings (e.g. ["pandas", "-r requirements.txt", "-c constraints.txt"]) or the string path to a pip requirements file on the local filesystem (e.g. "requirements.txt"). If provided, this describes additional pip requirements that are appended to a default set of pip requirements generated automatically based on the user’s current software environment. Both requirements and constraints are automatically parsed and written to requirements.txt and constraints.txt files, respectively, and stored as part of the model. Requirements are also written to the pip section of the model’s conda environment (conda.yaml) file.
Warning

The following arguments can’t be specified at the same time:
- conda_env
- pip_requirements
- extra_pip_requirements
This example demonstrates how to specify pip requirements using pip_requirements and extra_pip_requirements.
conda_env –
Either a dictionary representation of a Conda environment or the path to a conda environment yaml file. If provided, this describes the environment this model should be run in. At a minimum, it should specify the dependencies contained in get_default_conda_env(). If None, a conda environment with pip requirements inferred by mlflow.models.infer_pip_requirements() is added to the model. If the requirement inference fails, it falls back to using get_default_pip_requirements. pip requirements from conda_env are written to a pip requirements.txt file and the full conda environment is written to conda.yaml. The following is an example dictionary representation of a conda environment:
```
{
    "name": "mlflow-env",
    "channels": ["conda-forge"],
    "dependencies": [
        "python=3.8.15",
        {
            "pip": [
                "sentence_transformers==x.y.z"
            ],
        },
    ],
}
```
metadata – Custom metadata dictionary passed to the model and stored in the MLmodel file.

prompts –

A list of prompt URIs registered in the MLflow Prompt Registry, to be associated with the model. Each prompt URI should be in the form prompt:/<name>/<version>. The prompts should be registered in the MLflow Prompt Registry before being associated with the model.

This will create a mutual link between the model and the prompt. The associated prompts can be seen in the model’s metadata stored in the MLmodel file. From the Prompt Registry UI, you can navigate to the model as well.

import mlflow

prompt_template = "Hi, {name}! How are you doing today?"

# Register a prompt in the MLflow Prompt Registry
mlflow.prompts.register_prompt("my_prompt", prompt_template, description="A simple prompt")

# Log a model with the registered prompt
with mlflow.start_run():
    model_info = mlflow.pyfunc.log_model(
        name=MyModel(),
        name="model",
        prompts=["prompt:/my_prompt/1"]
    )

print(model_info.prompts)
# Output: ['prompt:/my_prompt/1']

# Load the prompt
prompt = mlflow.genai.load_prompt(model_info.prompts[0])

name – Model name.
params – A dictionary of parameters to log with the model.
tags – A dictionary of tags to log with the model.
model_type – The type of the model.
step – The step at which to log the model outputs and metrics
model_id – The ID of the model.

mlflow.sentence_transformers.save_model(model, path: str, task: Optional[str] = None, inference_config: Optional[dict[str, typing.Any]] = None, code_paths: Optional[list[str]] = None, mlflow_model: Optional[mlflow.models.model.Model] = None, signature: Optional[mlflow.models.signature.ModelSignature] = None, input_example: Optional[Union[pandas.core.frame.DataFrame, numpy.ndarray, dict, list, csr_matrix, csc_matrix, str, bytes, tuple]] = None, pip_requirements: Optional[Union[list[str], str]] = None, extra_pip_requirements: Optional[Union[list[str], str]] = None, conda_env=None, metadata: Optional[dict[str, typing.Any]] = None) → None[source]

Note

The ‘sentence_transformers’ MLflow Models integration is known to be compatible with 2.3.1 <= sentence-transformers <= 5.0.0. MLflow Models integrations with sentence_transformers may not succeed when used with package versions outside of this range.

Note

Saving Sentence Transformers models with custom code (i.e. models that require trust_remote_code=True) is supported in MLflow 2.12.0 and above.

Save a trained sentence-transformers model to a path on the local file system.

Parameters

model – A trained sentence-transformers model.
path – Local path destination for the serialized model to be saved.
task – MLflow inference task type for sentence-transformers model. Candidate task type is llm/v1/embeddings.
inference_config – A dict of valid inference parameters that can be applied to a sentence-transformer model instance during inference. These arguments are used exclusively for the case of loading the model as a pyfunc Model or for use in Spark. These values are not applied to a returned model from a call to mlflow.sentence_transformers.load_model()
code_paths –
A list of local filesystem paths to Python file dependencies (or directories containing file dependencies). These files are prepended to the system path when the model is loaded. Files declared as dependencies for a given model should have relative imports declared from a common root path if multiple files are defined with import dependencies between them to avoid import errors when loading the model.

For a detailed explanation of code_paths functionality, recommended usage patterns and limitations, see the code_paths usage guide.
mlflow_model – An MLflow model object that specifies the flavor that this model is being added to.
signature – an instance of the ModelSignature class that describes the model’s inputs and outputs. If not specified but an input_example is supplied, a signature will be automatically inferred based on the supplied input example and model. If both signature and input_example are not specified or the automatic signature inference fails, a default signature will be adopted. To prevent a signature from being adopted, set signature to False. To manually infer a model signature, call infer_signature() on datasets with valid model inputs and valid model outputs.
input_example – one or several instances of valid model input. The input example is used as a hint of what data to feed the model. It will be converted to a Pandas DataFrame and then serialized to json using the Pandas split-oriented format, or a numpy array where the example will be serialized to json by converting it to a list. Bytes are base64-encoded. When the signature parameter is None, the input example is used to infer a model signature.
pip_requirements – Either an iterable of pip requirement strings (e.g. ["sentence_transformers", "-r requirements.txt", "-c constraints.txt"]) or the string path to a pip requirements file on the local filesystem (e.g. "requirements.txt"). If provided, this describes the environment this model should be run in. If None, a default list of requirements is inferred by mlflow.models.infer_pip_requirements() from the current software environment. If the requirement inference fails, it falls back to using get_default_pip_requirements. Both requirements and constraints are automatically parsed and written to requirements.txt and constraints.txt files, respectively, and stored as part of the model. Requirements are also written to the pip section of the model’s conda environment (conda.yaml) file.
extra_pip_requirements –
Either an iterable of pip requirement strings (e.g. ["pandas", "-r requirements.txt", "-c constraints.txt"]) or the string path to a pip requirements file on the local filesystem (e.g. "requirements.txt"). If provided, this describes additional pip requirements that are appended to a default set of pip requirements generated automatically based on the user’s current software environment. Both requirements and constraints are automatically parsed and written to requirements.txt and constraints.txt files, respectively, and stored as part of the model. Requirements are also written to the pip section of the model’s conda environment (conda.yaml) file.
Warning

The following arguments can’t be specified at the same time:
- conda_env
- pip_requirements
- extra_pip_requirements
This example demonstrates how to specify pip requirements using pip_requirements and extra_pip_requirements.
conda_env –
Either a dictionary representation of a Conda environment or the path to a conda environment yaml file. If provided, this describes the environment this model should be run in. At a minimum, it should specify the dependencies contained in get_default_conda_env(). If None, a conda environment with pip requirements inferred by mlflow.models.infer_pip_requirements() is added to the model. If the requirement inference fails, it falls back to using get_default_pip_requirements. pip requirements from conda_env are written to a pip requirements.txt file and the full conda environment is written to conda.yaml. The following is an example dictionary representation of a conda environment:
```
{
    "name": "mlflow-env",
    "channels": ["conda-forge"],
    "dependencies": [
        "python=3.8.15",
        {
            "pip": [
                "sentence_transformers==x.y.z"
            ],
        },
    ],
}
```
metadata – Custom metadata dictionary passed to the model and stored in the MLmodel file.