跳到内容

使用 Databricks 和 instructor 获取结构化输出,一份完整指南

Databricks 提供了一个人工智能平台,可访问各种模型。本指南展示了如何将 instructor 与 Databricks 结合使用以获取结构化输出。

快速入门

首先,安装所需的软件包

pip install instructor

你需要一个 Databricks API 密钥和工作区 URL,可以将其设置为环境变量

export DATABRICKS_API_KEY=your_api_key_here
export DATABRICKS_HOST=your_workspace_url

基本示例

以下是如何从 Databricks 模型中提取结构化数据

import os
import instructor
from openai import OpenAI
from pydantic import BaseModel

# Initialize the client with Databricks base URL
client = instructor.from_openai(
    OpenAI(
        base_url="https://your-databricks-workspace-url/serving-endpoints/your-endpoint-name/invocations",
        api_key=os.environ["DATABRICKS_API_KEY"],
    ),
    mode=instructor.Mode.TOOLS,
)

# Define your data structure
class UserExtract(BaseModel):
    name: str
    age: int

# Extract structured data
user = client.chat.completions.create(
    model="databricks-model", # Your model name in Databricks
    response_model=UserExtract,
    messages=[
        {"role": "user", "content": "Extract jason is 25 years old"},
    ],
)

print(user)
# Output: UserExtract(name='Jason', age=25)

支持的模式

Databricks 支持与 OpenAI 相同的模式

  • Mode.TOOLS
  • Mode.JSON
  • Mode.FUNCTIONS
  • Mode.PARALLEL_TOOLS
  • Mode.MD_JSON
  • Mode.TOOLS_STRICT
  • Mode.JSON_O1

模型

Databricks 提供了对各种模型的访问,具体取决于你的设置,包括

  • 托管在 Databricks 上的基础模型
  • 自定义微调模型
  • 部署在 Databricks 上的开源模型