Add an output to OpenAI models

triton-inference-server · Mar 1, 2024 · abad69f · abad69f
1 parent 5421506
commit abad69f
Show file tree

Hide file tree

Showing 2 changed files with 16 additions and 4 deletions.
diff --git a/src/c++/perf_analyzer/client_backend/client_backend.cc b/src/c++/perf_analyzer/client_backend/client_backend.cc
@@ -522,6 +522,12 @@ InferRequestedOutput::Create(
     RETURN_IF_CB_ERROR(tritonremote::TritonInferRequestedOutput::Create(
         infer_output, name, class_count));
   }
+#ifdef TRITON_ENABLE_PERF_ANALYZER_OPENAI
+  else if (kind == OPENAI) {
+    RETURN_IF_CB_ERROR(
+        openai::OpenAiInferRequestedOutput::Create(infer_output, name));
+  }
+#endif  // TRITON_ENABLE_PERF_ANALYZER_OPENAI
 #ifdef TRITON_ENABLE_PERF_ANALYZER_TFS
   else if (kind == TENSORFLOW_SERVING) {
     RETURN_IF_CB_ERROR(

diff --git a/src/c++/perf_analyzer/model_parser.cc b/src/c++/perf_analyzer/model_parser.cc
@@ -277,10 +277,16 @@ ModelParser::InitOpenAI(
   max_batch_size_ = batch_size;
 
   // OpenAI will take a single json input with a fully formed payload
-  auto it = inputs_->emplace("payload", ModelTensor()).first;
-  it->second.name_ = "payload";
-  it->second.datatype_ = "JSON";
-  it->second.shape_.push_back(1);
+  auto in_it = inputs_->emplace("payload", ModelTensor()).first;
+  in_it->second.name_ = "payload";
+  in_it->second.datatype_ = "JSON";
+  in_it->second.shape_.push_back(1);
+
+  // OpenAI will reply with a single json output
+  auto out_it = outputs_->emplace("response", ModelTensor()).first;
+  out_it->second.name_ = "response";
+  out_it->second.datatype_ = "JSON";
+  out_it->second.shape_.push_back(1);
 
   return cb::Error::Success;
 }