跳到主要内容

将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 特性为 text-generation 管道类型添加了几个 MLflow 特有的键。

对于 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() 来完成。

在接下来的几个单元中,我们将学习如何使用本地 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" 作为任务参数传入。

# 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)集成。