wehub-resource-sync
57 lines
1.8 KiB
Python
57 lines
1.8 KiB
Python
import os
|
|
|
|
import openai
|
|
|
|
import mlflow
|
|
from mlflow.models.signature import ModelSignature
|
|
from mlflow.types.schema import ColSpec, ParamSchema, ParamSpec, Schema
|
|
|
|
assert "OPENAI_API_KEY" in os.environ, " OPENAI_API_KEY environment variable must be set"
|
|
|
|
print(
|
|
"""
|
|
# ******************************************************************************
|
|
# Completions indicating prompt template
|
|
# ******************************************************************************
|
|
"""
|
|
)
|
|
|
|
with mlflow.start_run():
|
|
model_info = mlflow.openai.log_model(
|
|
model="text-davinci-002",
|
|
task=openai.completions,
|
|
name="model",
|
|
prompt="Classify the following tweet's sentiment: '{tweet}'.",
|
|
)
|
|
|
|
model = mlflow.pyfunc.load_model(model_info.model_uri)
|
|
print(model.predict(["I believe in a better world"]))
|
|
|
|
|
|
print(
|
|
"""
|
|
# ******************************************************************************
|
|
# Completions using inference parameters
|
|
# ******************************************************************************
|
|
"""
|
|
)
|
|
with mlflow.start_run():
|
|
model_info = mlflow.openai.log_model(
|
|
model="text-davinci-002",
|
|
task=openai.completions,
|
|
name="model",
|
|
prompt="Classify the following tweet's sentiment: '{tweet}'.",
|
|
signature=ModelSignature(
|
|
inputs=Schema([ColSpec(type="string", name=None)]),
|
|
outputs=Schema([ColSpec(type="string", name=None)]),
|
|
params=ParamSchema([
|
|
ParamSpec(name="max_tokens", default=16, dtype="long"),
|
|
ParamSpec(name="temperature", default=0, dtype="float"),
|
|
ParamSpec(name="best_of", default=1, dtype="long"),
|
|
]),
|
|
),
|
|
)
|
|
|
|
model = mlflow.pyfunc.load_model(model_info.model_uri)
|
|
print(model.predict(["I believe in a better world"], params={"temperature": 1, "best_of": 5}))
|