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

下载此笔记本

欢迎参加我们的教程，学习如何使用 Transformers 和 MLflow 创建一个 OpenAI 兼容的聊天模型。在 MLflow 2.11 及更高版本中，MLflow 的 Transformers 功能支持特殊的任务类型 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 功能添加了一些 MLflow 特定的键，用于 text-generation 管道类型。

对于 text-generation 管道，您可以使用两个字符串字面量之一来代替指定 text-generation 作为任务类型，这两个字面量符合MLflow AI Gateway 的 endpoint_type 规范（“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() 加载。

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

首先，让我们通过原始流程来保存文本生成管道

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 部署服务器文档中描述的格式相同，不同之处在于 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）集成。

学习目标​

构建聊天模型​

部署聊天模型​

结论​

下一步是什么？​

学习目标

构建聊天模型

部署聊天模型

结论

下一步是什么？