Skip to main content

Databricks

Databricks 数据湖平台将数据、分析和人工智能统一在一个平台上。

这个示例笔记本展示了如何将 Databricks 端点封装为 LangChain 中的 LLMs。 它支持两种端点类型：

  • 服务端点，推荐用于生产和开发，
  • 集群驱动程序代理应用，推荐用于交互式开发。
from langchain.llms import Databricks

封装一个服务端点​

先决条件：

预期的 MLflow 模型签名是：

  • 输入：[{"name": "prompt", "type": "string"}, {"name": "stop", "type": "list[string]"}]
  • 输出：[{"type": "string"}]

如果模型签名不兼容或者您想插入额外的配置，可以相应地设置 transform_input_fn 和 transform_output_fn。

# 如果在 "single user" 或 "no isolation shared" 模式下运行附加到交互式集群的 Databricks 笔记本，
# 您只需要指定端点名称即可创建一个 `Databricks` 实例，以查询同一工作区中的服务端点。
llm = Databricks(endpoint_name="dolly")

llm("你好吗？")
    '很高兴听到你身体健康，一如既往地感谢你。'
llm("你好吗？", stop=["。"])
    '好'
# 否则，您可以手动指定 Databricks 工作区主机名和个人访问令牌，
# 或分别设置 `DATABRICKS_HOST` 和 `DATABRICKS_TOKEN` 环境变量。
# 请参阅 https://docs.databricks.com/dev-tools/auth.html#databricks-personal-access-tokens
# 我们强烈建议不要在笔记本中明确暴露 API 令牌。
# 您可以使用 Databricks 密钥管理器来安全地存储 API 令牌。
# 请参阅 https://docs.databricks.com/dev-tools/databricks-utils.html#secrets-utility-dbutilssecrets

import os

os.environ["DATABRICKS_TOKEN"] = dbutils.secrets.get("myworkspace", "api_token")

llm = Databricks(host="myworkspace.cloud.databricks.com", endpoint_name="dolly")

llm("你好吗？")
    '我很好。谢谢！'
# 如果服务端点接受额外的参数，比如 `temperature`，
# 您可以在 `model_kwargs` 中设置它们。
llm = Databricks(endpoint_name="dolly", model_kwargs={"temperature": 0.1})

llm("你好吗？")
    '我很好。'
# 如果服务端点期望不同的输入模式并且不返回 JSON 字符串，
# 或者您想在其上应用提示模板，则使用 `transform_input_fn` 和 `transform_output_fn`。


def transform_input(**request):
    full_prompt = f"""{request["prompt"]}
    言简意赅。
    """
    request["prompt"] = full_prompt
    return request


llm = Databricks(endpoint_name="dolly", transform_input_fn=transform_input)

llm("你好吗？")
    '我很好。你呢？'

封装一个集群驱动程序代理应用​

先决条件：

  • 在 "single user" 或 "no isolation shared" 模式下的 Databricks 交互式集群上加载了一个 LLM。
  • 在驱动程序节点上运行一个本地 HTTP 服务器，使用 HTTP POST 和 JSON 输入/输出来为模型提供服务，路径为 "/"。
  • 它使用一个介于 [3000, 8000] 之间的端口号，并监听驱动程序 IP 地址或仅仅是 0.0.0.0，而不是仅限于 localhost。
  • 您对集群具有 "Can Attach To" 权限。

预期的服务器模式（使用 JSON 模式）是：

  • 输入：
    {"type": "object",
     "properties": {
        "prompt": {"type": "string"},
         "stop": {"type": "array", "items": {"type": "string"}}},
      "required": ["prompt"]}
  • 输出：{"type": "string"}

如果服务器模式不兼容或者您想插入额外的配置，可以使用 transform_input_fn 和 transform_output_fn。

以下是一个运行驱动程序代理应用来提供 LLM 服务的最小示例：

from flask import Flask, request, jsonify
import torch
from transformers import pipeline, AutoTokenizer, StoppingCriteria

model = "databricks/dolly-v2-3b"
tokenizer = AutoTokenizer.from_pretrained(model, padding_side="left")
dolly = pipeline(model=model, tokenizer=tokenizer, trust_remote_code=True, device_map="auto")
device = dolly.device

class CheckStop(StoppingCriteria):
    def __init__(self, stop=None):
        super().__init__()
        self.stop = stop or []
        self.matched = ""
        self.stop_ids = [tokenizer.encode(s, return_tensors='pt').to(device) for s in self.stop]
    def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor, **kwargs):
        for i, s in enumerate(self.stop_ids):
            if torch.all((s == input_ids[0][-s.shape[1]:])).item():
                self.matched = self.stop[i]
                return True
        return False

def llm(prompt, stop=None, **kwargs):
  check_stop = CheckStop(stop)
  result = dolly(prompt, stopping_criteria=[check_stop], **kwargs)
  return result[0]["generated_text"].rstrip(check_stop.matched)

app = Flask("dolly")

@app.route('/', methods=['POST'])
def serve_llm():
  resp = llm(**request.json)
  return jsonify(resp)

app.run(host="0.0.0.0", port="7777")

一旦服务器运行起来，您可以创建一个 Databricks 实例来封装它作为一个 LLM。

# 如果在运行应用的集群上附加了一个 Databricks 笔记本，
# 您只需要指定驱动程序端口即可创建一个 `Databricks` 实例。
llm = Databricks(cluster_driver_port="7777")

llm("你好吗？")
    '你好，谢谢你的关心。很高兴听到你身体健康。'
# 否则，您可以手动指定要使用的集群 ID，
# 以及 Databricks 工作区主机名和个人访问令牌。

llm = Databricks(cluster_id="0000-000000-xxxxxxxx", cluster_driver_port="7777")

llm("你好吗？")
    '我很好。你呢？'
# 如果应用接受额外的参数，比如 `temperature`，
# 您可以在 `model_kwargs` 中设置它们。
llm = Databricks(cluster_driver_port="7777", model_kwargs={"temperature": 0.1})

llm("你好吗？")
    '我很好。很高兴认识你。'
# 如果应用期望不同的输入模式并且不返回 JSON 字符串，
# 或者您想在其上应用提示模板，则使用 `transform_input_fn` 和 `transform_output_fn`。


def transform_input(**request):
    full_prompt = f"""{request["prompt"]}
    言简意赅。
    """
    request["prompt"] = full_prompt
    return request


def transform_output(response):
    return response.upper()


llm = Databricks(
    cluster_driver_port="7777",
    transform_input_fn=transform_input,
    transform_output_fn=transform_output,
)

llm("你好吗？")
    '我很好，谢谢。'