Use um modelo de incorporação personalizado

O exemplo a seguir mostra como log um modelo Python personalizado para que Databricks o utilize na geração de embeddings.

O esquema de entrada deve ser uma única string ColSpec e o esquema de saída deve ser TensorSpec como sua assinatura. Os embeddings devem ser retornados em um array 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
    )