与 Databricks LLM API 集成。

前置条件

• Databricks 个人访问令牌，用于查询和访问 Databricks 模型服务端点。

• 位于受支持区域的 Databricks 工作区，用于按 token 付费的 Foundation Model API。

设置

如果你在 Colab 中打开这个 Notebook，可能需要先安装 LlamaIndex 🦙。

% pip install llama-index-llms-databricks

!pip install llama-index

from llama_index.llms.databricks import Databricks

None of PyTorch, TensorFlow >= 2.0, or Flax have been found. Models won't be available and only tokenizers, configuration and file/data utilities can be used.

终端窗口

export DATABRICKS_TOKEN=<your api key>
export DATABRICKS_SERVING_ENDPOINT=<your api serving endpoint>

或者，你也可以在初始化 LLM 时传入 API 密钥和服务端点：

llm = Databricks(
    model="databricks-dbrx-instruct",
    api_key="your_api_key",
    api_base="https://[your-work-space].cloud.databricks.com/serving-endpoints/",
)

可用的 LLM 模型列表见此处。

response = llm.complete("Explain the importance of open source LLMs")

print(response)

使用消息列表调用 chat

from llama_index.core.llms import ChatMessage

messages = [
    ChatMessage(
        role="system", content="You are a pirate with a colorful personality"
    ),
    ChatMessage(role="user", content="What is your name"),
]
resp = llm.chat(messages)

print(resp)

流式

使用 stream_complete 端点

response = llm.stream_complete("Explain the importance of open source LLMs")

for r in response:
    print(r.delta, end="")

使用 stream_chat 端点

from llama_index.core.llms import ChatMessage

messages = [
    ChatMessage(
        role="system", content="You are a pirate with a colorful personality"
    ),
    ChatMessage(role="user", content="What is your name"),
]
resp = llm.stream_chat(messages)

for r in resp:
    print(r.delta, end="")