カスタム埋め込みモデルを使用する

次の例は、埋め込みを生成するために Databricks が使用するカスタム Python モデルをログに記録する方法を示しています。

入力スキーマは単一のColSpec文字列である必要があり、出力スキーマは署名としてTensorSpecである必要があります。埋め込みは NumPy 配列で返される必要があります。

Python

class CustomEmbeddingModel(PythonModel):

    def load_context(self, context):
        import os
        self.databricks_token = os.environ.get('DATABRICKS_TOKEN')
        self.base_url = os.environ.get("BASE_URL")
        self.model_name = "databricks-gte-large-en"

    def predict(self, context, model_input, params=None):
        from openai import OpenAI
        import numpy as np
        client = OpenAI(api_key=self.databricks_token, base_url=self.base_url)
        embeddings = client.embeddings.create(
          input=model_input.iloc[:, 0], model=self.model_name)
        results = np.stack([e.embedding for e in embeddings.data])
        return results

with mlflow.start_run() as run:
    input_schema = Schema([ColSpec(DataType.string)])
    output_schema = Schema([TensorSpec(np.dtype(np.float32), [-1, 1024])])
    signature = ModelSignature(inputs=input_schema, outputs=output_schema)
    mlflow.pyfunc.log_model(
        artifact_path="model",
        python_model=CustomEmbeddingModel(),
        pip_requirements=["mlflow==2.21.3", "openai==1.69.0"],
        signature=signature,
        input_example=input_example
    )