Skip to main content

Invoke a model

This recipe invokes a model hosted on Mosaic AI Model Serving using the Databricks SDK for Python and returns the result. Choose either a traditional ML model or a large language model (LLM).

Code snippets​

Traditional Machine Learning​

Using dataframe_split (JSON-serialized DataFrame in split orientation)​

app.py
from databricks.sdk import WorkspaceClient

w = WorkspaceClient()

response = w.serving_endpoints.query(
    name="custom-regression-model",
    dataframe_split={
        "columns": ["feature1", "feature2"],
        "data": [[1.5, 2.5]]
    }
)

Using dataframe_records (JSON-serialized DataFrame in records orientation)​

app.py
from databricks.sdk import WorkspaceClient

w = WorkspaceClient()

response = w.serving_endpoints.query(
    name="custom-regression-model",
    dataframe_records={
        "feature1": [1.5],
        "feature2": [2.5]
    }
)

Using instances (Tensor inputs in row format for TensorFlow/PyTorch models)​

app.py
from databricks.sdk import WorkspaceClient

w = WorkspaceClient()

tensor_input = [[1.0, 2.0, 3.0]]
response = w.serving_endpoints.query(
    name="tensor-processing-model",
    instances=tensor_input,
)

Using inputs (Tensor inputs in columnar format for TensorFlow/PyTorch models)​

app.py
from databricks.sdk import WorkspaceClient

w = WorkspaceClient()

tensor_input = {
    "input1": [1.0, 2.0, 3.0],
    "input2": [4.0, 5.0, 6.0],
}
response = w.serving_endpoints.query(
    name="tensor-processing-model",
    inputs=tensor_input,
)

Large language models (LLMs)​

Using prompt (Input text for completion tasks)​

app.py
from databricks.sdk import WorkspaceClient

w = WorkspaceClient()

response = w.serving_endpoints.query(
    name="llm-text-completions-model",
    prompt="Generate a recipe for building scalable Databricks Apps.",
    temperature=0.5,
)

Using messages (List of chat messages for conversational models)​

app.py
from databricks.sdk import WorkspaceClient
from databricks.sdk.service.serving import ChatMessage, ChatMessageRole

w = WorkspaceClient()

response = w.serving_endpoints.query(
    name="chat-assistant-model",
    messages=[
        ChatMessage(
            role=ChatMessageRole.SYSTEM,
            content="You are a helpful assistant.",
        ),
        ChatMessage(
            role=ChatMessageRole.USER,
            content="Provide tips for deploying Databricks Apps.",
        ),
    ],
)

Using input (Input text for embedding tasks)​

app.py
from databricks.sdk import WorkspaceClient

w = WorkspaceClient()

response = w.serving_endpoints.query(
    name="embedding-model",
    input="Databricks provides a unified analytics platform.",
)

Resources​

Permissions​

Your app service principal needs the following permissions:

  • CAN QUERY on the model serving endpoint

See Manage permissions on your model serving endpoint for more information.

Dependencies​

requirements.txt
databricks-sdk
dash