如何在AWS SageMaker多模型端点的Triton服务器返回字符串响应？

可以返回字符串，但不能直接返回原始字符串，必须按照Triton要求封装成pb_utils.InferenceResponse对象后返回。以下是完整的实现示例：

完整model.py示例（关键部分补全） #

import numpy as np
import sys
import os
import json
from pathlib import Path

import torch
from transformers import DistilBartForConditionalGeneration, DistilBartTokenizer

import triton_python_backend_utils as pb_utils

class TritonPythonModel:

    def initialize(self, args):
        # 从模型仓库加载模型和分词器
        model_dir = args['model_repository'] + '/' + args['model_version']
        self.tokenizer = DistilBartTokenizer.from_pretrained(model_dir)
        self.model = DistilBartForConditionalGeneration.from_pretrained(model_dir).cuda()
        self.model.eval()

    def execute(self, requests):
        responses = []
        for request in requests:
            # 获取输入张量（注意与config.pbtxt中的输入名称一致）
            input_ids = pb_utils.get_input_tensor_by_name(request, "INPUT").as_numpy()
            input_ids = torch.as_tensor(input_ids).long().cuda()

            # 生成摘要（关闭梯度计算提升性能）
            with torch.no_grad():
                translation = self.model.generate(input_ids=input_ids, num_beams=1)
            
            # 将模型输出的token ID解码为字符串
            summary_str = self.tokenizer.decode(translation[0], skip_special_tokens=True)
            
            # 将字符串转换为符合TYPE_STRING要求的numpy数组
            # 批量场景下需对应每个输入样本生成字符串，再组合成数组
            output_data = np.array([summary_str], dtype=np.object_)
            
            # 构建输出张量与响应对象
            output_tensor = pb_utils.Tensor("OUTPUT", output_data)
            inference_response = pb_utils.InferenceResponse(output_tensors=[output_tensor])
            responses.append(inference_response)
        
        return responses

关键注意事项 #

• 输入/输出张量名称必须与config.pbtxt中的定义完全匹配
• TYPE_STRING对应的numpy数据类型需为np.object_或np.str_，不能直接传入原始字符串，必须包装成numpy数组
• 批量请求场景下，需确保输出数组的维度与输入批量维度对应，每个输入样本生成一个对应字符串

补全后的config.pbtxt #

name: 'somename'
backend: 'python'
max_batch_size: 16
input [{
  name: "INPUT"
  data_type: TYPE_INT32
  dims: [ -1 ]
}]
output[{
   name: "OUTPUT"
   data_type: TYPE_STRING
   dims: [ 1 ]
}]

内容的提问来源于stack exchange，提问作者haju