Add list_models option for ASR clients

Author: rmittal-github

riva/clients/asr/riva_asr_client.cc

@@ -52,6 +52,7 @@ DEFINE_int32(num_parallel_requests, 10, "Number of parallel requests to keep in
 DEFINE_bool(print_transcripts, true, "Print final transcripts");
 DEFINE_string(output_filename, "", "Filename to write output transcripts");
 DEFINE_string(model_name, "", "Name of the TRTIS model to use");
+DEFINE_bool(list_models, false, "List available models on server");
 DEFINE_bool(output_ctm, false, "If true, output format should be NIST CTM");
 DEFINE_string(language_code, "en-US", "Language code of the model to use");
 DEFINE_string(boosted_words_file, "", "File with a list of words to boost. One line per word.");
@@ -455,6 +456,8 @@ main(int argc, char** argv)
   str_usage << "           --boosted_words_score=<float>" << std::endl;
   str_usage << "           --ssl_cert=<filename>" << std::endl;
   str_usage << "           --speaker_diarization=<true|false>" << std::endl;
+  str_usage << "           --model_name=<model>" << std::endl;
+  str_usage << "           --list_models" << std::endl;
   str_usage << "           --metadata=<key,value,...>" << std::endl;
   str_usage << "           --start_history=<int>" << std::endl;
   str_usage << "           --start_threshold=<float>" << std::endl;
@@ -503,6 +506,30 @@ main(int argc, char** argv)
     return 1;
   }
 
+  if (FLAGS_list_models) {
+    std::unique_ptr<nr_asr::RivaSpeechRecognition::Stub> asr_stub_(
+        nr_asr::RivaSpeechRecognition::NewStub(grpc_channel));
+    grpc::ClientContext asr_context;
+    nr_asr::RivaSpeechRecognitionConfigRequest asr_request;
+    nr_asr::RivaSpeechRecognitionConfigResponse asr_response;
+    asr_stub_->GetRivaSpeechRecognitionConfig(&asr_context, asr_request, &asr_response);
+
+    std::multimap<std::string, std::string> model_map;
+    for (int i = 0; i < asr_response.model_config_size(); i++) {
+      if (asr_response.model_config(i).parameters().find("type")->second == "offline") {
+        model_map.insert(std::make_pair(
+            asr_response.model_config(i).parameters().find("language_code")->second,
+            asr_response.model_config(i).model_name()));
+      }
+    }
+
+    for (auto& m : model_map) {
+      std::cout << "'" << m.first << "': '" << m.second << "'" << std::endl;
+    }
+
+    return 0;
+  }
+
   RecognizeClient recognize_client(
       grpc_channel, FLAGS_language_code, FLAGS_max_alternatives, FLAGS_profanity_filter,
       FLAGS_word_time_offsets, FLAGS_automatic_punctuation,

riva/clients/asr/riva_streaming_asr_client.cc

@@ -63,6 +63,7 @@ DEFINE_string(
     output_filename, "final_transcripts.json",
     "Filename of .json file containing output transcripts");
 DEFINE_string(model_name, "", "Name of the TRTIS model to use");
+DEFINE_bool(list_models, false, "List available models on server");
 DEFINE_string(language_code, "en-US", "Language code of the model to use");
 DEFINE_string(boosted_words_file, "", "File with a list of words to boost. One line per word.");
 DEFINE_double(boosted_words_score, 10., "Score by which to boost the boosted words");
@@ -133,6 +134,8 @@ main(int argc, char** argv)
   str_usage << "           --boosted_words_file=<string>" << std::endl;
   str_usage << "           --boosted_words_score=<float>" << std::endl;
   str_usage << "           --ssl_cert=<filename>" << std::endl;
+  str_usage << "           --model_name=<model>" << std::endl;
+  str_usage << "           --list_models" << std::endl;
   str_usage << "           --metadata=<key,value,...>" << std::endl;
   str_usage << "           --start_history=<int>" << std::endl;
   str_usage << "           --start_threshold=<float>" << std::endl;
@@ -182,6 +185,29 @@ main(int argc, char** argv)
     return 1;
   }
 
+  if (FLAGS_list_models) {
+    std::unique_ptr<nr_asr::RivaSpeechRecognition::Stub> asr_stub_(
+        nr_asr::RivaSpeechRecognition::NewStub(grpc_channel));
+    grpc::ClientContext asr_context;
+    nr_asr::RivaSpeechRecognitionConfigRequest asr_request;
+    nr_asr::RivaSpeechRecognitionConfigResponse asr_response;
+    asr_stub_->GetRivaSpeechRecognitionConfig(&asr_context, asr_request, &asr_response);
+
+    std::multimap<std::string, std::string> model_map;
+    for (int i = 0; i < asr_response.model_config_size(); i++) {
+      if (asr_response.model_config(i).parameters().find("type")->second == "online") {
+        model_map.insert(std::make_pair(
+            asr_response.model_config(i).parameters().find("language_code")->second,
+            asr_response.model_config(i).model_name()));
+      }
+    }
+
+    for (auto& m : model_map) {
+      std::cout << "'" << m.first << "': '" << m.second << "'" << std::endl;
+    }
+    return 0;
+  }
+
   StreamingRecognizeClient recognize_client(
       grpc_channel, FLAGS_num_parallel_requests, FLAGS_language_code, FLAGS_max_alternatives,
       FLAGS_profanity_filter, FLAGS_word_time_offsets, FLAGS_automatic_punctuation,
@@ -230,4 +256,4 @@ main(int argc, char** argv)
   }
 
   return 0;
-}
+}

riva/clients/nmt/riva_nmt_streaming_s2s_client.cc

@@ -276,7 +276,7 @@ main(int argc, char** argv)
     nr_nmt::AvailableLanguageRequest request;
     nr_nmt::AvailableLanguageResponse response;
 
-    request.set_model("s2s_model"); // this is optional, if empty returns all available models/languages
+    request.set_model("s2s_model");  // get only S2S supported languages
     nmt_s2s->ListSupportedLanguagePairs(&context, request, &response);
     std::cout << response.DebugString() << std::endl;
     return 0;

riva/clients/nmt/riva_nmt_streaming_s2t_client.cc

@@ -157,7 +157,7 @@ main(int argc, char** argv)
     nr_nmt::AvailableLanguageRequest request;
     nr_nmt::AvailableLanguageResponse response;
 
-    request.set_model("s2t_model"); // this is optional, if empty returns all available models/languages
+    request.set_model("s2t_model");  // get only S2T supported languages
     nmt_s2t->ListSupportedLanguagePairs(&context, request, &response);
     std::cout << response.DebugString() << std::endl;
     return 0;