diff --git a/README.md b/README.md index 6092fe8..eb05551 100644 --- a/README.md +++ b/README.md @@ -85,6 +85,19 @@ $ riva_asr_client --audio_file audio_folder Note that only single-channel audio files in the `.wav` format are currently supported. +Final transcripts can be exported as structured JSON, SubRip, or WebVTT by choosing the +corresponding extension for `--output_filename`: +``` +$ riva_asr_client --audio_file audio.wav --output_filename transcript.json +$ riva_asr_client --audio_file audio.wav --output_filename transcript.srt +$ riva_streaming_asr_client --audio_file audio.wav --output_filename transcript.vtt +``` +The format can also be selected explicitly with `--output_format=json|srt|vtt`. JSON output uses +one object per recognition request (JSON Lines when processing multiple requests). SRT and VTT +automatically request word timestamps and currently require one audio file with +`--num_iterations=1`; microphone input is not supported for subtitle export. Each final recognition +result becomes one subtitle cue. + Other options and information can be found by running the clients with `-help` ### Speech Synthesis (TTS) Client diff --git a/riva/clients/asr/BUILD b/riva/clients/asr/BUILD index d88c46b..87505fd 100644 --- a/riva/clients/asr/BUILD +++ b/riva/clients/asr/BUILD @@ -27,6 +27,16 @@ cc_library( ], ) +cc_library( + name = "transcript_exporter", + srcs = ["transcript_exporter.cc"], + hdrs = ["transcript_exporter.h"], + deps = [ + ":asr_client_helper", + "@nvriva_common//riva/proto:riva_grpc_asr", + ], +) + cc_library( name = "client_call", srcs = ["client_call.h", "client_call.cc"], @@ -39,6 +49,7 @@ cc_library( ], }) + [ ":asr_client_helper", + ":transcript_exporter", "@com_github_grpc_grpc//:grpc++", "@nvriva_common//riva/proto:riva_grpc_asr", "@glog//:glog", @@ -54,6 +65,7 @@ cc_library( deps = [ ":asr_client_helper", ":client_call", + ":transcript_exporter", "//riva/utils/wav:reader", "//riva/utils/opus", "@glog//:glog", @@ -79,6 +91,7 @@ cc_binary( deps = [ ":asr_client_helper", ":client_call", + ":transcript_exporter", "@nvriva_common//riva/proto:riva_grpc_asr", "//riva/utils:stamping", "//riva/utils/files:files", @@ -97,6 +110,7 @@ cc_binary( ":asr_client_helper", ":client_call", ":streaming_recognize_client", + ":transcript_exporter", "@nvriva_common//riva/proto:riva_grpc_asr", "//riva/utils/files:files", "//riva/utils/wav:reader", @@ -126,3 +140,13 @@ cc_test( ], tags = ["needs_alsa"] ) + +cc_test( + name = "transcript_exporter_test", + srcs = ["transcript_exporter_test.cc"], + deps = [ + ":transcript_exporter", + "@googletest//:gtest_main", + "@nvriva_common//riva/proto:riva_grpc_asr", + ], +) diff --git a/riva/clients/asr/client_call.cc b/riva/clients/asr/client_call.cc index 536c411..4f75c01 100644 --- a/riva/clients/asr/client_call.cc +++ b/riva/clients/asr/client_call.cc @@ -43,6 +43,13 @@ ClientCall::AppendResult(const nr_asr::StreamingRecognitionResult& result) bool is_final = result.is_final(); if (is_final) { int num_alternatives = result.alternatives_size(); + if (num_alternatives > 0) { + Results::Segment segment; + for (int a = 0; a < num_alternatives; ++a) { + segment.alternatives.push_back(result.alternatives(a)); + } + latest_result_.segments.push_back(std::move(segment)); + } latest_result_.final_transcripts.resize(num_alternatives); latest_result_.final_scores.resize(num_alternatives); latest_result_.final_time_stamps.resize(num_alternatives); @@ -88,7 +95,7 @@ ClientCall::AppendResult(const nr_asr::StreamingRecognitionResult& result) } void -ClientCall::PrintResult(bool audio_device, std::ofstream& output_file) +ClientCall::PrintResult(bool audio_device) { std::cout << "-----------------------------------------------------------" << std::endl; @@ -100,16 +107,8 @@ ClientCall::PrintResult(bool audio_device, std::ofstream& output_file) std::cout << std::endl; std::cout << "Final transcripts: " << std::endl; - if (latest_result_.final_transcripts.size() == 0) { - output_file << "{\"audio_filepath\": \"" << filename << "\","; - output_file << "\"text\": \"\"}" << std::endl; - } else { + if (latest_result_.final_transcripts.size() > 0) { for (uint32_t a = 0; a < latest_result_.final_transcripts.size(); ++a) { - if (a == 0) { - output_file << "{\"audio_filepath\": \"" << filename << "\","; - output_file << "\"text\": \"" << EscapeTranscript(latest_result_.final_transcripts[a]) - << "\"}" << std::endl; - } std::cout << a << " : " << latest_result_.final_transcripts[a] << latest_result_.partial_transcript << std::endl; std::cout << std::endl; diff --git a/riva/clients/asr/client_call.h b/riva/clients/asr/client_call.h index f639d07..d2bd759 100644 --- a/riva/clients/asr/client_call.h +++ b/riva/clients/asr/client_call.h @@ -41,7 +41,7 @@ class ClientCall { void AppendResult(const nr_asr::StreamingRecognitionResult& result); - void PrintResult(bool audio_device, std::ofstream& output_file); + void PrintResult(bool audio_device); // Container for the data we expect from the server. nr_asr::StreamingRecognizeResponse response; diff --git a/riva/clients/asr/riva_asr_client.cc b/riva/clients/asr/riva_asr_client.cc index 6e2cfaf..33cdf03 100644 --- a/riva/clients/asr/riva_asr_client.cc +++ b/riva/clients/asr/riva_asr_client.cc @@ -28,6 +28,7 @@ #include "riva/utils/stamping.h" #include "riva/utils/wav/wav_reader.h" #include "riva_asr_client_helper.h" +#include "transcript_exporter.h" using grpc::Status; using grpc::StatusCode; @@ -51,6 +52,9 @@ DEFINE_int32(num_iterations, 1, "Number of times to loop over audio files"); DEFINE_int32(num_parallel_requests, 10, "Number of parallel requests to keep in flight"); DEFINE_bool(print_transcripts, true, "Print final transcripts"); DEFINE_string(output_filename, "", "Filename to write output transcripts"); +DEFINE_string( + output_format, "", + "Transcript output format: json, srt, or vtt. By default, infer from output_filename"); 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"); @@ -100,7 +104,8 @@ class RecognizeClient { const std::string& boosted_phrases_file, float boosted_phrases_score, bool speaker_diarization, int32_t diarization_max_speakers, int32_t start_history, float start_threshold, int32_t stop_history, int32_t stop_history_eou, float stop_threshold, - float stop_threshold_eou, std::string custom_configuration) + float stop_threshold_eou, std::string custom_configuration, + TranscriptOutputFormat output_format) : stub_(nr_asr::RivaSpeechRecognition::NewStub(channel)), language_code_(language_code), max_alternatives_(max_alternatives), profanity_filter_(profanity_filter), word_time_offsets_(word_time_offsets), automatic_punctuation_(automatic_punctuation), @@ -113,14 +118,15 @@ class RecognizeClient { start_history_(start_history), start_threshold_(start_threshold), stop_history_(stop_history), stop_history_eou_(stop_history_eou), stop_threshold_(stop_threshold), stop_threshold_eou_(stop_threshold_eou), - custom_configuration_(custom_configuration) + custom_configuration_(custom_configuration), output_format_(output_format), cue_index_(1) { if (!output_filename.empty()) { output_file_.open(output_filename); if (ctm) { write_fn_ = &RecognizeClient::WriteCTM; } else { - write_fn_ = &RecognizeClient::WriteJSON; + write_fn_ = &RecognizeClient::WriteTranscript; + InitializeTranscriptOutput(output_file_, output_format_); } } @@ -164,19 +170,13 @@ class RecognizeClient { } } - void WriteJSON(const Results& result, const std::string& filename) + void WriteTranscript(const Results& result, const std::string& filename) { - if (result.final_transcripts.size() == 0) { - output_file_ << "{\"audio_filepath\": \"" << filename << "\","; - output_file_ << "\"text\": \"\"}" << std::endl; - } else { - for (size_t a = 0; a < result.final_transcripts.size(); ++a) { - if (a == 0) { - output_file_ << "{\"audio_filepath\": \"" << filename << "\","; - output_file_ << "\"text\": \"" << EscapeTranscript(result.final_transcripts.at(a)) - << "\"}" << std::endl; - } - } + std::string error; + if (!WriteTranscriptOutput( + output_file_, output_format_, result, filename, &cue_index_, &error)) { + std::cerr << "Unable to export transcript: " << error << std::endl; + num_failed_requests_++; } } @@ -439,6 +439,8 @@ class RecognizeClient { float stop_threshold_; float stop_threshold_eou_; std::string custom_configuration_; + TranscriptOutputFormat output_format_; + size_t cue_index_; }; int @@ -459,6 +461,7 @@ main(int argc, char** argv) str_usage << " --num_parallel_requests= " << std::endl; str_usage << " --print_transcripts= " << std::endl; str_usage << " --output_filename=" << std::endl; + str_usage << " --output_format=" << std::endl; str_usage << " --output-ctm=" << std::endl; str_usage << " --verbatim_transcripts=" << std::endl; str_usage << " --language_code=" << std::endl; @@ -501,6 +504,23 @@ main(int argc, char** argv) return 1; } + if (FLAGS_output_ctm && !FLAGS_output_format.empty()) { + std::cerr << "output_ctm and output_format cannot be used together." << std::endl; + return 1; + } + TranscriptOutputFormat output_format = TranscriptOutputFormat::kJson; + if (!FLAGS_output_filename.empty() && !FLAGS_output_ctm) { + std::string error; + if (!ParseTranscriptOutputFormat( + FLAGS_output_filename, FLAGS_output_format, &output_format, &error)) { + std::cerr << error << std::endl; + return 1; + } + if (TranscriptOutputRequiresWordTimeOffsets(output_format)) { + FLAGS_word_time_offsets = true; + } + } + bool flag_set = gflags::GetCommandLineFlagInfoOrDie("riva_uri").is_default; const char* riva_uri = getenv("RIVA_URI"); @@ -541,15 +561,6 @@ main(int argc, char** argv) return 0; } - RecognizeClient recognize_client( - grpc_channel, FLAGS_language_code, FLAGS_max_alternatives, FLAGS_profanity_filter, - FLAGS_word_time_offsets, FLAGS_automatic_punctuation, - /* separate_recognition_per_channel*/ false, FLAGS_print_transcripts, FLAGS_output_filename, - FLAGS_model_name, FLAGS_output_ctm, FLAGS_verbatim_transcripts, FLAGS_boosted_words_file, - (float)FLAGS_boosted_words_score, FLAGS_speaker_diarization, FLAGS_diarization_max_speakers, - FLAGS_start_history, FLAGS_start_threshold, FLAGS_stop_history, FLAGS_stop_history_eou, - FLAGS_stop_threshold, FLAGS_stop_threshold_eou, FLAGS_custom_configuration); - // Preload all wav files, sort by size to reduce tail effects std::vector> all_wav; try { @@ -563,6 +574,20 @@ main(int argc, char** argv) std::cout << "No audio files specified. Exiting." << std::endl; return 1; } + if (TranscriptOutputRequiresWordTimeOffsets(output_format) && + (all_wav.size() != 1 || FLAGS_num_iterations != 1)) { + std::cerr << "SRT and VTT export require one audio file and num_iterations=1." << std::endl; + return 1; + } + + RecognizeClient recognize_client( + grpc_channel, FLAGS_language_code, FLAGS_max_alternatives, FLAGS_profanity_filter, + FLAGS_word_time_offsets, FLAGS_automatic_punctuation, + /* separate_recognition_per_channel*/ false, FLAGS_print_transcripts, FLAGS_output_filename, + FLAGS_model_name, FLAGS_output_ctm, FLAGS_verbatim_transcripts, FLAGS_boosted_words_file, + (float)FLAGS_boosted_words_score, FLAGS_speaker_diarization, FLAGS_diarization_max_speakers, + FLAGS_start_history, FLAGS_start_threshold, FLAGS_stop_history, FLAGS_stop_history_eou, + FLAGS_stop_threshold, FLAGS_stop_threshold_eou, FLAGS_custom_configuration, output_format); uint32_t all_wav_max = all_wav.size() * FLAGS_num_iterations; std::vector> all_wav_repeated; @@ -614,5 +639,5 @@ main(int argc, char** argv) } } - return 0; + return recognize_client.NumFailedRequests() ? 1 : 0; } diff --git a/riva/clients/asr/riva_asr_client_helper.cc b/riva/clients/asr/riva_asr_client_helper.cc index 068a37f..179a8b7 100644 --- a/riva/clients/asr/riva_asr_client_helper.cc +++ b/riva/clients/asr/riva_asr_client_helper.cc @@ -125,6 +125,13 @@ AppendResult( } int num_alternatives = result.alternatives_size(); + if (num_alternatives > 0) { + Results::Segment segment; + for (int a = 0; a < num_alternatives; ++a) { + segment.alternatives.push_back(result.alternatives(a)); + } + output_result.segments.push_back(std::move(segment)); + } output_result.final_transcripts.resize(num_alternatives); output_result.final_scores.resize(num_alternatives); output_result.final_time_stamps.resize(num_alternatives); @@ -244,4 +251,4 @@ ReadCustomConfiguration(std::string& custom_configuration) } } return custom_configuration_map; -} \ No newline at end of file +} diff --git a/riva/clients/asr/riva_asr_client_helper.h b/riva/clients/asr/riva_asr_client_helper.h index 9abfc33..93039a8 100644 --- a/riva/clients/asr/riva_asr_client_helper.h +++ b/riva/clients/asr/riva_asr_client_helper.h @@ -37,6 +37,10 @@ std::string static inline EscapeTranscript(const std::string& input_str) } struct Results { + struct Segment { + std::vector alternatives; + }; + std::vector final_transcripts; std::vector final_scores; std::string partial_transcript; @@ -45,6 +49,7 @@ struct Results { std::vector partial_time_stamps; int request_cnt; float audio_processed; + std::vector segments; }; void AppendResult( diff --git a/riva/clients/asr/riva_streaming_asr_client.cc b/riva/clients/asr/riva_streaming_asr_client.cc index 465f09c..b72d6b6 100644 --- a/riva/clients/asr/riva_streaming_asr_client.cc +++ b/riva/clients/asr/riva_streaming_asr_client.cc @@ -61,7 +61,10 @@ DEFINE_bool(print_transcripts, true, "Print final transcripts"); DEFINE_bool(interim_results, true, "Print intermediate transcripts"); DEFINE_string( output_filename, "final_transcripts.json", - "Filename of .json file containing output transcripts"); + "Filename containing output transcripts"); +DEFINE_string( + output_format, "", + "Transcript output format: json, srt, or vtt. By default, infer from output_filename"); 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, "", "Language code of the model to use"); @@ -137,6 +140,7 @@ main(int argc, char** argv) str_usage << " --num_parallel_requests= " << std::endl; str_usage << " --print_transcripts= " << std::endl; str_usage << " --output_filename=" << std::endl; + str_usage << " --output_format=" << std::endl; str_usage << " --verbatim_transcripts=" << std::endl; str_usage << " --language_code=" << std::endl; str_usage << " --boosted_words_file=" << std::endl; @@ -179,6 +183,21 @@ main(int argc, char** argv) return 1; } + TranscriptOutputFormat output_format; + std::string output_error; + if (!ParseTranscriptOutputFormat( + FLAGS_output_filename, FLAGS_output_format, &output_format, &output_error)) { + std::cerr << output_error << std::endl; + return 1; + } + if (TranscriptOutputRequiresWordTimeOffsets(output_format)) { + FLAGS_word_time_offsets = true; + if (!FLAGS_audio_device.empty()) { + std::cerr << "SRT and VTT export currently support audio_file input only." << std::endl; + return 1; + } + } + bool flag_set = gflags::GetCommandLineFlagInfoOrDie("riva_uri").is_default; const char* riva_uri = getenv("RIVA_URI"); @@ -227,7 +246,7 @@ main(int argc, char** argv) FLAGS_verbatim_transcripts, FLAGS_boosted_words_file, FLAGS_boosted_words_score, FLAGS_start_history, FLAGS_start_threshold, FLAGS_stop_history, FLAGS_stop_history_eou, FLAGS_stop_threshold, FLAGS_stop_threshold_eou, FLAGS_custom_configuration, - FLAGS_speaker_diarization, FLAGS_diarization_max_speakers); + FLAGS_speaker_diarization, FLAGS_diarization_max_speakers, output_format); if (FLAGS_audio_file.size()) { return recognize_client.DoStreamingFromFile( diff --git a/riva/clients/asr/streaming_recognize_client.cc b/riva/clients/asr/streaming_recognize_client.cc index 84581af..6b166d4 100644 --- a/riva/clients/asr/streaming_recognize_client.cc +++ b/riva/clients/asr/streaming_recognize_client.cc @@ -60,7 +60,8 @@ StreamingRecognizeClient::StreamingRecognizeClient( bool verbatim_transcripts, const std::string& boosted_phrases_file, float boosted_phrases_score, int32_t start_history, float start_threshold, int32_t stop_history, int32_t stop_history_eou, float stop_threshold, float stop_threshold_eou, std::string custom_configuration, - bool speaker_diarization, int32_t diarization_max_speakers) + bool speaker_diarization, int32_t diarization_max_speakers, + TranscriptOutputFormat output_format) : print_latency_stats_(true), stub_(nr_asr::RivaSpeechRecognition::NewStub(channel)), language_code_(language_code), max_alternatives_(max_alternatives), profanity_filter_(profanity_filter), word_time_offsets_(word_time_offsets), @@ -68,19 +69,27 @@ StreamingRecognizeClient::StreamingRecognizeClient( separate_recognition_per_channel_(separate_recognition_per_channel), print_transcripts_(print_transcripts), chunk_duration_ms_(chunk_duration_ms), interim_results_(interim_results), total_audio_processed_(0.), num_streams_started_(0), - model_name_(model_name), simulate_realtime_(simulate_realtime), + write_transcripts_(!output_filename.empty()), model_name_(model_name), + simulate_realtime_(simulate_realtime), verbatim_transcripts_(verbatim_transcripts), boosted_phrases_score_(boosted_phrases_score), start_history_(start_history), start_threshold_(start_threshold), stop_history_(stop_history), stop_history_eou_(stop_history_eou), stop_threshold_(stop_threshold), stop_threshold_eou_(stop_threshold_eou), custom_configuration_(custom_configuration), - speaker_diarization_(speaker_diarization), diarization_max_speakers_(diarization_max_speakers) + speaker_diarization_(speaker_diarization), diarization_max_speakers_(diarization_max_speakers), + output_format_(output_format), cue_index_(1), export_failed_(false) { num_active_streams_.store(0); num_streams_finished_.store(0); thread_pool_.reset(new ThreadPool(4 * num_parallel_requests)); - if (print_transcripts_) { + if (write_transcripts_) { output_file_.open(output_filename); + if (!output_file_.is_open()) { + std::cerr << "Unable to open transcript output file: " << output_filename << std::endl; + export_failed_.store(true); + } else { + InitializeTranscriptOutput(output_file_, output_format_); + } } boosted_phrases_ = ReadPhrasesFromFile(boosted_phrases_file); @@ -88,7 +97,7 @@ StreamingRecognizeClient::StreamingRecognizeClient( StreamingRecognizeClient::~StreamingRecognizeClient() { - if (print_transcripts_) { + if (write_transcripts_) { output_file_.close(); } } @@ -264,6 +273,11 @@ StreamingRecognizeClient::DoStreamingFromFile( std::cout << "No audio files specified. Exiting." << std::endl; return 1; } + if (TranscriptOutputRequiresWordTimeOffsets(output_format_) && + (all_wav.size() != 1 || num_iterations != 1)) { + std::cerr << "SRT and VTT export require one audio file and num_iterations=1." << std::endl; + return 1; + } uint32_t all_wav_max = all_wav.size() * num_iterations; std::vector> all_wav_repeated; @@ -315,7 +329,7 @@ StreamingRecognizeClient::DoStreamingFromFile( std::cout << "Throughput: " << total_processed * 1000. / diff_time << " RTFX" << std::endl; } - return 0; + return export_failed_.load() ? 1 : 0; } void @@ -340,8 +354,17 @@ StreamingRecognizeClient::PostProcessResults(std::shared_ptr call, b latencies_.push_back(lat); } } + if (write_transcripts_) { + const std::string filename = audio_device ? "microphone" : call->stream->wav->filename; + std::string error; + if (!WriteTranscriptOutput( + output_file_, output_format_, call->latest_result_, filename, &cue_index_, &error)) { + std::cerr << "Unable to export transcript: " << error << std::endl; + export_failed_.store(true); + } + } if (print_transcripts_) { - call->PrintResult(audio_device, output_file_); + call->PrintResult(audio_device); } } @@ -376,7 +399,7 @@ StreamingRecognizeClient::ReceiveResponses(std::shared_ptr call, boo call->latest_result_.audio_processed = result.audio_processed(); - if (print_transcripts_) { + if (print_transcripts_ || write_transcripts_) { call->AppendResult(result); } } @@ -454,7 +477,7 @@ StreamingRecognizeClient::DoStreamingFromMicrophone( CloseAudioDevice(&alsa_handle); std::cout << "\nExiting with 0" << std::flush << std::endl; - return 0; + return export_failed_.load() ? 1 : 0; } void @@ -500,4 +523,4 @@ StreamingRecognizeClient::PrintStats() << std::endl; return 1; } -} \ No newline at end of file +} diff --git a/riva/clients/asr/streaming_recognize_client.h b/riva/clients/asr/streaming_recognize_client.h index 9232cff..b2e038d 100644 --- a/riva/clients/asr/streaming_recognize_client.h +++ b/riva/clients/asr/streaming_recognize_client.h @@ -31,6 +31,7 @@ #include "riva/utils/thread_pool.h" #include "riva/utils/wav/wav_reader.h" #include "riva_asr_client_helper.h" +#include "transcript_exporter.h" using grpc::Status; using grpc::StatusCode; @@ -50,7 +51,8 @@ class StreamingRecognizeClient { float boosted_phrases_score, int32_t start_history, float start_threshold, int32_t stop_history, int32_t stop_history_eou, float stop_threshold, float stop_threshold_eou, std::string custom_configuration, - bool speaker_diarization, int32_t diarization_max_speakers); + bool speaker_diarization, int32_t diarization_max_speakers, + TranscriptOutputFormat output_format); ~StreamingRecognizeClient(); @@ -114,6 +116,7 @@ class StreamingRecognizeClient { std::unique_ptr thread_pool_; std::ofstream output_file_; + bool write_transcripts_; std::string model_name_; bool simulate_realtime_; @@ -131,4 +134,7 @@ class StreamingRecognizeClient { std::string custom_configuration_; bool speaker_diarization_; int32_t diarization_max_speakers_; -}; \ No newline at end of file + TranscriptOutputFormat output_format_; + size_t cue_index_; + std::atomic export_failed_; +}; diff --git a/riva/clients/asr/transcript_exporter.cc b/riva/clients/asr/transcript_exporter.cc new file mode 100644 index 0000000..ea8b821 --- /dev/null +++ b/riva/clients/asr/transcript_exporter.cc @@ -0,0 +1,302 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: MIT + */ + +#include "transcript_exporter.h" + +#include +#include +#include +#include + +namespace { + +std::string +Lowercase(std::string value) +{ + std::transform(value.begin(), value.end(), value.begin(), [](unsigned char character) { + return static_cast(std::tolower(character)); + }); + return value; +} + +std::string +Extension(const std::string& filename) +{ + const size_t dot = filename.find_last_of('.'); + if (dot == std::string::npos || dot + 1 == filename.size()) { + return ""; + } + return Lowercase(filename.substr(dot + 1)); +} + +void +WriteJsonString(std::ostream& output, const std::string& value) +{ + output << '"'; + for (unsigned char character : value) { + switch (character) { + case '"': output << "\\\""; break; + case '\\': output << "\\\\"; break; + case '\b': output << "\\b"; break; + case '\f': output << "\\f"; break; + case '\n': output << "\\n"; break; + case '\r': output << "\\r"; break; + case '\t': output << "\\t"; break; + default: + if (character < 0x20) { + output << "\\u00" << std::hex << std::setw(2) << std::setfill('0') + << static_cast(character) << std::dec << std::setfill(' '); + } else { + output << character; + } + } + } + output << '"'; +} + +void +WriteWord(std::ostream& output, const nr_asr::WordInfo& word) +{ + output << "{\"word\":"; + WriteJsonString(output, word.word()); + output << ",\"start_time_ms\":" << word.start_time() + << ",\"end_time_ms\":" << word.end_time() << ",\"confidence\":" + << word.confidence() << ",\"speaker_tag\":" << word.speaker_tag() + << ",\"language_code\":"; + WriteJsonString(output, word.language_code()); + output << '}'; +} + +void +WriteAlternative(std::ostream& output, const nr_asr::SpeechRecognitionAlternative& alternative) +{ + output << "{\"transcript\":"; + WriteJsonString(output, alternative.transcript()); + output << ",\"confidence\":" << alternative.confidence() << ",\"words\":["; + for (int word_index = 0; word_index < alternative.words_size(); ++word_index) { + if (word_index > 0) { + output << ','; + } + WriteWord(output, alternative.words(word_index)); + } + output << "],\"language_codes\":["; + for (int language_index = 0; language_index < alternative.language_code_size(); + ++language_index) { + if (language_index > 0) { + output << ','; + } + WriteJsonString(output, alternative.language_code(language_index)); + } + output << "]}"; +} + +std::string +TranscriptText(const Results& result) +{ + std::string transcript; + for (const auto& segment : result.segments) { + if (segment.alternatives.empty() || segment.alternatives.front().transcript().empty()) { + continue; + } + if (!transcript.empty()) { + transcript += ' '; + } + transcript += segment.alternatives.front().transcript(); + } + if (transcript.empty() && !result.final_transcripts.empty()) { + transcript = result.final_transcripts.front(); + } + return transcript; +} + +void +WriteJson(std::ostream& output, const Results& result, const std::string& audio_filename) +{ + output << "{\"schema_version\":\"1.0\",\"audio_filepath\":"; + WriteJsonString(output, audio_filename); + const std::string transcript = TranscriptText(result); + output << ",\"text\":"; + WriteJsonString(output, transcript); + output << ",\"transcript\":"; + WriteJsonString(output, transcript); + output << ",\"segments\":["; + bool first_segment = true; + for (const auto& segment : result.segments) { + if (segment.alternatives.empty()) { + continue; + } + if (!first_segment) { + output << ','; + } + first_segment = false; + const auto& top_alternative = segment.alternatives.front(); + output << "{\"transcript\":"; + WriteJsonString(output, top_alternative.transcript()); + if (top_alternative.words_size() > 0) { + output << ",\"start_time_ms\":" << top_alternative.words(0).start_time() + << ",\"end_time_ms\":" + << top_alternative.words(top_alternative.words_size() - 1).end_time(); + } else { + output << ",\"start_time_ms\":null,\"end_time_ms\":null"; + } + output << ",\"confidence\":" << top_alternative.confidence() << ",\"words\":["; + for (int word_index = 0; word_index < top_alternative.words_size(); ++word_index) { + if (word_index > 0) { + output << ','; + } + WriteWord(output, top_alternative.words(word_index)); + } + output << "],\"alternatives\":["; + for (size_t alternative_index = 0; alternative_index < segment.alternatives.size(); + ++alternative_index) { + if (alternative_index > 0) { + output << ','; + } + WriteAlternative(output, segment.alternatives[alternative_index]); + } + output << "]}"; + } + output << "]}" << std::endl; +} + +std::string +CollapseWhitespace(const std::string& text) +{ + std::string collapsed; + bool pending_space = false; + for (unsigned char character : text) { + if (std::isspace(character)) { + pending_space = !collapsed.empty(); + } else { + if (pending_space) { + collapsed += ' '; + } + collapsed += static_cast(character); + pending_space = false; + } + } + return collapsed; +} + +std::string +FormatTimestamp(int64_t timestamp_ms, bool webvtt) +{ + const int64_t hours = timestamp_ms / 3600000; + timestamp_ms %= 3600000; + const int64_t minutes = timestamp_ms / 60000; + timestamp_ms %= 60000; + const int64_t seconds = timestamp_ms / 1000; + const int64_t milliseconds = timestamp_ms % 1000; + std::ostringstream formatted; + formatted << std::setfill('0') << std::setw(2) << hours << ':' << std::setw(2) << minutes << ':' + << std::setw(2) << seconds << (webvtt ? '.' : ',') << std::setw(3) << milliseconds; + return formatted.str(); +} + +bool +ValidateSubtitles(const Results& result, std::string* error) +{ + for (const auto& segment : result.segments) { + if (segment.alternatives.empty() || segment.alternatives.front().transcript().empty()) { + continue; + } + const auto& alternative = segment.alternatives.front(); + if (alternative.words_size() == 0) { + *error = "SRT and VTT export require word time offsets, but a finalized transcript segment " + "did not contain them."; + return false; + } + const int64_t start_time = alternative.words(0).start_time(); + const int64_t end_time = alternative.words(alternative.words_size() - 1).end_time(); + if (start_time < 0 || end_time < start_time) { + *error = "SRT and VTT export require non-negative, ordered timestamps."; + return false; + } + } + return true; +} + +void +WriteSubtitles(std::ostream& output, const Results& result, bool webvtt, size_t* cue_index) +{ + for (const auto& segment : result.segments) { + if (segment.alternatives.empty()) { + continue; + } + const auto& alternative = segment.alternatives.front(); + const std::string text = CollapseWhitespace(alternative.transcript()); + if (text.empty()) { + continue; + } + const int64_t start_time = alternative.words(0).start_time(); + const int64_t end_time = + std::max(alternative.words(alternative.words_size() - 1).end_time(), start_time + 1); + output << (*cue_index)++ << '\n' + << FormatTimestamp(start_time, webvtt) << " --> " << FormatTimestamp(end_time, webvtt) + << '\n' + << text << "\n\n"; + } +} + +} // namespace + +bool +ParseTranscriptOutputFormat( + const std::string& output_filename, const std::string& requested_format, + TranscriptOutputFormat* output_format, std::string* error) +{ + const std::string format = + requested_format.empty() ? Extension(output_filename) : Lowercase(requested_format); + if (format == "json") { + *output_format = TranscriptOutputFormat::kJson; + } else if (format == "srt") { + *output_format = TranscriptOutputFormat::kSrt; + } else if (format == "vtt") { + *output_format = TranscriptOutputFormat::kVtt; + } else { + *error = "Unable to determine transcript output format. Use a .json, .srt, or .vtt file " + "extension, or pass --output_format."; + return false; + } + return true; +} + +bool +TranscriptOutputRequiresWordTimeOffsets(TranscriptOutputFormat output_format) +{ + return output_format == TranscriptOutputFormat::kSrt || output_format == TranscriptOutputFormat::kVtt; +} + +void +InitializeTranscriptOutput(std::ostream& output, TranscriptOutputFormat output_format) +{ + if (output_format == TranscriptOutputFormat::kVtt) { + output << "WEBVTT\n\n"; + } +} + +bool +WriteTranscriptOutput( + std::ostream& output, TranscriptOutputFormat output_format, const Results& result, + const std::string& audio_filename, size_t* cue_index, std::string* error) +{ + if (!output.good()) { + *error = "Transcript output file is not writable."; + return false; + } + if (output_format == TranscriptOutputFormat::kJson) { + WriteJson(output, result, audio_filename); + } else { + if (!ValidateSubtitles(result, error)) { + return false; + } + WriteSubtitles(output, result, output_format == TranscriptOutputFormat::kVtt, cue_index); + } + if (!output.good()) { + *error = "Failed while writing transcript output."; + return false; + } + return true; +} diff --git a/riva/clients/asr/transcript_exporter.h b/riva/clients/asr/transcript_exporter.h new file mode 100644 index 0000000..99a658b --- /dev/null +++ b/riva/clients/asr/transcript_exporter.h @@ -0,0 +1,26 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: MIT + */ + +#pragma once + +#include +#include +#include + +#include "riva_asr_client_helper.h" + +enum class TranscriptOutputFormat { kJson, kSrt, kVtt }; + +bool ParseTranscriptOutputFormat( + const std::string& output_filename, const std::string& requested_format, + TranscriptOutputFormat* output_format, std::string* error); + +bool TranscriptOutputRequiresWordTimeOffsets(TranscriptOutputFormat output_format); + +void InitializeTranscriptOutput(std::ostream& output, TranscriptOutputFormat output_format); + +bool WriteTranscriptOutput( + std::ostream& output, TranscriptOutputFormat output_format, const Results& result, + const std::string& audio_filename, size_t* cue_index, std::string* error); diff --git a/riva/clients/asr/transcript_exporter_test.cc b/riva/clients/asr/transcript_exporter_test.cc new file mode 100644 index 0000000..2d4ab96 --- /dev/null +++ b/riva/clients/asr/transcript_exporter_test.cc @@ -0,0 +1,121 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: MIT + */ + +#include "transcript_exporter.h" + +#include +#include + +#include "gtest/gtest.h" + +namespace { + +Results +MakeResults() +{ + Results result; + Results::Segment segment; + nr_asr::SpeechRecognitionAlternative alternative; + alternative.set_transcript("hello \"world\""); + alternative.set_confidence(0.9F); + auto* first_word = alternative.add_words(); + first_word->set_word("hello"); + first_word->set_start_time(1000); + first_word->set_end_time(1200); + first_word->set_confidence(0.8F); + auto* second_word = alternative.add_words(); + second_word->set_word("world"); + second_word->set_start_time(1300); + second_word->set_end_time(1500); + second_word->set_confidence(0.7F); + segment.alternatives.push_back(alternative); + result.segments.push_back(segment); + result.final_transcripts.push_back(alternative.transcript()); + return result; +} + +TEST(TranscriptExporterTest, InfersOutputFormat) +{ + TranscriptOutputFormat format; + std::string error; + EXPECT_TRUE(ParseTranscriptOutputFormat("transcript.SRT", "", &format, &error)); + EXPECT_EQ(format, TranscriptOutputFormat::kSrt); + EXPECT_TRUE(ParseTranscriptOutputFormat("transcript.data", "vtt", &format, &error)); + EXPECT_EQ(format, TranscriptOutputFormat::kVtt); + EXPECT_FALSE(ParseTranscriptOutputFormat("transcript.txt", "", &format, &error)); +} + +TEST(TranscriptExporterTest, WritesSrt) +{ + std::ostringstream output; + size_t cue_index = 1; + std::string error; + + EXPECT_TRUE(WriteTranscriptOutput( + output, TranscriptOutputFormat::kSrt, MakeResults(), "audio.wav", &cue_index, &error)); + EXPECT_EQ( + output.str(), + "1\n00:00:01,000 --> 00:00:01,500\nhello \"world\"\n\n"); +} + +TEST(TranscriptExporterTest, WritesWebVtt) +{ + std::ostringstream output; + InitializeTranscriptOutput(output, TranscriptOutputFormat::kVtt); + size_t cue_index = 1; + std::string error; + + EXPECT_TRUE(WriteTranscriptOutput( + output, TranscriptOutputFormat::kVtt, MakeResults(), "audio.wav", &cue_index, &error)); + EXPECT_EQ( + output.str(), + "WEBVTT\n\n1\n00:00:01.000 --> 00:00:01.500\nhello \"world\"\n\n"); +} + +TEST(TranscriptExporterTest, WritesStructuredJson) +{ + std::ostringstream output; + size_t cue_index = 1; + std::string error; + + EXPECT_TRUE(WriteTranscriptOutput( + output, TranscriptOutputFormat::kJson, MakeResults(), "audio.wav", &cue_index, &error)); + EXPECT_NE(output.str().find("\"schema_version\":\"1.0\""), std::string::npos); + EXPECT_NE(output.str().find("\"text\":\"hello \\\"world\\\"\""), std::string::npos); + EXPECT_NE(output.str().find("\"transcript\":\"hello \\\"world\\\"\""), std::string::npos); + EXPECT_NE(output.str().find("\"start_time_ms\":1000"), std::string::npos); +} + +TEST(TranscriptExporterTest, RejectsSubtitlesWithoutWordTimes) +{ + Results result; + Results::Segment segment; + nr_asr::SpeechRecognitionAlternative alternative; + alternative.set_transcript("hello"); + segment.alternatives.push_back(alternative); + result.segments.push_back(segment); + std::ostringstream output; + size_t cue_index = 1; + std::string error; + + EXPECT_FALSE(WriteTranscriptOutput( + output, TranscriptOutputFormat::kSrt, result, "audio.wav", &cue_index, &error)); + EXPECT_NE(error.find("require word time offsets"), std::string::npos); +} + +TEST(TranscriptExporterTest, IgnoresJsonSegmentsWithoutAlternatives) +{ + Results result; + result.segments.emplace_back(); + std::ostringstream output; + size_t cue_index = 1; + std::string error; + + EXPECT_TRUE(WriteTranscriptOutput( + output, TranscriptOutputFormat::kJson, result, "audio.wav", &cue_index, &error)); + EXPECT_NE(output.str().find("\"segments\":[]"), std::string::npos); +} + +} // namespace