如何在 Riva 上部署使用 NeMo 训练的自定义声学模型 (Citrinet)#

本教程将引导您完成在 NVIDIA Riva 上部署使用 NVIDIA NeMo 训练的自定义声学模型 (Citrinet) 的过程。

NVIDIA Riva 概述#

NVIDIA Riva 是一个 GPU 加速的 SDK，用于构建针对您的用例进行定制并提供实时性能的语音 AI 应用程序。
Riva 提供丰富的语音和自然语言理解服务，例如

自动语音识别 (ASR)。
文本到语音合成 (TTS)。
自然语言处理 (NLP) 服务集合，例如命名实体识别 (NER)、标点符号和意图分类。

在本教程中，我们将部署一个在 Riva 上使用 NeMo 训练的自定义声学模型 (Citrinet)。
要了解 Riva ASR API 的基础知识，请参阅 Python 中 Riva ASR 入门。

有关 Riva 的更多信息，请参阅 Riva 产品页面和 Riva 开发者文档。

NeMo (神经模块) 和 `nemo2riva`#

NVIDIA NeMo 是一个开源框架，用于构建、训练和微调 GPU 加速的语音 AI 和自然语言理解 (NLU) 模型，并具有简单的 Python 界面。要使用 NeMo 微调 Citrinet 声学模型，请参阅 Citrinet 微调教程。

nemo2riva 命令行工具提供了导出 .nemo 模型的功能，该模型可以使用 NVIDIA Riva 进行部署，NVIDIA Riva 是一个高性能应用程序框架，用于使用 GPU 的多模态对话式 AI 服务。nemo2riva 的 Python .whl 文件包含在 Riva 快速入门资源文件夹中。您还可以使用 pip 安装 nemo2riva，如 Citrinet 微调教程中所示。

本教程探讨如何获取 .riva 模型 — 调用 nemo2riva CLI 工具的结果 (请参阅 Citrinet 微调教程) — 并利用 Riva ServiceMaker 框架聚合 Riva 部署到目标环境所需的所有必要工件。一旦模型部署在 Riva 中，您就可以向服务器发出推理请求。我们将演示整个过程是多么快速和直接。在本教程中，您将学习如何

使用 Riva ServiceMaker 从 .riva 文件构建 .rmir 模型管道。
在本地 Riva 服务器上部署模型。
使用 Riva API 绑定从演示客户端发送推理请求。

先决条件#

在我们开始之前，请确保您拥有

访问 NVIDIA NGC 并能够下载 Riva 快速入门资源。
您要部署的 .riva 模型文件。您可以使用 nemo2riva 命令将 .nemo 模型文件转换为 .riva 模型文件。有关使用 NeMo 自定义 Citrinet 声学模型并使用 nemo2riva 导出结果模型的信息，请参阅 Citrinet 微调教程。或者，您可以从此处获取英语 ASR 的预训练 Citrinet .riva 模型。

Riva ServiceMaker#

Riva ServiceMaker 是一组工具，用于聚合 Riva 部署到目标环境所需的所有必要工件（模型、文件、配置和用户设置）。它有两个主要组件

Riva-Build#

此步骤有助于构建 Riva 就绪的模型版本。它唯一的输出是 Riva 中支持服务的端到端管道的中间格式（称为 RMIR）。让我们考虑一个 ASR Citrinet 模型。

riva-build 负责将一个或多个导出的模型（.riva 文件）组合成一个包含中间格式的文件，称为 Riva 模型中间表示 (.rmir)。此文件包含整个端到端管道的部署无关规范，以及最终部署和推理所需的所有资产。有关更多信息，请参阅文档。

# IMPORTANT: UPDATE THESE PATHS 

# Riva Docker container
RIVA_CONTAINER = "<add container name>"

# Example: 
# RIVA_CONTAINER = f"nvcr.io/nvidia/riva/riva-speech:{__riva_version__}"

# Directory where the .riva model is stored $MODEL_LOC/*.riva
MODEL_LOC = "<add path to model location>"

# Name of the .riva file
MODEL_NAME = "<add model name>"

# Get the Riva docker
! docker pull $RIVA_CONTAINER

如果 MODEL_LOC 中尚不存在子目录，请创建一个子目录来存储您的 .rmir 文件。

! mkdir -p $MODEL_LOC/rmir

构建 `.rmir` 文件。#

注释

如果您在调用 nemo2riva 时通过添加 --key 标志加密了您的声学模型和/或语言模型，或者您从 NGC 下载了 2023 年之前的预训练模型，则需要在 riva-build 命令中在模型名称后附加一个冒号，然后是密钥的值，如下所示。您可能会发现设置名为 KEY 的字符串变量并将其作为 $KEY 传递到相应的 riva-build 参数中很方便。旧预训练模型的标准加密密钥是 tlt_encode。
如果您想为美国英语以外的受支持语言构建 ASR 管道，请参阅 Riva ASR 管道配置文档。要获得特定应用程序的正确 riva-build 参数，请从页面第一部分底部的交互式 Web 菜单中选择声学模型、语言和管道类型（本教程的目的为离线）。

# Syntax: riva-build <task-name> output-dir-for-rmir/model.rmir[:key] dir-for-riva/model.riva[:key]
! docker run --rm --gpus 1 -v $MODEL_LOC:/data $RIVA_CONTAINER -- \
    riva-build speech_recognition \
        /data/rmir/asr_offline_citrinet.rmir \
        /data/$MODEL_NAME \
        --offline \
        --streaming=False \
        --name=asr_offline_citrinet_pipeline \
        --chunk_size=300 \
        --left_padding_size=0. \
        --right_padding_size=0. \
        --ms_per_timestep=80 \
        --greedy_decoder.asr_model_delay=-1 \
        --featurizer.use_utterance_norm_params=False \
        --featurizer.precalc_norm_time_steps=0 \
        --featurizer.precalc_norm_params=False \
        --decoder_type=greedy \
        --language_code=en-US

Riva-Deploy#

部署工具将一个或多个 RMIR 文件和一个目标模型存储库目录作为输入。它创建一个整体配置，指定执行的管道，最后将所有这些资产写入输出模型存储库目录。

注意： 如果您在使用 riva-build 构建 .rmir 文件时向其添加了加密密钥，请确保在 riva-deploy 命令中在模型名称后附加一个冒号，然后是密钥的值，如下所示。

# Syntax: riva-deploy -f dir-for-rmir/model.rmir[:key] output-dir-for-repository
! docker run --rm --gpus 0 -v $MODEL_LOC:/data $RIVA_CONTAINER -- \
    riva-deploy -f /data/rmir/asr_offline_citrinet.rmir /data/models/

启动 Riva 服务器#

生成模型存储库后，我们就可以启动 Riva 服务器了。首先，从 NGC 下载 Riva 快速入门资源。在此处设置目录的路径

# Set the Riva Quick Start directory
RIVA_DIR = "<Path to the uncompressed folder downloaded from quickstart(include the folder name)>"

接下来，我们修改 config.sh 文件以启用相关的 Riva 服务 (Citrinet 模型的 ASR)、提供加密密钥和模型存储库的路径 (riva_model_loc) (在上一步中生成)，以及其他配置。

例如，如果上面的模型存储库在 $MODEL_LOC/models 生成，那么您可以将 riva_model_loc 指定为与 MODEL_LOC 相同的目录。

models_asr/nlp/tts 中指定的模型的预训练版本是从 NGC 获取的。由于我们使用的是自定义模型，我们可以在 models_asr 中注释它（以及任何其他与您的用例无关的模型）。

config.sh 代码片段#

# Enable or Disable Riva Services
service_enabled_asr=true 
service_enabled_nlp=true # MAKE CHANGES HERE - SET TO FALSE
service_enabled_tts=true # MAKE CHANGES HERE - SET TO FALSE
service_enabled_nmt=true # MAKE CHANGES HERE - SET TO FALSE

...

# Locations to use for storing models artifacts
#
# If an absolute path is specified, the data will be written to that location
# Otherwise, a Docker volume will be used (default).
#
# riva_init.sh will create a `rmir` and `models` directory in the volume or
# path specified.
#
# RMIR ($riva_model_loc/rmir)
# Riva uses an intermediate representation (RMIR) for models
# that are ready to deploy but not yet fully optimized for deployment. Pretrained
# versions can be obtained from NGC (by specifying NGC models below) and will be
# downloaded to $riva_model_loc/rmir by `riva_init.sh`
#
# Custom models produced by NeMo or TLT and prepared using riva-build
# may also be copied manually to this location $(riva_model_loc/rmir).
#
# Models ($riva_model_loc/models)
# During the riva_init process, the RMIR files in $riva_model_loc/rmir
# are inspected and optimized for deployment. The optimized versions are
# stored in $riva_model_loc/models. The riva server exclusively uses these
# optimized versions.
riva_model_loc="riva-model-repo"  # MAKE CHANGES HERE - REPLACE WITH $MODEL_LOC

if [[ $riva_target_gpu_family == "tegra" ]]; then
    riva_model_loc="`pwd`/model_repository"
fi

# The default RMIRs are downloaded from NGC by default in the above $riva_rmir_loc directory
# If you'd like to skip the download from NGC and use the existing RMIRs in the $riva_rmir_loc
# then set the below $use_existing_rmirs flag to true. You can also deploy your set of custom
# RMIRs by keeping them in the riva_rmir_loc dir and use this quickstart script with the
# below flag to deploy them all together.
use_existing_rmirs=false        # MAKE CHANGES HERE - set to true                    

注意 在继续之前，请确保执行以下操作

或者手动执行这些任务

在 Jupyter Lab 的文件导航器中，导航到 $RIVA_DIR 并打开 config.sh
按照上面的代码片段配置设置
- 将 NLP、TTS 和 NMT 服务设置为 false
- 将 riva_model_loc 路径设置为也分配给 MODEL_LOC 的路径
- 将变量 use_existing_rmirs 设置为 true

或者运行下面的单元格

ENABLE_ASR = 'true'
ENABLE_NLP = 'false'
ENABLE_TTS = 'false'
ENABLE_NMT = 'false'

!sed -i "s|service_enabled_asr=.*|service_enabled_asr=$ENABLE_ASR|g" $RIVA_DIR/config.sh
!sed -i "s|service_enabled_nlp=.*|service_enabled_nlp=$ENABLE_NLP|g" $RIVA_DIR/config.sh
!sed -i "s|service_enabled_tts=.*|service_enabled_tts=$ENABLE_TTS|g" $RIVA_DIR/config.sh
!sed -i "s|service_enabled_nmt=.*|service_enabled_nmt=$ENABLE_NMT|g" $RIVA_DIR/config.sh

!sed -i "/\sriva_model_loc=.*/! s|riva_model_loc=.*|riva_model_loc=\"$MODEL_LOC\"|g" $RIVA_DIR/config.sh

!sed -i "s|use_existing_rmirs=.*|use_existing_rmirs=true|g" $RIVA_DIR/config.sh

# Ensure you have permission to execute these scripts
! cd $RIVA_DIR && chmod +x ./riva_init.sh && chmod +x ./riva_start.sh

# Run Riva Init. This will fetch the containers/models
# YOU CAN SKIP THIS STEP IF YOU DID RIVA DEPLOY
! cd $RIVA_DIR && ./riva_init.sh config.sh

# Run Riva Start. This will deploy your model.
! cd $RIVA_DIR && ./riva_start.sh config.sh

运行推理#

在 Riva 服务器启动并运行您的模型后，您可以发送推理请求来查询服务器。

要发送 gRPC 请求，您可以安装客户端的 Riva Python API 绑定。这在 PyPI 上的 Python 模块中提供。

# Install the Client API Bindings
! pip install nvidia-riva-client

import riva.client

连接到 Riva 服务器并运行推理#

调用此推理函数会查询 Riva 服务器（使用 gRPC）以转录音频文件。

def run_inference(audio_file, server='localhost:50051', print_full_response=False):
    with open(audio_file, 'rb') as fh:
        data = fh.read()
    
    auth = riva.client.Auth(uri=server)
    client = riva.client.ASRService(auth)
    config = riva.client.RecognitionConfig(
        language_code="en-US",
        max_alternatives=1,
        enable_automatic_punctuation=False,
    )
    
    response = client.offline_recognize(data, config)
    if print_full_response: 
        print(response)
    else:
        print(response.results[0].alternatives[0].transcript)

现在我们可以实际查询 Riva 服务器了。

audio_file = "<add path to .wav (PCM-, A-Law-, or U-Law-encoded), .flac, .opus, or .ogg (Opus-encoded) file>"
run_inference(audio_file)

您可以在关闭 Jupyter 内核之前停止 Riva 服务器。

! cd $RIVA_DIR && ./riva_stop.sh 

NVIDIA Riva

如何在 Riva 上部署使用 NeMo 训练的自定义声学模型 (Citrinet)

目录