将 Transformer 模型部署为兼容 OpenAI 的聊天机器人

下载此笔记本

欢迎阅读我们的教程，该教程介绍如何使用 Transformers 和 MLflow 创建兼容 OpenAI 的聊天模型。在 MLflow 2.11 及更高版本中，MLflow 的 Transformers flavor 支持特殊的任务类型 llm/v1/chat，这会将 Hugging Face 上的数千个文本生成模型转换为与 OpenAI 模型互操作的对话聊天机器人。这使您可以无缝地替换聊天应用程序的后端 LLM，或者轻松评估不同的模型，而无需编辑客户端代码。

如果您还没有看过，您可能会发现浏览我们的关于聊天和 Transformers 的入门笔记本很有帮助，然后再继续学习本教程，因为本教程的级别稍高，并且不会深入探讨 Transformers 或 MLflow Tracking 的内部工作原理。

注意：本页面介绍如何将 Transformers 模型部署为聊天机器人。如果您使用的是不同的框架或自定义 python 模型，请改用 ChatModel 来构建兼容 OpenAI 的聊天机器人。

学习目标

在本教程中，您将

使用 TinyLLama-1.1B-Chat 创建兼容 OpenAI 的聊天模型
将模型记录到 MLflow 并重新加载它以进行本地推理。
使用 MLflow 模型服务提供模型

%pip install mlflow>=2.11.0 -q -U
# OpenAI-compatible chat model support is available for Transformers 4.34.0 and above
%pip install transformers>=4.34.0 -q -U

# Disable tokenizers warnings when constructing pipelines
%env TOKENIZERS_PARALLELISM=false

import warnings

# Disable a few less-than-useful UserWarnings from setuptools and pydantic
warnings.filterwarnings("ignore", category=UserWarning)

env: TOKENIZERS_PARALLELISM=false

构建聊天模型

MLflow 的原生 Transformers 集成允许您在保存或记录管道时指定 task 参数。最初，此参数接受任何 Transformers 管道任务类型，但是 mlflow.transformers flavor 为 text-generation 管道类型添加了一些特定于 MLflow 的键。

对于 text-generation 管道，您可以提供两个符合 MLflow AI Gateway 的 endpoint_type 规范的字符串文字，而不是将 text-generation 指定为任务类型（"llm/v1/embeddings" 可以指定为使用 mlflow.sentence_transformers 保存的模型的任务）

"llm/v1/chat" 用于聊天式应用程序
"llm/v1/completions" 用于通用完成

指定这些键之一后，MLflow 将自动处理服务聊天或完成模型所需的一切。这包括

在模型上设置聊天/完成兼容签名
执行数据预处理和后处理，以确保输入和输出符合聊天/完成 API 规范，该规范与 OpenAI 的 API 规范兼容。

请注意，这些修改仅在模型通过 mlflow.pyfunc.load_model() 加载时才适用（例如，在使用 mlflow models serve CLI 工具提供模型时）。如果只想加载基本管道，则始终可以通过 mlflow.transformers.load_model() 进行加载。

在接下来的几个单元格中，我们将学习如何使用本地 Transformers 管道和 MLflow 提供聊天模型，并以 TinyLlama-1.1B-Chat 为例。

首先，让我们回顾一下保存文本生成管道的原始流程

from transformers import pipeline

import mlflow

generator = pipeline(
  "text-generation",
  model="TinyLlama/TinyLlama-1.1B-Chat-v1.0",
)

# save the model using the vanilla `text-generation` task type
mlflow.transformers.save_model(
  path="tinyllama-text-generation", transformers_model=generator, task="text-generation"
)

/var/folders/qd/9rwd0_gd0qs65g4sdqlm51hr0000gp/T/ipykernel_55429/4268198845.py:11: FutureWarning: The 'transformers' MLflow Models integration is known to be compatible with the following package version ranges: ``4.25.1`` -  ``4.37.1``. MLflow Models integrations with transformers may not succeed when used with package versions outside of this range.
mlflow.transformers.save_model(

现在，让我们加载模型并将其用于推理。我们加载的模型是一个 text-generation 管道，让我们看一下它的签名，以查看其期望的输入和输出。

# load the model for inference
model = mlflow.pyfunc.load_model("tinyllama-text-generation")

model.metadata.signature

Loading checkpoint shards:   0%|          | 0/5 [00:00<?, ?it/s]

2024/02/26 21:06:51 WARNING mlflow.transformers: Could not specify device parameter for this pipeline type

Loading checkpoint shards:   0%|          | 0/5 [00:00<?, ?it/s]

inputs: 
[string (required)]
outputs: 
[string (required)]
params: 
None

不幸的是，它仅接受 string 作为输入，这与聊天界面并不直接兼容。例如，在与 OpenAI 的 API 交互时，我们希望能够简单地输入消息列表。为了使用我们当前的模型执行此操作，我们将需要编写一些额外的样板代码

# first, apply the tokenizer's chat template, since the
# model is tuned to accept prompts in a chat format. this
# also converts the list of messages to a string.
messages = [{"role": "user", "content": "Write me a hello world program in python"}]
prompt = generator.tokenizer.apply_chat_template(
  messages, tokenize=False, add_generation_prompt=True
)

model.predict(prompt)

['<|user|>
Write me a hello world program in python</s>
<|assistant|>
Here's a simple hello world program in Python:

```python
print("Hello, world!")
```

This program prints the string "Hello, world!" to the console. You can run this program by typing it into the Python interpreter or by running the command `python hello_world.py` in your terminal.']

现在我们已经取得了一些进展，但是在推理之前格式化我们的消息非常麻烦。

此外，输出格式也不与 OpenAI API 规范兼容 - 它只是一个字符串列表。如果我们正在寻找为我们的聊天应用程序评估不同的模型后端，我们将必须重写一些客户端代码，既要格式化输入，又要解析这个新响应。

为了简化所有这些，让我们在保存模型时仅传入 "llm/v1/chat" 作为 task 参数。

# save the model using the `"llm/v1/chat"`
# task type instead of `text-generation`
mlflow.transformers.save_model(
  path="tinyllama-chat", transformers_model=generator, task="llm/v1/chat"
)

/var/folders/qd/9rwd0_gd0qs65g4sdqlm51hr0000gp/T/ipykernel_55429/609241782.py:3: FutureWarning: The 'transformers' MLflow Models integration is known to be compatible with the following package version ranges: ``4.25.1`` -  ``4.37.1``. MLflow Models integrations with transformers may not succeed when used with package versions outside of this range.
mlflow.transformers.save_model(

再次，让我们加载模型并检查签名

model = mlflow.pyfunc.load_model("tinyllama-chat")

model.metadata.signature

Loading checkpoint shards:   0%|          | 0/5 [00:00<?, ?it/s]

2024/02/26 21:10:04 WARNING mlflow.transformers: Could not specify device parameter for this pipeline type

Loading checkpoint shards:   0%|          | 0/5 [00:00<?, ?it/s]

inputs: 
['messages': Array({content: string (required), name: string (optional), role: string (required)}) (required), 'temperature': double (optional), 'max_tokens': long (optional), 'stop': Array(string) (optional), 'n': long (optional), 'stream': boolean (optional)]
outputs: 
['id': string (required), 'object': string (required), 'created': long (required), 'model': string (required), 'choices': Array({finish_reason: string (required), index: long (required), message: {content: string (required), name: string (optional), role: string (required)} (required)}) (required), 'usage': {completion_tokens: long (required), prompt_tokens: long (required), total_tokens: long (required)} (required)]
params: 
None

现在，在执行推理时，我们可以像与 OpenAI API 交互时一样，在字典中传递我们的消息。此外，我们从模型收到的响应也符合规范。

messages = [{"role": "user", "content": "Write me a hello world program in python"}]

model.predict({"messages": messages})

[{'id': '8435a57d-9895-485e-98d3-95b1cbe007c0',
'object': 'chat.completion',
'created': 1708949437,
'model': 'TinyLlama/TinyLlama-1.1B-Chat-v1.0',
'usage': {'prompt_tokens': 24, 'completion_tokens': 71, 'total_tokens': 95},
'choices': [{'index': 0,
  'finish_reason': 'stop',
  'message': {'role': 'assistant',
   'content': 'Here's a simple hello world program in Python:

```python
print("Hello, world!")
```

This program prints the string "Hello, world!" to the console. You can run this program by typing it into the Python interpreter or by running the command `python hello_world.py` in your terminal.'}}]}]

提供聊天模型

为了进一步说明这个例子，让我们使用 MLflow 来提供我们的聊天模型，以便我们可以像 Web API 一样与它进行交互。为此，我们可以使用 mlflow models serve CLI 工具。

在终端 shell 中，运行

$ mlflow models serve -m tinyllama-chat

服务器完成初始化后，您应该能够通过 HTTP 请求与模型进行交互。输入格式与 MLflow Deployments Server 文档中描述的格式几乎相同，但 temperature 默认为 1.0 而不是 0.0。

这是一个快速示例

%%sh
curl http://127.0.0.1:5000/invocations   -H 'Content-Type: application/json'   -d '{ "messages": [{"role": "user", "content": "Write me a hello world program in python"}] }'   | jq

  % Total    % Received % Xferd  Average Speed   Time    Time     Time  Current
                               Dload  Upload   Total   Spent    Left  Speed
100   706  100   617  100    89     25      3  0:00:29  0:00:23  0:00:06   160

[
{
  "id": "fc3d08c3-d37d-420d-a754-50f77eb32a92",
  "object": "chat.completion",
  "created": 1708949465,
  "model": "TinyLlama/TinyLlama-1.1B-Chat-v1.0",
  "usage": {
    "prompt_tokens": 24,
    "completion_tokens": 71,
    "total_tokens": 95
  },
  "choices": [
    {
      "index": 0,
      "finish_reason": "stop",
      "message": {
        "role": "assistant",
        "content": "Here's a simple hello world program in Python:

```python
print("Hello, world!")
```

This program prints the string "Hello, world!" to the console. You can run this program by typing it into the Python interpreter or by running the command `python hello_world.py` in your terminal."
      }
    }
  ]
}
]

就这么简单！

您还可以使用一些可选的推理参数来调用 API，以调整模型的响应。这些映射到 Transformers 管道参数，并在推理时直接传入。

max_tokens（映射到 max_new_tokens）：模型应生成的新 token 的最大数量。
temperature（映射到 temperature）：控制模型响应的创造性。请注意，并非所有模型都保证支持此参数，并且为了使此参数生效，必须使用 do_sample=True 创建管道。
stop（映射到 stopping_criteria）：停止生成的 token 列表。

注意：n 没有等效的 Transformers 管道参数，并且在查询中不受支持。但是，您可以使用自定义 Pyfunc 实现一个消耗 n 参数的模型（详见下文）。

结论

在本教程中，您学习了如何通过在保存 Transformers 管道时指定 "llm/v1/chat" 作为任务来创建兼容 OpenAI 的聊天模型。

下一步是什么？

了解自定义 ChatModel。如果您正在寻找更深层次的定制或 Transformers 之外的模型，则链接的页面提供了有关如何使用 MLflow 的 ChatModel 类构建聊天机器人的实践指导。
更多关于 MLflow AI Gateway 的信息。在本教程中，我们看到了如何使用本地服务器部署模型，但是 MLflow 提供了许多其他方法来将模型部署到生产环境。查看此页面以了解有关不同选项的更多信息。
更多关于 MLflow 的 Transformers 集成的信息。本页面提供了关于 MLflow 的 Transformers 集成的全面概述，以及大量的实践指南和笔记本。了解如何微调模型、使用提示模板等！
其他 LLM 集成。除了 Transformers 之外，MLflow 还集成了许多其他流行的 LLM 库，例如 Langchain 和 OpenAI。

学习目标​

构建聊天模型​

提供聊天模型​

结论​

下一步是什么？​

学习目标

构建聊天模型

提供聊天模型

结论

下一步是什么？