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
219 lines
8.2 KiB
C++
219 lines
8.2 KiB
C++
#include "engine/framework/audio/wav_reader.h"
|
|
#include "engine/framework/core/backend.h"
|
|
#include "engine/framework/io/filesystem.h"
|
|
#include "engine/framework/runtime/model.h"
|
|
#include "engine/framework/runtime/registry.h"
|
|
#include "engine/framework/runtime/session.h"
|
|
|
|
#include <algorithm>
|
|
#include <cstdio>
|
|
#include <filesystem>
|
|
#include <stdexcept>
|
|
#include <string>
|
|
|
|
#ifndef ENGINE_REPO_ROOT
|
|
#define ENGINE_REPO_ROOT "."
|
|
#endif
|
|
|
|
namespace {
|
|
|
|
constexpr int kExitPass = 0;
|
|
constexpr int kExitFail = 1;
|
|
constexpr int kExitSkip = 125;
|
|
constexpr const char* kExpectedText =
|
|
"Well, I don't wish to see it any more, observed Phoebe, turning away "
|
|
"her eyes. It is certainly very like the old portrait.";
|
|
|
|
std::filesystem::path repo_path(const std::string& relative) {
|
|
return std::filesystem::path(ENGINE_REPO_ROOT) / relative;
|
|
}
|
|
|
|
std::string arg_value(int argc, char** argv, const std::string& name, const std::string& fallback) {
|
|
for (int i = 1; i + 1 < argc; ++i) {
|
|
if (argv[i] == name) {
|
|
return argv[i + 1];
|
|
}
|
|
}
|
|
return fallback;
|
|
}
|
|
|
|
engine::runtime::TaskResult run_stream(
|
|
engine::runtime::IStreamingVoiceTaskSession& session,
|
|
const engine::audio::WavData& wav,
|
|
bool& saw_partial) {
|
|
engine::runtime::TaskRequest stream_request;
|
|
session.start_stream(stream_request);
|
|
constexpr size_t chunk_samples = 8000;
|
|
for (size_t offset = 0; offset < wav.samples.size(); offset += chunk_samples) {
|
|
const size_t count = std::min(chunk_samples, wav.samples.size() - offset);
|
|
engine::runtime::AudioChunk chunk;
|
|
chunk.sample_rate = wav.sample_rate;
|
|
chunk.channels = 1;
|
|
chunk.start_sample = static_cast<int64_t>(offset);
|
|
chunk.samples.assign(
|
|
wav.samples.begin() + static_cast<std::ptrdiff_t>(offset),
|
|
wav.samples.begin() + static_cast<std::ptrdiff_t>(offset + count));
|
|
const auto event = session.process_audio_chunk(chunk);
|
|
saw_partial = saw_partial ||
|
|
(event.partial_text.has_value() && !event.partial_text->text.empty());
|
|
}
|
|
return session.finalize();
|
|
}
|
|
|
|
} // namespace
|
|
|
|
int main(int argc, char** argv) {
|
|
const std::filesystem::path model_path = arg_value(
|
|
argc, argv, "--model", repo_path("models/parakeet-tdt-0.6b-v3").string());
|
|
const std::filesystem::path audio_path = arg_value(
|
|
argc, argv, "--audio", repo_path("tests/parakeet_tdt/assets/2086-149220-0033.wav").string());
|
|
const std::string weight_type = arg_value(argc, argv, "--weight-type", "q8_0");
|
|
const bool model_available =
|
|
engine::io::is_existing_file(model_path) ||
|
|
engine::io::is_existing_file(model_path / "config.json");
|
|
if (!model_available ||
|
|
!engine::io::is_existing_file(audio_path)) {
|
|
std::fprintf(stderr, "SKIP: Parakeet streaming test requires model weights\n");
|
|
return kExitSkip;
|
|
}
|
|
|
|
try {
|
|
const auto wav = engine::audio::read_wav_f32(audio_path);
|
|
auto registry = engine::runtime::make_default_registry();
|
|
engine::runtime::ModelLoadRequest load_request;
|
|
load_request.model_path = model_path;
|
|
load_request.family_hint = "parakeet_tdt";
|
|
auto model = registry.load(load_request);
|
|
const engine::runtime::TaskSpec task{
|
|
engine::runtime::VoiceTaskKind::Asr,
|
|
engine::runtime::RunMode::Streaming,
|
|
};
|
|
engine::runtime::SessionOptions options;
|
|
options.backend.type = engine::core::BackendType::Cpu;
|
|
options.backend.threads = 12;
|
|
options.options["parakeet_tdt.matmul_weight_type"] = weight_type;
|
|
options.options["parakeet_tdt.audio_chunk_duration_sec"] = "2";
|
|
options.options["parakeet_tdt.left_context_sec"] = "2";
|
|
options.options["parakeet_tdt.right_context_sec"] = "1";
|
|
|
|
auto require_rejected_session_option = [&](
|
|
const engine::runtime::TaskSpec& invalid_task,
|
|
const std::string& key,
|
|
const std::string& value) {
|
|
auto invalid_options = options;
|
|
invalid_options.options[key] = value;
|
|
bool rejected = false;
|
|
try {
|
|
(void)model->create_task_session(invalid_task, invalid_options);
|
|
} catch (const std::exception&) {
|
|
rejected = true;
|
|
}
|
|
if (!rejected) {
|
|
throw std::runtime_error(
|
|
"invalid session option was accepted: " + key + "=" + value);
|
|
}
|
|
};
|
|
require_rejected_session_option(
|
|
task,
|
|
"parakeet_tdt.full_context_max_duration_sec",
|
|
"30");
|
|
require_rejected_session_option(
|
|
task,
|
|
"parakeet_tdt.audio_chunk_duration_sec",
|
|
"0.0001");
|
|
require_rejected_session_option(
|
|
task,
|
|
"parakeet_tdt.audio_chunk_threshold_sec",
|
|
"0.0001");
|
|
require_rejected_session_option(
|
|
task,
|
|
"parakeet_tdt.offline_mode",
|
|
"typo");
|
|
require_rejected_session_option(
|
|
{
|
|
engine::runtime::VoiceTaskKind::Asr,
|
|
engine::runtime::RunMode::Offline,
|
|
},
|
|
"parakeet_tdt.streaming_attention_mode",
|
|
"typo");
|
|
|
|
auto base = model->create_task_session(task, options);
|
|
auto* session =
|
|
dynamic_cast<engine::runtime::IStreamingVoiceTaskSession*>(base.get());
|
|
if (session == nullptr) {
|
|
throw std::runtime_error("model did not create a streaming session");
|
|
}
|
|
engine::runtime::TaskRequest prepare_request;
|
|
prepare_request.audio_input = engine::runtime::AudioBuffer{
|
|
wav.sample_rate,
|
|
wav.channels,
|
|
wav.samples,
|
|
};
|
|
session->prepare(engine::runtime::build_preparation_request(prepare_request));
|
|
|
|
session->start_stream({});
|
|
bool discontinuity_failed = false;
|
|
try {
|
|
(void)session->process_audio_chunk({
|
|
wav.sample_rate,
|
|
1,
|
|
1,
|
|
{},
|
|
});
|
|
} catch (const std::exception&) {
|
|
discontinuity_failed = true;
|
|
}
|
|
if (!discontinuity_failed) {
|
|
throw std::runtime_error("non-contiguous stream chunk did not fail");
|
|
}
|
|
|
|
bool saw_partial = false;
|
|
const auto first = run_stream(*session, wav, saw_partial);
|
|
if (!saw_partial) {
|
|
throw std::runtime_error("stream produced no partial result before finalize");
|
|
}
|
|
const std::string first_text =
|
|
first.text_output.has_value() ? first.text_output->text : "";
|
|
if (first_text != kExpectedText) {
|
|
throw std::runtime_error("buffered-streaming final transcript mismatch");
|
|
}
|
|
|
|
bool repeated_finalize_failed = false;
|
|
try {
|
|
(void)session->finalize();
|
|
} catch (const std::exception&) {
|
|
repeated_finalize_failed = true;
|
|
}
|
|
if (!repeated_finalize_failed) {
|
|
throw std::runtime_error("repeated finalize did not fail");
|
|
}
|
|
bool post_finalize_chunk_failed = false;
|
|
try {
|
|
(void)session->process_audio_chunk({
|
|
wav.sample_rate,
|
|
1,
|
|
static_cast<int64_t>(wav.samples.size()),
|
|
{},
|
|
});
|
|
} catch (const std::exception&) {
|
|
post_finalize_chunk_failed = true;
|
|
}
|
|
if (!post_finalize_chunk_failed) {
|
|
throw std::runtime_error("post-finalize chunk did not fail");
|
|
}
|
|
|
|
bool second_saw_partial = false;
|
|
const auto second = run_stream(*session, wav, second_saw_partial);
|
|
const std::string second_text =
|
|
second.text_output.has_value() ? second.text_output->text : "";
|
|
if (second_text != first_text || !second_saw_partial) {
|
|
throw std::runtime_error("reset stream did not reproduce the first result");
|
|
}
|
|
|
|
std::printf("PASS: buffered streaming partials, finalize, and reset are stable\n");
|
|
return kExitPass;
|
|
} catch (const std::exception& ex) {
|
|
std::fprintf(stderr, "FAIL: %s\n", ex.what());
|
|
return kExitFail;
|
|
}
|
|
}
|