基于 AMAZON SAGEMAKER 训练一个关键词抽取的模型以及部署

浏览本文需要约 20 分钟，建议按照章节分段阅读。

前言

Amazon Sagemaker 是一个通过完全托管的基础设施、工具和工作流程为任何用例构建、训练和部署机器学习（ML）模型的集成开发环境。

很多行业都有丰富的内容发布平台，最终用户的阅读内容与业务行为之间会有一定的关联性，算法团队因此可以构建一个阅读内容与业务行为的分析模型并理解其中的可解释性。NLP 是这里的主要工作，其中需要借助信息抽取框架来从非结构化的文本中识别跟输出结构化的信息，UIE 就是其中一个框架。

这篇文章我们介绍在 Amazon Sagemaker 上如何使用一个信息抽取的统一的 text-to-structure 生成框架 UIE，实现了对不同信息抽取任务的统一建模。文章主要分为三部分：

第一部分：关于 UIE 的介绍；
第二部分：关于 Amazon Web Services AI/ML 的基础架构；
第三部分：关于使用 notebook 以及 Sagemaker 的资源来做 UIE 的关键词抽取

什么是信息抽取任务

在自然语言处理领域，信息抽取任务（information extraction，下文中简称 IE）是 NLP 一个常见的任务，从非结构化的文本中识别并输出结构化的信息，包括命名实体识别，关系抽取等。而处理海量的文本文件最关键的是要把用户最关心的问题提取出来。而无论是对于长文本还是短文本，往往可以通过一些关键词窥探整个文本的主题思想。与此同时，不管是基于文本的推荐还是基于文本的搜索，对于文本关键词的依赖也很大，关键词提取的准确程度直接关系到推荐系统或者搜索系统的最终效果。因此，关键词提取在文本挖掘领域是一个很重要的部分。

现在常见的框架有 3 个：KeyBert，T5-prompt，UIE。三者的特性如下：

KeyBert

无监督（自监督）
基于语意信息非统计频次
用于无标注场景下的使用，替换传统的 textrank

T5-prompt

需要训练
基于 prompt-learning
可扩展

UIE

小样本
实体抽取预训练模型强大
专业领域表现优秀

这次以 UIE 为例，在 Sagemaker 上用 UIE 做关键词抽取。

信息抽取的使用场景

曾几何时，当我们面对各种复杂多样的信息抽取（IE）任务，我们总会造各式各样 IE 模型的轮子，来满足不同复杂任务的多变需求。

真实的情况是：针对不同任务设定，需要针对特定领域 schema 建模，不同 IE 模型被单个训练、不共享，一个公司可能需要管理众多 IE 模型。

因而，随着 NLP 的发展，统一/通用的 IE 是一个众望所归的模型。

UIE 来自 2022 ACL，它可以做到：

统一地建模不同的 IE 任务
自适应地生成目标结构
从不同的知识来源统一学习通用的信息抽取能力

UIE 原理介绍

UIE 提出的统一生成框架，基于 T5 模型进行了 IE 预训练，在实体、关系、事件和情感等 4 个信息抽取任务、13 个数据集的全监督、低资源和少样本设置下均取得了 SOTA 性能。

设计结构化抽取语言（SEL，Structured Extraction Language）来统一编码易购提取结构，及编码实体，关系，事件统一表示。

构建结构化模式提示器（SSI，Structural Schema Instructor），一个基于 schema 的 prompt 机制，用于控制不同的生成需求。

UIE 的优势

使用简单：用户可以使用自然语言自定义抽取目标，无需训练即可统一抽取输入文本中的对应信息。实现开箱即用，并满足各类信息抽取需求。

降本增效：以往的信息抽取技术需要大量标注数据才能保证信息抽取的效果，为了提高开发过程中的开发效率，减少不必要的重复工作时间，开放域信息抽取可以实现零样本（zero-shot）或者少样本（few-shot）抽取，大幅度降低对标注数据依赖，在降低成本的同时，还提升了效果。

效果领先：开放域信息抽取在多种场景、多种任务上，均有不俗的表现。

机器学习基础架构

在这次的信息抽取任务中，我们使用下面的基础架构来完成 UIE 的相关工作。

Amazon SageMaker 是一项完全托管的机器学习服务。借助 SageMaker 生产就绪的托管环境中，数据科学家和开发人员可以快速、轻松地构建和训练机器学习模型，然后直接将模型部署到生产就绪托管环境中。它提供了一个集成的 Jupyter 编写 Notebook 实例，供您轻松访问数据源以便进行探索和分析，因此您无需管理服务器。此外，它还可以提供常见的机器学习算法，这些算法经过了优化，可以在分布式环境中高效处理非常大的数据。Sagemaker 也可以管理相关的训练任务以及模型。

Amazon Sagemaker Notebook Instance 是一个运行 Jupyter 笔记本应用程序的机器学习（ML）计算实例。SageMaker 管理创建实例和相关资源，在笔记本实例中使用 Jupyter 笔记本来准备和处理数据、编写用于训练模型的代码、将模型部署到 SageMaker 托管以及测试或验证模型。数据都是根据不同的数据以及使用情况（训练数据，测试数据，结果数据等）存储在不同的 S3 桶中。

下图说明了 Amazon 如何训练和部署模型 SageMaker：

使用 Sagemaker Notebook 以及导入项目代码

1. 创建 Notebook Instance，可以根据不同需求创建不同计算大小的实例，在 Sagemaker 的环境中，数据处理、模型训练、模型部署都不在 Notebook Instance 上，所以一般 Notebook Instance 也不需要选用太大的机型，以此也可以节省成本。

2. Notebook Instance 创建完成

3. 在 Notebook Instance 右边打开 JupyterLab

4. Notebook Instance 集成很多不同的环境以及实用工具（包括各种版本的 Python环境、Conda 环境、命令行，甚至有 R 语音以及 PySpark 等的支持），这里可以打开 Command Line Terminal (Launcher – Other – Terminal) 来进行 git 的操作。

5. 从代码长库里复制代码到 Notebook Instance 环境

sh-4.2$ pwd
/home/ec2-user
sh-4.2$ ls
anaconda3  LICENSE                  Nvidia_Cloud_EULA.pdf  SageMaker         sample-notebooks-1667739735  tools
sh-4.2$ cd SageMaker/
sh-4.2$ git clone https://github.com/jackie930/financial-Forecast-RCA.git
Cloning into 'financial-Forecast-RCA'...
remote: Enumerating objects: 612, done.
remote: Counting objects: 100% (612/612), done.
remote: Compressing objects: 100% (500/500), done.
remote: Total 612 (delta 108), reused 587 (delta 86), pack-reused 0
Receiving objects: 100% (612/612), 21.17 MiB | 16.92 MiB/s, done.
Resolving deltas: 100% (108/108), done.

6. 项目代码结构如下。这是一个比较大的项目（包括了 3 个 IE 的框架，SHAP 归因分析等），UIE 只是其中的一部分，这篇文章只集中在 UIE 部分。

sh-4.2$ pwd
Project-	|  --code--	|  --keyword extraction--	|  --KeyBert
       	|		|			|  --UIE
       	|		|			|  --T5-Prompt
        	|		|  --forecast & RCA
        	|		|  --SHAP

UIE 关键词抽取的开发工作

7. 双击 notebook 文件 ~/code/keyword_extraction/UIE/uie_byos_gpu.ipynb，在此就是 notebook 里打开了 UIE 的开发环境了。

这里也可以看到数据的数据的加载部分

WORK_DIRECTORY = './data/'
data_location = sagemaker_session.upload_data(WORK_DIRECTORY, key_prefix=prefix)
inputs = {'training': data_location}

print(inputs) 可以查看 S3 上的数据目录

{'training': 's3://sagemaker-us-east-1-301554342947/sagemaker/DEMO-PaddleNLP'}

数据也会 copy 到 Notebook Instance，可以在目录/UIE/data/里查看，比如这是份标注好的数据例子：

{"content": "5 月 9 日交通费 29 元从北苑到望京 Soho", "result_list": [{"text": "5 月 9 日", "start": 0, "end": 4}], "prompt": "时间"}

{"content": "5 月 9 日交通费 29 元从北苑到望京 Soho", "result_list": [{"text": "29", "start": 7, "end": 9}], "prompt": "费用"}

{"content": "5 月 9 日交通费 29 元从北苑到望京 Soho", "result_list": [{"text": "北苑", "start": 11, "end": 13}], "prompt": "出发地"}

{"content": "5 月 9 日交通费 29 元从北苑到望京 Soho", "result_list": [{"text": "望京 Soho", "start": 14, "end": 18}], "prompt": "目的地"}

{"content": "5 月 17 号晚上 10 点 35 分加班打车回家，36 块五", "result_list": [{"text": "5 月 17 号晚上 10 点 35 分", "start": 0, "end": 13}], "prompt": "时间"}

8. 训练部分，用 Pytorch 的后端来训练 UIE。这里用在 notebook 里声明 ml.g4dn.8xlarge 的 GPU 机型做训练。可以根据不同的训练任务、需不需要 GPU 等来选择不同的机型，而 Amazon Web Services 也提供了各种不同的计算资源来满足需求，比如有通用型（ml.m5d）、CPU 计算密集型（ml.c5）、内存密集型（ml.r5），以及 GPU 加持的实例类型（ml.p3, ml.g4dn, ml.g5）。这些计算资源都是以声明的方式在训练的时候才拉起，训练完成后释放，这种模式使得训练的工作不会长期占有训练机器，以 pay as you go 的方式来使用并支付费用，若任务允许还可以使用 SPOT 实例来做训练，这可以节省比较多的训练费用。

训练任务的 hyperparameters 以及一些基本参数也是在这里声明，比如’learning_rate’: 1e-5，’batch_size’: 16，’max_seq_len’:512，’num_epochs’: 100， ‘model’: ‘uie-base’，以及众多数据路径参数等。

from sagemaker.pytorch import PyTorch
hyperparameters = {'train_path': '/opt/ml/input/data/training/train.txt', 
                   'dev_path': '/opt/ml/input/data/training/dev.txt', 
                   'save_dir': '/opt/ml/model', 
                   'learning_rate': 1e-5, 
                   'batch_size': 16, 
                   'max_seq_len':512, 
                   'num_epochs': 100, 
                   'model': 'uie-base',
                   'seed': 1000,
                   'logging_steps': 10,
                   'valid_steps': 100,
                   'device': 'gpu'}
instance_type = 'ml.g4dn.8xlarge'  # 'ml.p3.2xlarge' or 'ml.p3.8xlarge' or ...
#git_config = {'repo': 'https://github.com/whn09/paddlenlp_sagemaker.git', 'branch': 'main'}
print(datetime.datetime.now())
estimator = PyTorch(entry_point='finetune.py',
                    source_dir='./',
#                             source_dir='model_zoo/uie/',
                            #git_config=git_config,
                            role=role,
                            hyperparameters=hyperparameters,
                            framework_version='1.9.1',
                            py_version='py38',
                            script_mode=True,
                            instance_count=1,  # 1 or 2 or ...
                            instance_type=instance_type)
print(datetime.datetime.now())
estimator.fit(inputs)
print(datetime.datetime.now())

最终会在 output 里看到训练完成：

2022-11-22 08:47:36,390 sagemaker-training-toolkit INFO     Reporting training SUCCESS2022-11-22 08:47:52 Uploading - Uploading generated training model
2022-11-22 08:49:33 Completed - Training job completed
Training seconds: 799
Billable seconds: 799
2022-11-22 08:49:46.907558

9. 训练完成后，就会把模型部署到推理服务器，推理的节点也是声明式的，借助 Amazon SageMaker 可以部署机器学习（ML）模型来进行预测，也称为推理。 SageMaker 提供广泛的 ML 基础架构和模型部署选项，以帮助满足 ML 的推理需求。它是一项完全托管的服务，与 MLOps 工具集成，因此可以扩展模型部署，降低推理成本，在生产中更有效地管理模型并减轻运营负担。在构建和训练了机器学习模型后，可以使用 SageMaker 推理开始从模型中获取预测或推论。借助 SageMaker 推理，可以设置返回推论的端点，也可以从模型中运行批量推理。

instance_type = 'ml.p3.2xlarge'
# predictor = estimator.deploy(initial_instance_count=1, instance_type=instance_type)
from sagemaker.pytorch.model import PyTorchModel
pytorch_model = PyTorchModel(model_data='s3://{}/{}/output/model-inference-gpu.tar.gz'.format(bucket, training_job_name), role=role,entry_point='infer_gpu_cn.py', framework_version='1.9.0', py_version='py38', model_server_workers=4)  # TODO [For GPU], model_server_workers=6
print(datetime.datetime.now())
predictor = pytorch_model.deploy(instance_type=instance_type, initial_instance_count=1)
print(datetime.datetime.now())

会看到部署的进度条如下（当有感叹号时表示部署已完成）：

2022-11-23 09:17:07.803458
--------------------!

这时在 Sagemaker Endpoints 里面也就可以看到通过 Notebook 利用 Sagemaker 部署推理服务如下：

10. 再来做推理：

texts = ['"北京市海淀区人民法院\n民事判决书\n(199x)建初字第xxx号\n原告：张三。\n委托代理人李四，北京市 A律师事务所律师。\n被告：B公司，法定代表人王五，开发公司总经理。\n委托代理人赵六，北京市 C律师事务所律师。"', '原告赵六，2022年5月29日生\n委托代理人孙七，深圳市C律师事务所律师。\n被告周八，1990年7月28日出生\n委托代理人吴九，山东D律师事务所律师']
print(predictor.predict(texts))

实时推理非常适合具有实时、交互式、低延迟要求的推理工作负载。您可以将模型部署到 SageMaker 托管服务并获得可用于推理的终端节点。这些端点是完全托管的，支持自动扩展（参见自动扩展 Amazon SageMaker 模型）。

使用 SageMaker 托管服务部署模型有多种选择。您可以使用 AWS SDK（例如，适用于 SDK for Python (Boto3)、P SageMaker ython SDK 等，以编程方式部署模型，也可以使用 SageMaker 控制台以交互方式部署模型。AWS CLIAWS 软件开发工具包是一个低级 API，支持 Java、C++、Go、Node.js、PHP JavaScript、Ruby 和 Python，而 Pyth SageMaker on SDK 是高级别 Python API。以下文档演示了如何使用 AWS SDK for Python (Boto3) 和 SageMaker Python SDK 部署模型。

如果使用 AWS SDK for Python (Boto3) 或 SageMaker 控制台，则使用 SageMaker 托管服务部署模型的 AWS CLI 过程分为三步：

在 SageMaker 中创建 SageMaker 模型
为 HTTPS 终端节点创建终节点配置
创建一个 HTTPS 终端节点

使用 SageMaker Python SDK 部署模型不需要创建端点配置。因此，这是一个两步过程：

从 Model 类中创建一个可以部署到 HTTPS 端点的模型对象
使用模型对象的预建 deploy() 方法创建 HTTPS 端点

要使用后续代码片段，请将示例代码中替换为新项目信息。

11. 实验最后可以用下面代码清除资源：

# estimator.delete_endpoint()
predictor.delete_endpoint()

也可以在 Amazon Web Services 的 web console 里面删除，请参考官网。

总结

通过上面简单的例子，我们可以快速了解到 UIE 是什么，能做什么，怎么用 Sagemaker 以及 Notebook 方便地进行 UIE 相关模型的开发。但限于篇幅，无法对 UIE 所有的功能进行介绍，其中模型的调优以及完善也需要根据真实项目来推进，有兴趣的读者可以参考 UIE 官方文档进行深入阅读。

Amazon SageMaker 是一个端到端的机器学习平台，SageMaker 支持模型开发全流程，从数据准备、数据处理、算法构建、模型训练、超参调优、模型部署与监控等环节，SageMaker 都提供了相应的功能帮助算法工程师们专注于业务和模型本身，提高开发效率。

Amazon SageMaker 除了提供 Notebook 实例之外，还提供了模型训练和模型部署等功能，算法工程师可以高效、经济的开展机器学习项目。其中训练实例的计费是按秒计费，并且支持 Spot 实例，节省高达 90% 的费用。