wehub-resource-sync
54 lines
1.5 KiB
Python
54 lines
1.5 KiB
Python
import os
|
|
|
|
import numpy as np
|
|
import openai
|
|
|
|
import mlflow
|
|
from mlflow.models.signature import ModelSignature
|
|
from mlflow.types.schema import ColSpec, ParamSchema, ParamSpec, Schema, TensorSpec
|
|
|
|
assert "OPENAI_API_KEY" in os.environ, " OPENAI_API_KEY environment variable must be set"
|
|
|
|
|
|
print(
|
|
"""
|
|
# ******************************************************************************
|
|
# Text embeddings
|
|
# ******************************************************************************
|
|
"""
|
|
)
|
|
|
|
with mlflow.start_run():
|
|
model_info = mlflow.openai.log_model(
|
|
model="text-embedding-ada-002",
|
|
task=openai.embeddings,
|
|
name="model",
|
|
)
|
|
|
|
model = mlflow.pyfunc.load_model(model_info.model_uri)
|
|
print(model.predict(["hello", "world"]))
|
|
|
|
|
|
print(
|
|
"""
|
|
# ******************************************************************************
|
|
# Text embeddings with batch_size parameter
|
|
# ******************************************************************************
|
|
"""
|
|
)
|
|
|
|
with mlflow.start_run():
|
|
mlflow.openai.log_model(
|
|
model="text-embedding-ada-002",
|
|
task=openai.embeddings,
|
|
name="model",
|
|
signature=ModelSignature(
|
|
inputs=Schema([ColSpec(type="string", name=None)]),
|
|
outputs=Schema([TensorSpec(type=np.dtype("float64"), shape=(-1,))]),
|
|
params=ParamSchema([ParamSpec(name="batch_size", dtype="long", default=1024)]),
|
|
),
|
|
)
|
|
|
|
model = mlflow.pyfunc.load_model(model_info.model_uri)
|
|
print(model.predict(["hello", "world"], params={"batch_size": 16}))
|