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
370 lines
17 KiB
C++
370 lines
17 KiB
C++
#include "engine/community_models/whistle_asr/assets.h"
|
|
#include "engine/community_models/whistle_asr/encoder.h"
|
|
#include "engine/community_models/whistle_asr/frontend.h"
|
|
#include "engine/community_models/whistle_asr/runtime.h"
|
|
#include "engine/framework/audio/wav_reader.h"
|
|
#include "engine/framework/core/execution_context.h"
|
|
|
|
#include <algorithm>
|
|
#include <chrono>
|
|
#include <cmath>
|
|
#include <cstdint>
|
|
#include <cstdlib>
|
|
#include <filesystem>
|
|
#include <fstream>
|
|
#include <iostream>
|
|
#include <limits>
|
|
#include <sstream>
|
|
#include <stdexcept>
|
|
#include <string>
|
|
#include <utility>
|
|
#include <vector>
|
|
|
|
namespace {
|
|
|
|
namespace whistle = engine::community_models::whistle_asr;
|
|
|
|
struct Options {
|
|
std::string gguf;
|
|
std::string wav;
|
|
std::string expected;
|
|
bool full = false;
|
|
bool use_neon = false;
|
|
std::string reference_dir;
|
|
std::string dump_dir;
|
|
// Allowed max |difference| as a fraction of the reference tensor's largest magnitude.
|
|
double tolerance = 2.0e-3;
|
|
engine::core::BackendConfig backend{engine::core::BackendType::Cpu, 0, 4};
|
|
};
|
|
|
|
Options parse(int argc, char ** argv) {
|
|
Options options;
|
|
std::vector<std::string> positional;
|
|
for (int index = 1; index < argc; ++index) {
|
|
const std::string argument = argv[index];
|
|
const auto value = [&]() -> std::string {
|
|
if (index + 1 >= argc) {
|
|
throw std::invalid_argument(argument + " needs a value");
|
|
}
|
|
return argv[++index];
|
|
};
|
|
if (argument == "--full") {
|
|
options.full = true;
|
|
} else if (argument == "--reference") {
|
|
options.reference_dir = value();
|
|
} else if (argument == "--dump") {
|
|
options.dump_dir = value();
|
|
} else if (argument == "--backend") {
|
|
const std::string backend = value();
|
|
if (backend == "cpu") {
|
|
options.backend.type = engine::core::BackendType::Cpu;
|
|
} else if (backend == "metal") {
|
|
options.backend.type = engine::core::BackendType::Metal;
|
|
} else {
|
|
throw std::invalid_argument("Unsupported test backend: " + backend);
|
|
}
|
|
} else if (argument == "--threads") {
|
|
options.backend.threads = std::stoi(value());
|
|
} else if (argument == "--decoder-impl") {
|
|
const auto impl = value();
|
|
if (impl != "native" && impl != "neon") throw std::invalid_argument("Unknown decoder implementation: " + impl);
|
|
options.use_neon = impl == "neon";
|
|
} else if (argument == "--tolerance") {
|
|
options.tolerance = std::stod(value());
|
|
} else {
|
|
positional.push_back(argument);
|
|
}
|
|
}
|
|
if (positional.empty() || positional.size() > 3) {
|
|
throw std::invalid_argument(
|
|
"Pass <GGUF> [16k-mono.wav] [expected transcript] [--full] "
|
|
"[--reference <dump-dir>] [--tolerance <fraction>] [--dump <dir>] (encoder tensors and decoder logits/tokens/transcript) [--backend cpu|metal] [--threads n] [--decoder-impl native|neon]");
|
|
}
|
|
options.gguf = positional[0];
|
|
if (positional.size() > 1) options.wav = positional[1];
|
|
if (positional.size() > 2) options.expected = positional[2];
|
|
return options;
|
|
}
|
|
|
|
engine::runtime::AudioBuffer read_audio(const std::string & path, bool full) {
|
|
auto wave = engine::audio::read_wav_f32(std::filesystem::path(path));
|
|
if (wave.sample_rate != 16000 || wave.channels != 1) {
|
|
throw std::invalid_argument("Whistle test requires 16 kHz mono WAV");
|
|
}
|
|
if (!full) {
|
|
wave.samples.resize(std::min<size_t>(wave.samples.size(), 3 * 16000));
|
|
}
|
|
engine::runtime::AudioBuffer audio;
|
|
audio.sample_rate = wave.sample_rate;
|
|
audio.channels = wave.channels;
|
|
audio.samples = std::move(wave.samples);
|
|
return audio;
|
|
}
|
|
|
|
std::vector<float> read_f32(const std::string & path) {
|
|
std::ifstream stream(path, std::ios::binary | std::ios::ate);
|
|
const auto bytes = static_cast<std::streamoff>(stream.tellg());
|
|
if (!stream || bytes < 0 || bytes % sizeof(float) != 0) {
|
|
throw std::runtime_error("Invalid float reference: " + path);
|
|
}
|
|
std::vector<float> values(static_cast<size_t>(bytes) / sizeof(float));
|
|
stream.seekg(0);
|
|
if (!stream.read(reinterpret_cast<char *>(values.data()),
|
|
static_cast<std::streamsize>(values.size() * sizeof(float)))) {
|
|
throw std::runtime_error("Cannot read float reference: " + path);
|
|
}
|
|
return values;
|
|
}
|
|
|
|
// Reports the largest absolute difference against a reference dump, scaled by the
|
|
// reference's largest magnitude, and fails when that fraction exceeds the tolerance.
|
|
void compare(const std::string & name, const std::vector<float> & actual, const std::string & path, double tolerance) {
|
|
const auto expected = read_f32(path);
|
|
if (expected.size() != actual.size()) {
|
|
throw std::runtime_error(name + " reference has " + std::to_string(expected.size()) +
|
|
" values, runtime produced " + std::to_string(actual.size()));
|
|
}
|
|
double max_abs = 0.0;
|
|
double scale = 0.0;
|
|
for (size_t index = 0; index < actual.size(); ++index) {
|
|
if (!std::isfinite(actual[index]) || !std::isfinite(expected[index])) {
|
|
throw std::runtime_error(name + " has a non-finite value at index " + std::to_string(index));
|
|
}
|
|
max_abs = std::max(max_abs, std::abs(static_cast<double>(actual[index]) - expected[index]));
|
|
scale = std::max(scale, std::abs(static_cast<double>(expected[index])));
|
|
}
|
|
const double fraction = scale > 0.0 ? max_abs / scale : max_abs;
|
|
std::cout << name << ": max_abs=" << max_abs << " reference_max=" << scale
|
|
<< " fraction=" << fraction << "\n";
|
|
if (!(fraction <= tolerance)) {
|
|
throw std::runtime_error(name + " differs from the reference beyond the tolerance");
|
|
}
|
|
}
|
|
|
|
void write_f32(const std::string & path, const std::vector<float> & values) {
|
|
std::ofstream stream(path, std::ios::binary);
|
|
if (!stream.write(reinterpret_cast<const char *>(values.data()),
|
|
static_cast<std::streamsize>(values.size() * sizeof(float)))) {
|
|
throw std::runtime_error("Cannot write float dump: " + path);
|
|
}
|
|
}
|
|
|
|
// Frontend and encoder tensors of the full recording, either written as raw
|
|
// little-endian f32 dumps (--dump) or compared against such dumps (--reference).
|
|
void check_encoder_tensors(const Options & options, const std::shared_ptr<const whistle::WhistleAssets> & assets,
|
|
engine::core::ExecutionContext & execution_context) {
|
|
const auto audio = read_audio(options.wav, true);
|
|
const whistle::WhistleFrontend frontend(assets->mel_filterbank);
|
|
const auto mel = frontend.extract(audio.samples);
|
|
whistle::WhistleEncoderRuntime encoder(assets, execution_context);
|
|
auto start = std::chrono::steady_clock::now();
|
|
const auto output = encoder.encode(mel);
|
|
const double first_ms = std::chrono::duration<double, std::milli>(std::chrono::steady_clock::now() - start).count();
|
|
start = std::chrono::steady_clock::now();
|
|
(void)encoder.encode(mel);
|
|
const double warm_ms = std::chrono::duration<double, std::milli>(std::chrono::steady_clock::now() - start).count();
|
|
std::cout << "encoder: " << output.frames << " frames, first " << first_ms << " ms, warm " << warm_ms << " ms\n";
|
|
const std::vector<std::pair<std::string, const std::vector<float> *>> tensors = {
|
|
{"mel", &mel.values},
|
|
{"memory", &output.memory},
|
|
{"cross_k_0", &output.cross_k[0]},
|
|
{"cross_v_0", &output.cross_v[0]},
|
|
{"cross_k_7", &output.cross_k[7]},
|
|
{"cross_v_7", &output.cross_v[7]},
|
|
};
|
|
if (!options.dump_dir.empty()) {
|
|
for (const auto & [name, values] : tensors) {
|
|
write_f32(options.dump_dir + "/" + name + ".f32", *values);
|
|
}
|
|
std::cout << "Wrote " << tensors.size() << " reference dumps to " << options.dump_dir << "\n";
|
|
}
|
|
if (!options.reference_dir.empty()) {
|
|
for (const auto & [name, values] : tensors) {
|
|
compare(name, *values, options.reference_dir + "/" + name + ".f32", options.tolerance);
|
|
}
|
|
}
|
|
}
|
|
|
|
std::vector<int32_t> read_i32(const std::string & path) {
|
|
std::ifstream stream(path, std::ios::binary | std::ios::ate);
|
|
const auto bytes = static_cast<std::streamoff>(stream.tellg());
|
|
if (!stream || bytes < 0 || bytes % sizeof(int32_t) != 0) {
|
|
throw std::runtime_error("Invalid token reference: " + path);
|
|
}
|
|
std::vector<int32_t> values(static_cast<size_t>(bytes) / sizeof(int32_t));
|
|
stream.seekg(0);
|
|
if (!stream.read(reinterpret_cast<char *>(values.data()),
|
|
static_cast<std::streamsize>(values.size() * sizeof(int32_t)))) {
|
|
throw std::runtime_error("Cannot read token reference: " + path);
|
|
}
|
|
return values;
|
|
}
|
|
|
|
std::string read_text(const std::string & path) {
|
|
std::ifstream stream(path, std::ios::binary);
|
|
if (!stream) {
|
|
throw std::runtime_error("Cannot read transcript reference: " + path);
|
|
}
|
|
std::ostringstream contents;
|
|
contents << stream.rdbuf();
|
|
return contents.str();
|
|
}
|
|
|
|
// Per-step decoder logits (decoder_logits_<position>.f32), the input token sequence
|
|
// (decoder_tokens.i32, BOS and language token first) and the transcript of the full
|
|
// recording, written by --dump or required to match by --reference.
|
|
void check_decoder(const Options & options, const std::shared_ptr<const whistle::WhistleAssets> & assets,
|
|
engine::core::ExecutionContext & execution_context) {
|
|
const auto audio = read_audio(options.wav, true);
|
|
whistle::WhistleRuntime runtime(assets, execution_context, options.use_neon);
|
|
std::vector<whistle::WhistleDecodeStep> steps;
|
|
const auto transcript = runtime.transcribe(audio, "", [&](const whistle::WhistleDecodeStep & step) {
|
|
steps.push_back(step);
|
|
});
|
|
std::vector<int32_t> tokens;
|
|
for (const auto & step : steps) {
|
|
tokens.push_back(step.input_token);
|
|
}
|
|
const auto logits_path = [](const std::string & dir, size_t position) {
|
|
return dir + "/decoder_logits_" + std::to_string(position) + ".f32";
|
|
};
|
|
if (!options.dump_dir.empty()) {
|
|
for (const auto & step : steps) {
|
|
write_f32(logits_path(options.dump_dir, step.position), step.logits);
|
|
}
|
|
std::ofstream token_stream(options.dump_dir + "/decoder_tokens.i32", std::ios::binary);
|
|
if (!token_stream.write(reinterpret_cast<const char *>(tokens.data()),
|
|
static_cast<std::streamsize>(tokens.size() * sizeof(int32_t)))) {
|
|
throw std::runtime_error("Cannot write decoder tokens");
|
|
}
|
|
std::ofstream text_stream(options.dump_dir + "/decoder_transcript.txt", std::ios::binary);
|
|
if (!text_stream.write(transcript.text.data(), static_cast<std::streamsize>(transcript.text.size()))) {
|
|
throw std::runtime_error("Cannot write decoder transcript");
|
|
}
|
|
}
|
|
if (!options.reference_dir.empty()) {
|
|
if (read_i32(options.reference_dir + "/decoder_tokens.i32") != tokens) {
|
|
throw std::runtime_error("Decoder token sequence differs from the reference");
|
|
}
|
|
if (read_text(options.reference_dir + "/decoder_transcript.txt") != transcript.text) {
|
|
throw std::runtime_error("Decoder transcript differs from the reference");
|
|
}
|
|
for (const auto & step : steps) {
|
|
compare("decoder_logits_" + std::to_string(step.position), step.logits,
|
|
logits_path(options.reference_dir, step.position), options.tolerance);
|
|
}
|
|
}
|
|
std::cout << "decoder: " << steps.size() << " steps, transcript: " << transcript.text << "\n";
|
|
}
|
|
|
|
void check_transcription(const Options & options, const std::shared_ptr<const whistle::WhistleAssets> & assets,
|
|
engine::core::ExecutionContext & execution_context) {
|
|
const auto audio = read_audio(options.wav, options.full);
|
|
const auto start = std::chrono::steady_clock::now();
|
|
whistle::WhistleRuntime runtime(assets, execution_context, options.use_neon);
|
|
const auto transcript = runtime.transcribe(audio);
|
|
std::cout << "Whistle: " << transcript.text << " [" << transcript.language << "] in "
|
|
<< std::chrono::duration_cast<std::chrono::milliseconds>(
|
|
std::chrono::steady_clock::now() - start).count() << " ms\n";
|
|
if (!options.expected.empty() && transcript.text != options.expected) {
|
|
throw std::runtime_error("Whistle transcript differs from the expected speech");
|
|
}
|
|
const auto repeated = runtime.transcribe(audio);
|
|
if (repeated.text != transcript.text || repeated.language != transcript.language) {
|
|
throw std::runtime_error("Whistle retained decoder state across recordings");
|
|
}
|
|
auto quiet = audio;
|
|
for (float & sample : quiet.samples) {
|
|
sample *= 0.0005f;
|
|
}
|
|
if (runtime.transcribe(quiet).text.rfind("Some call me", 0) != 0) {
|
|
throw std::runtime_error("Whistle dropped quiet speech despite front-end AGC");
|
|
}
|
|
auto constant_rms = audio;
|
|
constant_rms.samples.resize(std::min<size_t>(constant_rms.samples.size(), 3 * 16000));
|
|
constexpr size_t kRmsBlockSamples = 320;
|
|
constexpr float kTargetRms = 1000.0f / 32768.0f;
|
|
for (size_t offset = 0; offset < constant_rms.samples.size(); offset += kRmsBlockSamples) {
|
|
const size_t count = std::min(kRmsBlockSamples, constant_rms.samples.size() - offset);
|
|
double sum_squares = 0.0;
|
|
for (size_t index = 0; index < count; ++index) {
|
|
const double sample = constant_rms.samples[offset + index];
|
|
sum_squares += sample * sample;
|
|
}
|
|
const float rms = static_cast<float>(std::sqrt(sum_squares / count));
|
|
if (rms > 0.0f) {
|
|
const float scale = kTargetRms / rms;
|
|
for (size_t index = 0; index < count; ++index) {
|
|
constant_rms.samples[offset + index] *= scale;
|
|
}
|
|
}
|
|
}
|
|
if (runtime.transcribe(constant_rms).text != "Some call me nature.") {
|
|
throw std::runtime_error("Whistle dropped speech with constant block RMS");
|
|
}
|
|
const auto silence = runtime.transcribe(
|
|
engine::runtime::AudioBuffer{16000, 1, std::vector<float>(16000, 0.0f)});
|
|
if (!silence.text.empty() || !silence.language.empty()) {
|
|
throw std::runtime_error("Whistle invented words for silence");
|
|
}
|
|
bool rejected = false;
|
|
try {
|
|
(void)runtime.transcribe(engine::runtime::AudioBuffer{
|
|
16000, 1, std::vector<float>(30 * 16000 + 1, 0.0f)});
|
|
} catch (const std::invalid_argument &) {
|
|
rejected = true;
|
|
}
|
|
if (!rejected) {
|
|
throw std::runtime_error("Whistle accepted audio longer than 30 seconds");
|
|
}
|
|
rejected = false;
|
|
try {
|
|
(void)runtime.transcribe(engine::runtime::AudioBuffer{
|
|
16000, 1, std::vector<float>(1, std::numeric_limits<float>::quiet_NaN())});
|
|
} catch (const std::invalid_argument &) {
|
|
rejected = true;
|
|
}
|
|
if (!rejected) {
|
|
throw std::runtime_error("Whistle accepted a non-finite short recording");
|
|
}
|
|
}
|
|
|
|
} // namespace
|
|
|
|
int main(int argc, char ** argv) {
|
|
try {
|
|
const Options options = parse(argc, argv);
|
|
const auto assets = whistle::load_whistle_assets(options.gguf);
|
|
if (assets->mel_filterbank.size() != 257 * 80 ||
|
|
assets->hadamard_permutations[0][0] != 396 ||
|
|
assets->hadamard_permutations[1][0] != 458 ||
|
|
whistle::decode_whistle_tokens(assets->tokenizer_pieces, {3633, 1313, 409, 8116}) != " Some call me.") {
|
|
throw std::runtime_error("Whistle GGUF assets do not match the reference model");
|
|
}
|
|
std::cout << "Loaded " << assets->weights->tensors().size()
|
|
<< " GGUF tensors and " << assets->tokenizer_pieces.size()
|
|
<< " tokenizer pieces\n";
|
|
if (options.wav.empty()) {
|
|
if (!options.reference_dir.empty() || !options.dump_dir.empty()) {
|
|
throw std::invalid_argument("--dump and --reference need a WAV file");
|
|
}
|
|
return 0;
|
|
}
|
|
engine::core::ExecutionContext execution_context(options.backend);
|
|
if (!options.reference_dir.empty() || !options.dump_dir.empty()) {
|
|
if (!options.reference_dir.empty() && !options.dump_dir.empty() &&
|
|
std::filesystem::weakly_canonical(options.reference_dir) ==
|
|
std::filesystem::weakly_canonical(options.dump_dir)) {
|
|
throw std::invalid_argument("--dump would overwrite the --reference directory");
|
|
}
|
|
check_encoder_tensors(options, assets, execution_context);
|
|
check_decoder(options, assets, execution_context);
|
|
return 0;
|
|
}
|
|
check_transcription(options, assets, execution_context);
|
|
return 0;
|
|
} catch (const std::exception & error) {
|
|
std::cerr << error.what() << '\n';
|
|
return 1;
|
|
}
|
|
}
|