Fabio
c36510c21e
Some checks failed
Linux Build / Linux cpu (push) Has been cancelled
Linux Build / Linux cpu-clang (push) Has been cancelled
Linux Build / Linux vulkan (push) Has been cancelled
mac-build / macOS CPU (push) Has been cancelled
Nix Build / Nix (cpu) (push) Has been cancelled
Nix Build / Nix (vulkan) (push) Has been cancelled
windows-build / Windows CPU (push) Has been cancelled
139 lines
7.7 KiB
C++
139 lines
7.7 KiB
C++
#include "engine/community_models/parakeet_tdt/assets.h"
|
|
#include "engine/community_models/parakeet_tdt/frontend.h"
|
|
#include "engine/community_models/sortformer_diar/assets.h"
|
|
#include "engine/community_models/sortformer_diar/frontend.h"
|
|
#include "engine/community_models/granite5asr/frontend.h"
|
|
#include "engine/framework/audio/conversion.h"
|
|
#include "engine/framework/audio/wav_reader.h"
|
|
#include "engine/framework/debug/trace.h"
|
|
#include "engine/models/hviske_asr/assets.h"
|
|
#include "engine/models/hviske_asr/frontend.h"
|
|
#include "engine/models/citrinet_asr/runtime.h"
|
|
#include "engine/models/canary_asr/model.h"
|
|
#include "engine/models/cohere_asr/model.h"
|
|
#include "engine/models/sortformer_diar/assets.h"
|
|
#include "engine/models/sortformer_diar/frontend.h"
|
|
#include "engine/models/nemotron_asr/frontend.h"
|
|
|
|
#include <filesystem>
|
|
#include <fstream>
|
|
#include <iostream>
|
|
#include <stdexcept>
|
|
#include <string>
|
|
#include <vector>
|
|
|
|
namespace {
|
|
|
|
std::string argument(int argc, char ** argv, const std::string & name) {
|
|
for (int i = 1; i + 1 < argc; ++i) {
|
|
if (argv[i] == name) return argv[i + 1];
|
|
}
|
|
throw std::runtime_error("missing required argument: " + name);
|
|
}
|
|
|
|
void write_features(
|
|
const std::filesystem::path & prefix,
|
|
const std::vector<float> & values,
|
|
int64_t frames,
|
|
int64_t valid_frames,
|
|
int64_t features,
|
|
const char * layout) {
|
|
if (frames < 0 || valid_frames < 0 || valid_frames > frames || features <= 0 ||
|
|
values.size() != static_cast<size_t>(frames * features)) {
|
|
throw std::runtime_error("frontend returned inconsistent shape or valid length");
|
|
}
|
|
std::ofstream data(prefix.string() + ".f32", std::ios::binary | std::ios::trunc);
|
|
if (!data) throw std::runtime_error("cannot open feature output");
|
|
data.write(reinterpret_cast<const char *>(values.data()),
|
|
static_cast<std::streamsize>(values.size() * sizeof(float)));
|
|
if (!data) throw std::runtime_error("cannot write feature output");
|
|
std::ofstream meta(prefix.string() + ".meta", std::ios::trunc);
|
|
if (!meta) throw std::runtime_error("cannot open metadata output");
|
|
meta << frames << ' ' << valid_frames << ' ' << features << ' ' << layout << '\n';
|
|
if (!meta) throw std::runtime_error("cannot write metadata output");
|
|
std::cout << prefix << ": frames=" << frames << " valid=" << valid_frames
|
|
<< " features=" << features << " layout=" << layout << '\n';
|
|
}
|
|
|
|
} // namespace
|
|
|
|
int main(int argc, char ** argv) try {
|
|
const auto family = argument(argc, argv, "--family");
|
|
const std::filesystem::path model = argument(argc, argv, "--model");
|
|
const std::filesystem::path audio_path = argument(argc, argv, "--audio");
|
|
const std::filesystem::path prefix = argument(argc, argv, "--out");
|
|
engine::debug::configure_logging({true, argument(argc, argv, "--log")});
|
|
const auto wav = engine::audio::read_wav_f32(audio_path);
|
|
const engine::runtime::AudioBuffer input{wav.sample_rate, wav.channels, wav.samples};
|
|
|
|
if (family == "parakeet_tdt") {
|
|
engine::community_models::parakeet_tdt::ParakeetFrontend frontend(
|
|
engine::community_models::parakeet_tdt::load_parakeet_assets(model));
|
|
for (const bool centered : {true, false}) {
|
|
const auto features = frontend.extract(input, centered);
|
|
write_features(prefix.string() + (centered ? ".center" : ".no_center"),
|
|
features.values, features.frames, features.valid_frames,
|
|
features.feature_dim, "time_major");
|
|
}
|
|
} else if (family == "sortformer_diar") {
|
|
const auto assets = engine::models::sortformer_diar::load_sortformer_assets(model);
|
|
const auto features = engine::models::sortformer_diar::compute_sortformer_features(input, *assets, 8);
|
|
write_features(prefix, features.time_major, features.frames, features.valid_frames,
|
|
assets->feature_config.num_mel_bins, "time_major");
|
|
} else if (family == "sortformer_diar_v2") {
|
|
const auto assets = engine::community_models::sortformer_diar::load_sortformer_v2_assets(
|
|
model, model.parent_path() / "model_spec.json");
|
|
const auto offline = engine::community_models::sortformer_diar::compute_sortformer_v2_features(input, *assets, 8);
|
|
write_features(prefix.string() + ".offline", offline.time_major, offline.frames,
|
|
offline.valid_frames, assets->feature_config.num_mel_bins, "time_major");
|
|
const auto mono = engine::audio::mixdown_interleaved_to_mono_average(input.samples, input.channels);
|
|
const auto streaming = engine::community_models::sortformer_diar::compute_sortformer_v2_stream_features(
|
|
mono, *assets, 8, false);
|
|
write_features(prefix.string() + ".stream", streaming.time_major, streaming.frames,
|
|
streaming.valid_frames, assets->feature_config.num_mel_bins, "time_major");
|
|
} else if (family == "hviske_asr") {
|
|
engine::models::hviske_asr::HviskeFrontend frontend(
|
|
engine::models::hviske_asr::load_hviske_asr_assets(model));
|
|
const auto features = frontend.extract(input);
|
|
write_features(prefix, features.values, features.frames, features.valid_frames,
|
|
features.feature_dim, "feature_major");
|
|
} else if (family == "citrinet_asr") {
|
|
const auto weights = engine::models::citrinet_asr::load_citrinet_weights_cached(model);
|
|
const auto features = engine::models::citrinet_asr::extract_citrinet_frontend(input, *weights);
|
|
write_features(prefix, features.values, features.padded_frames, features.raw_frames,
|
|
weights->config.n_mels, "time_major");
|
|
} else if (family == "canary_asr") {
|
|
const auto assets = engine::models::canary_asr::load_canary_assets(model);
|
|
const auto mono = engine::audio::mixdown_interleaved_to_mono_average(input.samples, input.channels);
|
|
const auto features = engine::models::canary_asr::extract_canary_frontend(mono, *assets, 8);
|
|
write_features(prefix, features.values, features.raw_frames, features.valid_frames,
|
|
128, "feature_major");
|
|
} else if (family == "cohere_asr") {
|
|
const auto assets = engine::models::cohere_asr::load_cohere_assets(model);
|
|
const auto mono = engine::audio::mixdown_interleaved_to_mono_average(input.samples, input.channels);
|
|
const auto features = engine::models::cohere_asr::extract_cohere_frontend(mono, *assets, 8);
|
|
write_features(prefix, features.values, features.shape.at(2),
|
|
static_cast<int64_t>(mono.size()) / 160, 128, "feature_major");
|
|
} else if (family == "nemotron_asr") {
|
|
engine::models::nemotron_asr::NemotronFrontend frontend(
|
|
engine::models::nemotron_asr::load_nemotron_asr_assets(model));
|
|
for (const bool centered : {true, false}) {
|
|
const auto features = frontend.extract(input, centered);
|
|
write_features(prefix.string() + (centered ? ".center" : ".no_center"),
|
|
features.values, features.frames, features.valid_frames,
|
|
features.feature_dim, "time_major");
|
|
}
|
|
} else if (family == "granite5asr") {
|
|
engine::community_models::granite5asr::Granite5Frontend frontend(
|
|
engine::community_models::granite5asr::load_granite5asr_assets(model));
|
|
const auto features = frontend.extract(input);
|
|
write_features(prefix, features.values, features.frames, features.frames,
|
|
features.feature_dim, "time_major");
|
|
} else {
|
|
throw std::runtime_error("unsupported family: " + family);
|
|
}
|
|
return 0;
|
|
} catch (const std::exception & error) {
|
|
std::cerr << "nemo_mel_migration_probe: " << error.what() << '\n';
|
|
return 1;
|
|
}
|