在 Databricks 中使用 DBRX 的 AutoGen

DBRX 发布

2024年3月，Databricks 发布了DBRX，这是一个通用的 LLM，为开放的 LLM 设定了新的标准。虽然 DBRX 在 Hugging Face 上作为开源模型可用（databricks/dbrx-instruct 和 databricks/dbrx-base），但 Databricks 的客户也可以通过 Foundation Model APIs 来使用 DBRX，这使得 DBRX 可以通过与 OpenAI 兼容的自动缩放 REST API 来使用。

Autogen 正在成为创建代理的流行标准。它被设计为支持任何实现了 OpenAI SDK 的“LLM即服务”，并且可以轻松扩展以与强大的开源模型集成。

本笔记本将演示使用 DBRX 的 Autogen 的一些基本示例，包括使用 AssistantAgent、UserProxyAgent 和 ConversableAgent。这些演示并不旨在详尽无遗 - 您可以将它们作为基础进行扩展！

要求

在您的 Databricks 集群上必须安装 AutoGen，并且需要 Python>=3.8。本示例包括 %pip 魔术命令来安装：%pip install pyautogen，以及其他必要的库。

此代码已在以下环境中进行了测试：* 无服务器笔记本（截至2024年4月18日仍处于公共预览阶段）* Databricks Runtime 14.3 LTS ML 文档

此代码可以在任何支持通过按需付费 API（或预配吞吐量）使用 DBRX 的 Databricks 工作区中运行。要检查您的区域是否受支持，请参阅Foundation Model 区域可用性。如果以上条件成立，则工作区还必须由管理员启用 Foundation Model APIs。 文档.

提示

• 这个笔记本可以从 GitHub 导入到 Databricks 工作区并直接运行。使用 git 的稀疏检出模式 只导入这个笔记本或示例目录。

• Databricks 建议使用 Secrets 来替代以明文存储令牌。

贡献者

tj@databricks.com (Github: tj-cycyota)

%pip install pyautogen==0.2.25 openai==1.21.2 typing_extensions==4.11.0 --upgrade

建议在安装后重新启动 Python 内核 - 取消注释并运行以下代码：

# dbutils.library.restartPython()

设置 DBRX 配置列表

有关 config_list 的使用，请参阅 Autogen 文档：LLM 配置

import os

# 使用当前工作区主机和个人访问令牌设置环境变量
# 对于您已经设置的机密：dbutils.secrets.get('your_scope_name','databricks_host')

# DATABRICKS_HOST 格式: "https://{your workspace url}"（无尾部斜杠）
## AWS 工作区示例: "https://my-databricks-workspace.cloud.databricks.com"
## Azure 工作区示例: "https://adb-123456790123.12.azuredatabricks.net"
os.environ["DATABRICKS_HOST"] = "<填写您支持的 DBRX 区域中的工作区 URL>"

# DATABRICKS_TOKEN 格式: "dapi..."
## 临时令牌: dbutils.notebook.entry_point.getDbutils().notebook().getContext().apiToken().get()
os.environ["DATABRICKS_TOKEN"] = "dapi...."

llm_config = {
    "config_list": [
        {
            "model": "databricks-dbrx-instruct",
            "api_key": str(os.environ["DATABRICKS_TOKEN"]),
            "base_url": str(os.getenv("DATABRICKS_HOST")) + "/serving-endpoints",
        }
    ],
}

Hello World 示例

我们的第一个示例将使用一个简单的 UserProxyAgent 向一个 AssistantAgent 提问。这是基于教程演示 here。

发送问题并看到回答后，您可以输入 exit 结束聊天或继续对话。

import autogen

# 创建 Assistant 和 User
assistant = autogen.AssistantAgent(name="assistant", llm_config=llm_config)

user_proxy = autogen.UserProxyAgent(name="user", code_execution_config=False)

# 从 user_proxy 方面发起聊天
chat_result = user_proxy.initiate_chat(assistant, message="What is MLflow?")

用户（对助手）：

什么是MLflow？

--------------------------------------------------------------------------------
助手（对用户）：

当然，我很乐意为您解释MLflow。MLflow是一个用于管理机器学习工作流的开源平台。它由Databricks开发，并于2018年开源。MLflow提供了许多功能，帮助数据科学家和机器学习工程师管理端到端的机器学习生命周期，包括：

1. **MLflow Tracking**：这是一个记录和查询实验的日志API，包括代码、数据、配置和结果。
2. **MLflow Projects**：这是一种打包可重用和可复现的数据科学代码的格式，可以在不同的平台上运行。
3. **MLflow Models**：这是一种将机器学习模型打包成多种格式的约定，使其易于在不同环境中部署。
4. **MLflow Model Registry**：这是一个中央存储库，用于存储、管理和提供机器学习模型。

这是一个使用MLflow Tracking记录简单实验的Python代码示例：
```python
# 文件名：mlflow_example.py

import mlflow
import numpy as np

# 记录一个参数（例如，随机森林中的树的数量）
mlflow.log_param("num_trees", 100)

# 记录一个指标（例如，模型的准确率）
accuracy = np.random.rand()
mlflow.log_metric("accuracy", accuracy)

# 记录模型
mlflow.sklearn.log_model(model, "model")

# 结束运行
mlflow.end_run()
```
要运行这段代码，您需要在您的机器上安装并运行 MLflow。您可以使用 pip 安装 MLflow：
```
pip install mlflow
```
然后，您可以使用以下命令运行代码：
```
python mlflow_example.py
```
这将在 MLflow 中创建一个新的实验，并记录参数、指标和模型。然后，您可以在 MLflow UI 中查看实验。

希望这对您有所帮助！如果您有任何其他问题，请告诉我。

--------------------------------------------------------------------------------
# 文件名：count_primes.py

def is_prime(n):
    if n <= 1:
        return False
    if n <= 3:
        return True
    if n % 2 == 0 or n % 3 == 0:
        return False
    i = 5
    while i * i <= n:
        if n % i == 0 or n % (i + 2) == 0:
            return False
        i += 6
    return True

def count_primes(end):
    count = 0
    for num in range(1, end + 1):
        if is_prime(num):
            count += 1
    return count

print(count_primes(10000))
```

这个脚本首先定义了一个辅助函数 `is_prime(n)`，用于检查一个数 `n` 是否为质数。然后，它定义了主函数 `count_primes(end)`，用于计算从 1 到 `end` 之间的质数的数量。最后，它打印出从 1 到 10000 的质数的数量。

要运行这个脚本，将其保存到一个名为 `count_primes.py` 的文件中，然后在终端或命令提示符中使用命令 `python count_primes.py` 执行它。脚本将打印出从 1 到 10000 的质数的数量。

如果在运行脚本时遇到任何问题或错误，请告诉我。

--------------------------------------------------------------------------------

向 DBRX 助手提供反馈。按回车键跳过并使用自动回复，或输入 'exit' 结束对话：  

>>>>>>>> 未收到人类输入。

>>>>>>>> 使用自动回复...

>>>>>>>> 执行代码块（推测语言为 python）...
用户（对 DBRX 助手）：

退出码：0（执行成功）
代码输出：1229


--------------------------------------------------------------------------------
DBRX 助手（对用户）：

太棒了！执行成功，输出为 1229，这是从 1 到 10000 的质数的数量。这意味着在 1 到 10000 的范围内有 1229 个质数。

为了验证结果，我们可以手动检查在 1 到 10000 的范围内的质数。以下是前几个质数：

2、3、5、7、11、13、17、19、23、29、31、37、41、43、47、53、59、61、67、71、73、79、83、89、97、101、103、107、109、113、127、131、137、139、149、151、157、163、167、169、173、179、181、191、193、197、199...

我们可以看到，在 1 到 10000 的范围内确实有 1229 个质数。

因此，脚本的输出是正确的，任务已完成。

终止

--------------------------------------------------------------------------------

向 DBRX 助手提供反馈。按回车键跳过并使用自动回复，或输入 'exit' 结束对话：  exit

我们可以看到在我们的工作目录中创建了一个 Python 文件：

%sh ls coding

count_primes.py

%sh head coding/count_primes.py

# filename: count_primes.py

def is_prime(n):
    if n <= 1:
        return False
    if n <= 3:
        return True
    if n % 2 == 0 or n % 3 == 0:
        return False
    i = 5

使用DBRX在教师/学生交流中进行“自言自语”的对话模式：

from autogen import ConversableAgent

# 设置代理
student_agent = ConversableAgent(
    name="Student_Agent",
    system_message="你是一名愿意学习的学生。",
    llm_config=llm_config,
)

teacher_agent = ConversableAgent(
    name="Teacher_Agent",
    system_message="你是一名计算机科学教师。",
    llm_config=llm_config,
)

# 开始对话
chat_result = student_agent.initiate_chat(
    teacher_agent,
    message="深度学习与人工智能有什么关系？",
    summary_method="last_msg",
    max_turns=1,  # 设置更高的数字以控制来回对话
)

学生代理（对教师代理）：

深度学习与人工智能有什么关系？

--------------------------------------------------------------------------------

>>>>>>>> 使用自动回复...
教师代理（对学生代理）：

你好！很高兴你问到深度学习和人工智能之间的关系。

深度学习实际上是人工智能的一个子集，人工智能是一个更广泛的领域，致力于创建能够执行通常需要人类智能的任务的算法和系统。人工智能的其他子集包括基于规则的系统、机器学习、自然语言处理和计算机视觉等。

另一方面，深度学习是一种特定的构建人工智能系统的方法，它受到人脑结构和功能的启发。在深度学习中，我们使用人工神经网络，它由相互连接的节点或“神经元”组成，来学习数据中的模式并进行预测或决策，而无需明确编程。

近年来，深度学习取得了特别成功，原因包括大量数据的可用性、强大的计算资源以及神经网络架构和训练算法的进步。因此，深度学习在图像和语音识别、自然语言处理和游戏对战等各种任务中取得了最先进的性能。

因此，总结起来，深度学习是一种特定的构建人工智能系统的方法，在近年来因其在各种任务中的出色表现而受到广泛关注。然而，它只是构建人工智能系统的众多方法之一，要充分理解深度学习的潜力和局限性，还需要了解人工智能的更广泛背景。

--------------------------------------------------------------------------------

实现日志显示

在调试时，将聊天记录显示在笔记本中并将这些记录持久化到Delta表中可能很有用。以下部分演示了如何扩展默认的AutoGen日志库。 首先，我们将实现一个扩展了 autogen.runtime_logging 功能的 Python class。 文档：

class Databricks_AutoGenLogger:
    def __init__(self):
        from pyspark.sql import SparkSession
        import autogen

        self.spark = SparkSession.builder.getOrCreate()
        self.logger_config = {"dbname": "logs.db"}

    def start(self):
        import autogen.runtime_logging

        self.logging_session_id = autogen.runtime_logging.start(config=self.logger_config)
        print("Logging session ID: " + str(self.logging_session_id))

    def stop(self):
        import autogen.runtime_logging

        autogen.runtime_logging.stop()

    def _get_log(self, dbname="logs.db", table="chat_completions"):
        import sqlite3

        con = sqlite3.connect(dbname)
        query = f"SELECT * from {table} WHERE session_id == '{self.logging_session_id}' ORDER BY end_time DESC"
        cursor = con.execute(query)
        rows = cursor.fetchall()
        column_names = [description[0] for description in cursor.description]
        data = [dict(zip(column_names, row)) for row in rows]
        con.close()
        return data

    def display_session(self):
        import pandas as pd

        return pd.DataFrame(self._get_log())

    def persist_results(self, target_delta_table: str, mode="append"):
        import pandas as pd

        # 转换为 Spark DF
        sdf = self.spark.createDataFrame(pd.DataFrame(self._get_log()))

        try:
            sdf.write.format("delta").mode(mode).saveAsTable(target_delta_table)
            print(f"日志已成功写入表 {target_delta_table}，写入模式为 {mode}")
        except Exception as e:
            print(f"发生错误：{e}")

让我们在最简单的示例中使用上面的类。请注意添加了日志记录的 .start() 和 .stop()，以及用于错误处理的 try/except。

assistant = autogen.AssistantAgent(name="assistant", llm_config=llm_config)
user_proxy = autogen.UserProxyAgent(name="user", code_execution_config=False)

# 在开始对话之前，启动日志记录：
logs = Databricks_AutoGenLogger()
logs.start()
try:
    user_proxy.initiate_chat(assistant, message="What is MLflow?", max_turns=1)
except Exception as e:
    print(f"发生错误：{e}")
logs.stop()
# 显示日志
display(logs.display_session())

记录会话 ID: 6c389f5f-3619-4762-8118-bc98dd414f90 用户（对助手）：

什么是 MLflow？

---

助手（对用户）：

当然，我很乐意为您解释 MLflow。MLflow 是一个用于管理机器学习工作流的开源平台。它由 Databricks 开发，并于2018年开源。MLflow 提供了一系列功能，帮助数据科学家和机器学习工程师管理端到端的机器学习生命周期，包括：

1. MLflow Tracking：这是一个记录和查询实验的日志 API，包括代码、数据、配置和结果。
2. MLflow Projects：这是一种打包可重用和可复现的数据科学代码的格式，可以在不同平台上运行。
3. MLflow Models：这是一种将机器学习模型打包成多种格式的约定，使其可以轻松部署在不同的环境中。
4. MLflow Model Registry：这是一个中央存储库，用于存储、管理和提供机器学习模型。

这是一个使用 MLflow Tracking 记录简单实验的 Python 代码示例：

# 文件名：mlflow_example.py

import mlflow
import numpy as np

# 记录一个参数（例如，随机森林中的树的数量）
mlflow.log_param("num_trees", 100)

# 记录一个指标（例如，模型的准确率）
accuracy = np.random.rand()
mlflow.log_metric("accuracy", accuracy)

# 记录模型
mlflow.sklearn.log_model(model, "model")

# 结束运行
mlflow.end_run()

要运行这段代码，您需要在您的机器上安装并运行 MLflow。您可以使用 pip 安装 MLflow：

pip install mlflow

然后，您可以使用以下命令运行代码：

python mlflow_example.py

这将在 MLflow 中创建一个新的实验，并记录参数、指标和模型。然后，您可以在 MLflow UI 中查看实验。

希望这对您有所帮助！如果您有任何其他问题，请告诉我。

---

<div dangerouslySetInnerHTML={{ __html: quartoRawHtml[0] }} />

通过这个，我们有了一个简单的框架来审查和保存我们的聊天记录！请注意，在上面的 `request` 字段中，我们还可以看到 LLM 的系统提示 - 这对于提示工程和调试都很有用。

请注意，当您将其部署到 Databricks Model Serving 时，模型的响应会自动使用 [Lakehouse Monitoring](https://docs.databricks.com/en/lakehouse-monitoring/index.html) 进行记录；但上述方法提供了一个简单的机制来从**客户端**记录聊天记录。

现在，让我们将这些结果持久化到 [Unity Catalog](https://docs.databricks.com/en/data-governance/unity-catalog/index.html) 中的 Delta 表中：

```python
from pyspark.sql import SparkSession

spark = SparkSession.builder.getOrCreate()  # 在 Databricks 中不需要；笔记本中预先配置了会话

# 使用 3 层命名空间：catalog.schema.table。如果表不存在，将创建该表。
target_delta_table = "your_catalog.your_schema.autogen_logs"
logs.persist_results(target_delta_table=target_delta_table, mode="append")

# 显示表中当前的行
display(spark.table(target_delta_table))
```

``` text
成功将日志写入表 shared.tjc.autogent_logs，以追加模式。
```

<div dangerouslySetInnerHTML={{ __html: quartoRawHtml[1] }} />

## 总结

本笔记本提供了使用 DBRX 的 Autogen 的一些基本示例，我们很期待看到您如何将这个框架与领先的开源 LLM 结合使用！

### 限制

-   Databricks Foundation Model API 支持其他开源 LLM（Mixtral、Llama2 等），但上述代码尚未在这些 LLM 上进行测试。

-   截至 2024 年 4 月，DBRX 尚不支持工具/函数调用功能。要进一步讨论此功能，请联系您的 Databricks 代表。