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
130 lines
8.5 KiB
C++
130 lines
8.5 KiB
C++
#include "engine/community_models/parakeet_tdt/hotwords.h"
|
|
#include "engine/framework/audio/wav_reader.h"
|
|
#include "engine/framework/io/json.h"
|
|
#include "engine/framework/runtime/registry.h"
|
|
#include "engine/framework/runtime/session.h"
|
|
|
|
#include <algorithm>
|
|
#include <chrono>
|
|
#include <iostream>
|
|
#include <limits>
|
|
#include <stdexcept>
|
|
|
|
namespace j = engine::io::json;
|
|
using namespace engine::runtime;
|
|
using namespace engine::community_models::parakeet_tdt;
|
|
j::Value words_json(const std::vector<WordTimestamp> & words) {
|
|
j::Value::Array out;
|
|
for (const auto & w : words) out.push_back(j::Value::make_object({
|
|
{"word",j::Value::make_string(w.word)}, {"start_sample",j::Value::make_number(w.span.start_sample)},
|
|
{"end_sample",j::Value::make_number(w.span.end_sample)}}));
|
|
return j::Value::make_array(std::move(out));
|
|
}
|
|
int main(int argc, char ** argv) {
|
|
try {
|
|
if (argc != 2) throw std::runtime_error("usage: phonon_features_probe input.json");
|
|
const auto input = j::parse_file(argv[1]);
|
|
if (const auto * queries = input.find("queries")) {
|
|
std::vector<std::string> vocab;
|
|
for (const auto & v : input.require("vocabulary").as_array()) vocab.push_back(v.as_string());
|
|
PhononHotwordAutomaton a(parse_phonon_hotwords(j::stringify(input.require("hotwords"))), vocab,
|
|
static_cast<int32_t>(input.require("blank").as_i64()), input.require("strength").as_f32());
|
|
j::Value::Array rows;
|
|
for (const auto & query : queries->as_array()) {
|
|
int32_t state = 0;
|
|
for (const auto & t : query.require("prefix").as_array()) state = a.step(state, static_cast<int32_t>(t.as_i64()));
|
|
j::Value::Array bonuses;
|
|
for (size_t t = 0; t < vocab.size(); ++t) bonuses.push_back(j::Value::make_number(a.bonus(state, static_cast<int32_t>(t))));
|
|
rows.push_back(j::Value::make_object({{"state",j::Value::make_number(state)}, {"bonus",j::Value::make_array(std::move(bonuses))}}));
|
|
}
|
|
std::cout << j::stringify(j::Value::make_object({{"states",j::Value::make_number(a.states())}, {"rows",j::Value::make_array(std::move(rows))}}));
|
|
return 0;
|
|
}
|
|
auto registry = make_default_registry();
|
|
ModelLoadRequest load; load.family_hint = "parakeet_tdt"; load.model_path = input.require("model").as_string();
|
|
auto model = registry.load(load);
|
|
SessionOptions options; options.backend.threads = 8;
|
|
const auto backend = input.require("backend").as_string();
|
|
options.backend.type = backend == "cuda" ? engine::core::BackendType::Cuda :
|
|
backend == "vulkan" ? engine::core::BackendType::Vulkan : engine::core::BackendType::Cpu;
|
|
if (const auto * device = input.find("device")) options.backend.device = static_cast<int>(device->as_i64());
|
|
const bool streaming = input.require("streaming").as_bool();
|
|
auto session = model->create_task_session({VoiceTaskKind::Asr, streaming ? RunMode::Streaming : RunMode::Offline}, options);
|
|
const auto wav = engine::audio::read_wav_f32(std::filesystem::path(input.require("audio").as_string()));
|
|
TaskRequest request; request.audio_input = AudioBuffer{wav.sample_rate, wav.channels, wav.samples};
|
|
if (const auto * terms = input.find("hotwords")) request.options["hotwords"] = j::stringify(*terms);
|
|
if (const auto * strength = input.find("strength")) request.options["hotwords_score"] = j::stringify(*strength);
|
|
if (const auto * cap = input.find("max_tokens")) request.options["max_tokens"] = j::stringify(*cap);
|
|
session->prepare(build_preparation_request(request));
|
|
auto start = std::chrono::steady_clock::now();
|
|
j::Value::Array events;
|
|
TaskResult result;
|
|
if (streaming) {
|
|
auto * stream = dynamic_cast<IStreamingVoiceTaskSession *>(session.get());
|
|
stream->set_stream_event_sink([&](const StreamEvent & e) {
|
|
j::Value::Object row;
|
|
if (e.partial_text) row["delta"] = j::Value::make_string(e.partial_text->text);
|
|
if (e.partial_text_snapshot) row["snapshot"] = j::Value::make_string(e.partial_text_snapshot->text);
|
|
row["words"] = words_json(e.word_timestamps);
|
|
events.push_back(j::Value::make_object(std::move(row)));
|
|
});
|
|
stream->start_stream(request);
|
|
size_t chunk = 800;
|
|
if (const auto * size = input.find("chunk_samples")) chunk = static_cast<size_t>(size->as_i64());
|
|
if (!chunk) throw std::runtime_error("chunk size must be positive");
|
|
for (size_t offset = 0; offset < wav.samples.size(); offset += chunk) {
|
|
auto end = std::min(offset + chunk, wav.samples.size());
|
|
stream->process_audio_chunk({wav.sample_rate, wav.channels, static_cast<int64_t>(offset),
|
|
std::vector<float>(wav.samples.begin() + offset, wav.samples.begin() + end)});
|
|
}
|
|
result = stream->finalize();
|
|
} else result = dynamic_cast<IOfflineVoiceTaskSession *>(session.get())->run(request);
|
|
if (const auto * lifecycle = input.find("lifecycle"); streaming && lifecycle && lifecycle->as_bool()) {
|
|
auto * stream = dynamic_cast<IStreamingVoiceTaskSession *>(session.get());
|
|
const auto expected = result.text_output->text;
|
|
const auto expected_words = j::stringify(words_json(result.word_timestamps));
|
|
auto rejected = [](auto operation) {
|
|
bool failed = false;
|
|
try { operation(); } catch (const std::exception &) { failed = true; }
|
|
if (!failed) throw std::runtime_error("expected lifecycle rejection");
|
|
};
|
|
rejected([&] { stream->finalize(); });
|
|
rejected([&] { stream->process_audio_chunk({16000, 1, 0, {0.f}}); });
|
|
auto replay = [&](size_t chunk) {
|
|
for (size_t offset = 0; offset < wav.samples.size(); offset += chunk) {
|
|
const auto end = std::min(offset + chunk, wav.samples.size());
|
|
stream->process_audio_chunk({wav.sample_rate, wav.channels, static_cast<int64_t>(offset),
|
|
std::vector<float>(wav.samples.begin() + offset, wav.samples.begin() + end)});
|
|
}
|
|
const auto repeated = stream->finalize();
|
|
if (repeated.text_output->text != expected || j::stringify(words_json(repeated.word_timestamps)) != expected_words)
|
|
throw std::runtime_error("stream chunk/reuse result changed");
|
|
};
|
|
stream->set_stream_event_sink({});
|
|
for (auto chunk : {size_t{13}, size_t{2049}, wav.samples.size()}) {
|
|
stream->start_stream(request);
|
|
rejected([&] { stream->process_audio_chunk({16000, 1, 1, {0.f}}); });
|
|
replay(chunk); // Bad ordering is validated before mutation.
|
|
}
|
|
stream->start_stream(request);
|
|
rejected([&] { stream->process_audio_chunk({16000, 1, 0, std::vector<float>(800, std::numeric_limits<float>::quiet_NaN())}); });
|
|
rejected([&] { stream->finalize(); });
|
|
stream->start_stream(request); replay(800);
|
|
stream->set_stream_event_sink([](const StreamEvent &) { throw std::runtime_error("injected sink failure"); });
|
|
stream->start_stream(request);
|
|
rejected([&] { replay(800); });
|
|
stream->set_stream_event_sink({});
|
|
stream->start_stream(request); replay(800);
|
|
// Idle audio must remain bounded, yield no hallucinated segments,
|
|
// and permit immediate reuse after finalization.
|
|
stream->start_stream(request);
|
|
stream->process_audio_chunk({16000, 1, 0, std::vector<float>(33 * 16000, 0.f)});
|
|
const auto silent = stream->finalize();
|
|
if (!silent.text_output->text.empty() || !silent.word_timestamps.empty()) throw std::runtime_error("silent stream produced words");
|
|
stream->start_stream(request); replay(800);
|
|
}
|
|
const auto elapsed = std::chrono::duration<double>(std::chrono::steady_clock::now() - start).count();
|
|
std::cout << j::stringify(j::Value::make_object({{"text",j::Value::make_string(result.text_output->text)},
|
|
{"words",words_json(result.word_timestamps)}, {"events",j::Value::make_array(std::move(events))}, {"seconds",j::Value::make_number(elapsed)}})) << '\n';
|
|
} catch (const std::exception & e) { std::cerr << e.what() << '\n'; return 1; }
|
|
}
|