1  from sentence_transformers import SentenceTransformer
 2  
 3  import mlflow
 4  import mlflow.sentence_transformers
 5  
 6  model = SentenceTransformer("all-MiniLM-L6-v2")
 7  
 8  example_sentences = ["This is a sentence.", "This is another sentence."]
 9  
10  # Define the signature
11  signature = mlflow.models.infer_signature(
12      model_input=example_sentences,
13      model_output=model.encode(example_sentences),
14  )
15  
16  # Log the model using mlflow
17  with mlflow.start_run():
18      logged_model = mlflow.sentence_transformers.log_model(
19          model=model,
20          name="sbert_model",
21          signature=signature,
22          input_example=example_sentences,
23      )
24  
25  # Load option 1: mlflow.pyfunc.load_model returns a PyFuncModel
26  loaded_model = mlflow.pyfunc.load_model(logged_model.model_uri)
27  embeddings1 = loaded_model.predict(["hello world", "i am mlflow"])
28  
29  # Load option 2: mlflow.sentence_transformers.load_model returns a SentenceTransformer
30  loaded_model = mlflow.sentence_transformers.load_model(logged_model.model_uri)
31  embeddings2 = loaded_model.encode(["hello world", "i am mlflow"])
32  
33  print(embeddings1)
34  
35  """
36  >> [[-3.44772562e-02  3.10232025e-02  6.73496164e-03  2.61089969e-02
37    ...
38    2.37922110e-02 -2.28897743e-02  3.89375277e-02  3.02067865e-02]
39   [ 4.81191138e-03 -9.33756605e-02  6.95968643e-02  8.09735525e-03
40    ...
41     6.57437667e-02 -2.72239652e-02  4.02687863e-02 -1.05599344e-01]]
42  """