diff --git a/CMakeLists.txt b/CMakeLists.txt index 01473219d..8282f6118 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -141,6 +141,12 @@ option(LLAMA_BUILD_PARADEE "Build the Paradee-8M TTS subtree (tools/paradee/)" # subtree lives in tools/misaki-g2p/ and is added when present. option(LLAMA_BUILD_MISAKI_G2P "Build the misaki G2P subtree (tools/misaki-g2p/)" ON) +# Parakeet-TDT ASR (FastConformer + TDT) for FermionResearch/Phonon-2, the +# five-value re-quantisation of nvidia/parakeet-tdt-0.6b-v3 (weights +# CC-BY-4.0). Produces parakeet_lib + parakeet-cli/parakeet-bench; the fused +# `elizainference` target folds parakeet_lib in with `if(TARGET ...)`. +option(LLAMA_BUILD_PARAKEET "Build the Parakeet-TDT ASR subtree (tools/parakeet/)" ON) + # W3-3 deprecation: ELIZA_FUSE_OMNIVOICE=ON is the pre-W3-3 flag for the # fork-root `omnivoice/` graft. The graft is gone (see the deprecation # block below); turn the flag into a redirect now so the option ON state @@ -312,6 +318,12 @@ if (LLAMA_BUILD_PARADEE AND NOT (LLAMA_BUILD_COMMON AND LLAMA_BUILD_TOOLS)) endif() endif() +# Parakeet-TDT ASR — embed-as-library hook (mirrors kokoro above); precedes +# tools/omnivoice so the fused target can fold parakeet_lib in. +if (LLAMA_BUILD_PARAKEET AND NOT (LLAMA_BUILD_COMMON AND LLAMA_BUILD_TOOLS)) + add_subdirectory(tools/parakeet) +endif() + if (LLAMA_BUILD_OMNIVOICE) if (NOT EXISTS "${CMAKE_CURRENT_SOURCE_DIR}/tools/omnivoice/CMakeLists.txt") message(FATAL_ERROR diff --git a/tools/CMakeLists.txt b/tools/CMakeLists.txt index 08d148989..b10beca53 100644 --- a/tools/CMakeLists.txt +++ b/tools/CMakeLists.txt @@ -61,6 +61,12 @@ else() message(STATUS "LLAMA_BUILD_PARADEE: skipped because LLAMA_BUILD_KOKORO=OFF") endif() endif() + if (LLAMA_BUILD_PARAKEET) + # Parakeet-TDT ASR (Phonon-2 five-value weights). Like kokoro it + # precedes omnivoice so the fused `elizainference` target can fold + # parakeet_lib in. + add_subdirectory(parakeet) + endif() if (LLAMA_BUILD_OMNIVOICE) # The merged omnivoice subtree lives under tools/omnivoice/. The # top-level CMakeLists.txt has a parallel add_subdirectory hook for diff --git a/tools/espeak-shared/eliza-espeak-shared.h b/tools/espeak-shared/eliza-espeak-shared.h new file mode 100644 index 000000000..535129a73 --- /dev/null +++ b/tools/espeak-shared/eliza-espeak-shared.h @@ -0,0 +1,93 @@ +// SPDX-License-Identifier: MIT +// +// eliza-espeak-shared.h: the single owner of libespeak-ng's process-global +// state for every component in this tree that calls it (tools/kokoro and +// tools/misaki-g2p). +// +// libespeak-ng keeps its data path, loaded voice and translator in globals +// and is not thread-safe. Two components that each guard it with their own +// mutex and their own "initialised" flag race each other and can leave the +// other one running on a voice it did not select. Here there is one lock and +// one record of what libespeak-ng was initialised with: +// +// auto lock = eliza_espeak_shared::lock(); +// if (!eliza_espeak_shared::prepare(lock, data_dir, "en-us", err)) ...; +// espeak_TextToPhonemes(...); // still under `lock` +// +// prepare() initialises libespeak-ng on first use and re-selects the +// caller's voice whenever another component changed it, so every caller gets +// its own voice regardless of interleaving. Per-call options (phoneme mode, +// tie character) are arguments of espeak_TextToPhonemes and need no reset. +// +// Data path: an empty `data_dir` means "whatever the process already +// initialised with, else libespeak-ng's default" (compiled-in path or +// ESPEAK_DATA_PATH); Kokoro always asks for this. A non-empty, explicitly +// configured path (misaki_g2p_open's espeak_data_dir) wins: when it differs +// from the current one libespeak-ng is re-initialised with it, and later +// default-path callers keep using it. +// +// Header-only on purpose: the state lives in a C++17 inline function's +// static, which the linker folds into one instance per program or shared +// library, so kokoro_lib and misaki_g2p share it without a CMake target +// between them (each stays buildable on its own, with or without the other) +// and without one more static archive for the mobile slices to collect. +// Include it only when the component links libespeak-ng. +#pragma once + +#include + +#include +#include + +namespace eliza_espeak_shared { + +struct State { + std::mutex mutex; + bool initialized = false; + std::string data_dir; // "" = libespeak-ng's default path + std::string voice; // voice currently selected, "" when unknown +}; + +inline State & state() { + static State s; + return s; +} + +// Every libespeak-ng call in the process must be made while holding this. +inline std::unique_lock lock() { + return std::unique_lock(state().mutex); +} + +// With `held` from lock(): make libespeak-ng ready for `data_dir` (see the +// data-path rule above) and select `voice`. On failure sets `err`. +inline bool prepare(const std::unique_lock & held, const std::string & data_dir, const char * voice, + std::string & err) { + State & s = state(); + if (!held.owns_lock() || held.mutex() != &s.mutex) { + err = "eliza_espeak_shared::prepare called without the shared eSpeak lock"; + return false; + } + if (!s.initialized || (!data_dir.empty() && data_dir != s.data_dir)) { + s.initialized = false; + s.voice.clear(); + // AUDIO_OUTPUT_SYNCHRONOUS: no audio device is opened. + const int rate = espeak_Initialize(AUDIO_OUTPUT_SYNCHRONOUS, 0, data_dir.empty() ? nullptr : data_dir.c_str(), 0); + if (rate <= 0) { + err = "espeak_Initialize failed for data dir '" + data_dir + "'"; + return false; + } + s.initialized = true; + s.data_dir = data_dir; + } + if (s.voice != voice) { + if (espeak_SetVoiceByName(voice) != EE_OK) { + s.voice.clear(); + err = std::string("eSpeak NG voice ") + voice + " not found under '" + s.data_dir + "'"; + return false; + } + s.voice = voice; + } + return true; +} + +} // namespace eliza_espeak_shared diff --git a/tools/kokoro/CMakeLists.txt b/tools/kokoro/CMakeLists.txt index c34869e7a..b5831a812 100644 --- a/tools/kokoro/CMakeLists.txt +++ b/tools/kokoro/CMakeLists.txt @@ -85,6 +85,9 @@ if(KOKORO_ENABLE_ESPEAK) target_include_directories(kokoro_lib PRIVATE ${ESPEAK_NG_INCLUDE_DIR}) target_link_libraries(kokoro_lib PRIVATE ${ESPEAK_NG_LIBRARY}) target_compile_definitions(kokoro_lib PRIVATE KOKORO_USE_ESPEAK) + # One lock + init owner for the process-global libespeak-ng, shared + # with tools/misaki-g2p (header-only; see tools/espeak-shared). + target_include_directories(kokoro_lib PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/../espeak-shared) message(STATUS "Kokoro G2P: libespeak-ng found (${ESPEAK_NG_LIBRARY}) — real IPA path enabled") else() message(STATUS "Kokoro G2P: libespeak-ng not found — ASCII grapheme fallback; g2p_kind() reports ASCII so the TS layer feeds IPA via kokoro_synthesize_ipa (#11776)") diff --git a/tools/kokoro/src/kokoro-phonemes.cpp b/tools/kokoro/src/kokoro-phonemes.cpp index 1f9fe9619..1e577f157 100644 --- a/tools/kokoro/src/kokoro-phonemes.cpp +++ b/tools/kokoro/src/kokoro-phonemes.cpp @@ -24,7 +24,8 @@ #if defined(KOKORO_USE_ESPEAK) #include -#include +// one lock + init owner shared with tools/misaki-g2p (libespeak-ng is process-global) +#include "eliza-espeak-shared.h" #endif namespace eliza_kokoro { @@ -113,31 +114,21 @@ char32_t next_codepoint(const std::string & s, size_t & i) noexcept { #if defined(KOKORO_USE_ESPEAK) -// espeak-ng is process-global and not thread-safe; serialize init + calls. -std::mutex & espeak_mutex() { - static std::mutex m; - return m; -} - -bool ensure_espeak_init() { - static bool ok = [] { - // AUDIO_OUTPUT_SYNCHRONOUS so no audio device is opened. - const int rate = espeak_Initialize(AUDIO_OUTPUT_SYNCHRONOUS, 0, nullptr, 0); - if (rate < 0) { - return false; - } - return espeak_SetVoiceByName("en-us") == EE_OK; - }(); - return ok; -} +// espeak-ng is process-global and not thread-safe: every call goes through +// the shared lock, and the voice is re-selected under it on each use because +// misaki-g2p (same process, fused library) selects its own. Kokoro uses +// libespeak-ng's default data path, or the explicit one misaki was opened +// with when that happened first (see eliza-espeak-shared.h). +constexpr const char * kVoice = "en-us"; // Run espeak over text, accumulating IPA across all clauses. espeak returns // one clause per call (stopping at sentence/comma punctuation) and advances // the text pointer; we join clauses with a single space to match the // `espeak-ng -q --ipa` binary output the reference was derived from. std::string espeak_text_to_ipa(const std::string & text) { - std::lock_guard lock(espeak_mutex()); - if (!ensure_espeak_init()) { + auto lock = eliza_espeak_shared::lock(); + std::string err; + if (!eliza_espeak_shared::prepare(lock, std::string(), kVoice, err)) { return std::string(); } const void * inptr = static_cast(text.c_str()); @@ -231,8 +222,9 @@ std::vector phonemize_to_input_ids(const std::string & text) { bool espeak_available() noexcept { #if defined(KOKORO_USE_ESPEAK) - std::lock_guard lock(espeak_mutex()); - return ensure_espeak_init(); + auto lock = eliza_espeak_shared::lock(); + std::string err; + return eliza_espeak_shared::prepare(lock, std::string(), kVoice, err); #else return false; #endif diff --git a/tools/kokoro/tests/CMakeLists.txt b/tools/kokoro/tests/CMakeLists.txt index ce1fd13e4..d3a0bcbf8 100644 --- a/tools/kokoro/tests/CMakeLists.txt +++ b/tools/kokoro/tests/CMakeLists.txt @@ -50,3 +50,7 @@ add_test(NAME test-kokoro-layer-options COMMAND test-kokoro-layer-options) add_executable(test-kokoro-cancel test_kokoro_cancel.cpp) target_link_libraries(test-kokoro-cancel PRIVATE kokoro_lib) add_test(NAME test-kokoro-cancel COMMAND test-kokoro-cancel) + +# "main" is the label CI runs (ctest -L main); unlabelled tests are built but +# never executed there. Only the self-contained, weight-free, offline tests. +set_tests_properties(test-kokoro-cancel test-kokoro-layer-options test-kokoro-istft test-kokoro-layers-portable PROPERTIES LABELS main) diff --git a/tools/kokoro/tests/test_kokoro_layer_options.cpp b/tools/kokoro/tests/test_kokoro_layer_options.cpp index 7fdd8e520..533d0fef4 100644 --- a/tools/kokoro/tests/test_kokoro_layer_options.cpp +++ b/tools/kokoro/tests/test_kokoro_layer_options.cpp @@ -102,10 +102,19 @@ int main() { detail::KokoroLayerOptionsScope s(opts(true)); snake1d_forward(x1.data(), 4, 1000, a.data()); } + // Both paths round the final x + sin^2(a x) / a to float, so each + // can be up to half an ulp of the output away from the exact value. + // Subtract that one-ulp allowance, then scale by |a| so the result + // is in sin^2 units and comparable to the fast_sin2 bound. double e = 0; for (int c = 0; c < 4; ++c) - for (int t = 0; t < 1000; ++t) - e = std::max(e, std::fabs((double) x[c * 1000 + t] - x1[c * 1000 + t]) * std::fabs(a[c])); + for (int t = 0; t < 1000; ++t) { + const float r = x[c * 1000 + t]; + const double ulp = (double) std::nextafter(std::fabs(r), INFINITY) - std::fabs(r); + const double d = std::fabs((double) r - x1[c * 1000 + t]); + e = std::max(e, std::max(0.0, d - ulp) * std::fabs(a[c])); + } + std::printf("snake1d fast vs libm: max excess over output rounding %.3g\n", e); CHECK(e < 1e-6, "snake mismatch %.3g", e); } diff --git a/tools/misaki-g2p/CMakeLists.txt b/tools/misaki-g2p/CMakeLists.txt index 3e80f61cd..7b6db7b54 100644 --- a/tools/misaki-g2p/CMakeLists.txt +++ b/tools/misaki-g2p/CMakeLists.txt @@ -69,6 +69,9 @@ if (MISAKI_G2P_ESPEAK) endif() if (MISAKI_G2P_HAS_ESPEAK) target_compile_definitions(misaki_g2p PRIVATE MISAKI_G2P_USE_ESPEAK) + # One lock + init owner for the process-global libespeak-ng, shared + # with tools/kokoro (header-only; see tools/espeak-shared). + target_include_directories(misaki_g2p PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/../espeak-shared) message(STATUS "misaki-g2p: libespeak-ng linked (GPL-3.0-or-later): eSpeak + E2M fallback, as in Paradee training") else() message(STATUS "misaki-g2p: libespeak-ng not found: out-of-lexicon words are spelled out (fallback kind 0)") diff --git a/tools/misaki-g2p/include/misaki-g2p.h b/tools/misaki-g2p/include/misaki-g2p.h index 9a4dd6cd7..f136bd8a9 100644 --- a/tools/misaki-g2p/include/misaki-g2p.h +++ b/tools/misaki-g2p/include/misaki-g2p.h @@ -8,7 +8,8 @@ * * Thread safety: a misaki_g2p handle may be used from one thread at a time. * Separate handles may run concurrently; calls into eSpeak NG are serialised - * internally because libespeak-ng is process-global. + * internally because libespeak-ng is process-global, under the same lock + * Kokoro uses (tools/espeak-shared/eliza-espeak-shared.h). */ #ifndef MISAKI_G2P_H #define MISAKI_G2P_H diff --git a/tools/misaki-g2p/src/espeak.cpp b/tools/misaki-g2p/src/espeak.cpp index 8e62b175c..2db51fc7f 100644 --- a/tools/misaki-g2p/src/espeak.cpp +++ b/tools/misaki-g2p/src/espeak.cpp @@ -6,11 +6,12 @@ // then misaki espeak.EspeakFallback's eSpeak-to-misaki (E2M) table. #include "misaki-internal.h" -#include #include #if defined(MISAKI_G2P_USE_ESPEAK) # include +// one lock + init owner shared with tools/kokoro (libespeak-ng is process-global) +# include "eliza-espeak-shared.h" #endif namespace misaki { @@ -152,31 +153,8 @@ std::vector restore(std::vector text, std::vector marks) { } #if defined(MISAKI_G2P_USE_ESPEAK) -std::mutex & espeak_mutex() { - static std::mutex m; - return m; -} -bool g_espeak_ready = false; -std::string g_espeak_path; - -// Caller holds espeak_mutex(). (Re)initialises libespeak-ng for data_dir. -bool ensure_espeak_locked(const std::string & data_dir, std::string & err) { - if (g_espeak_ready && g_espeak_path == data_dir) return true; - g_espeak_ready = false; - int rate = espeak_Initialize(AUDIO_OUTPUT_SYNCHRONOUS, 0, data_dir.empty() ? nullptr : data_dir.c_str(), 0); - if (rate <= 0) { - err = "espeak_Initialize failed for data dir '" + data_dir + "'"; - return false; - } - // phonemizer selects the first voice whose language is en-us - if (espeak_SetVoiceByName("gmw/en-US") != EE_OK) { - err = "eSpeak NG voice gmw/en-US not found under '" + data_dir + "'"; - return false; - } - g_espeak_ready = true; - g_espeak_path = data_dir; - return true; -} +// phonemizer selects the first voice whose language is en-us +constexpr const char * kVoice = "gmw/en-US"; #endif } // namespace @@ -229,14 +207,14 @@ u32 EspeakFallback::e2m(const u32 & in) { std::unique_ptr EspeakFallback::open(const std::string & data_dir, std::string & err) { #if defined(MISAKI_G2P_USE_ESPEAK) - std::lock_guard lock(espeak_mutex()); + auto lock = eliza_espeak_shared::lock(); // espeak-ng keeps its data path in a 160-byte buffer (N_PATH_HOME) and // fails on longer ones with a misleading "No such file" error. if (data_dir.size() + sizeof("/espeak-ng-data/phontab") > 160) { err = "eSpeak NG data path is too long for libespeak-ng (max ~135 bytes): " + data_dir; return nullptr; } - if (!ensure_espeak_locked(data_dir, err)) return nullptr; + if (!eliza_espeak_shared::prepare(lock, data_dir, kVoice, err)) return nullptr; return std::unique_ptr(new EspeakFallback(data_dir)); #else (void) data_dir; @@ -249,9 +227,11 @@ EspeakFallback::~EspeakFallback() = default; u32 EspeakFallback::phonemize_line(const u32 & text) { #if defined(MISAKI_G2P_USE_ESPEAK) - std::lock_guard lock(espeak_mutex()); + // Re-applies this component's voice under the shared lock: Kokoro may + // have selected its own since the last call. + auto lock = eliza_espeak_shared::lock(); std::string init_err; - if (!ensure_espeak_locked(data_dir_, init_err)) throw std::runtime_error(init_err); + if (!eliza_espeak_shared::prepare(lock, data_dir_, kVoice, init_err)) throw std::runtime_error(init_err); std::string utf8 = to_utf8(text); const void * ptr = utf8.c_str(); // EspeakWrapper.text_to_phonemes(text, tie=True) @@ -290,6 +270,7 @@ u32 EspeakFallback::phonemize_line(const u32 & text) { return out; #else (void) text; + (void) data_dir_; return u32(); #endif } diff --git a/tools/misaki-g2p/src/g2p.cpp b/tools/misaki-g2p/src/g2p.cpp index e4743abc9..3cdbd57b8 100644 --- a/tools/misaki-g2p/src/g2p.cpp +++ b/tools/misaki-g2p/src/g2p.cpp @@ -425,12 +425,20 @@ class Lexicon { } } }; + // Numbers past the named scales are not read (as the guarded cardinal + // branch below): misaki's num2words raises there. + bool too_long = false; auto extend_num = [&](const u32 & num, bool first) { std::string d = to_utf8(num); + if (!number_fits(d)) { + too_long = true; + return; + } extend_words(number_cardinal(d), first); }; const bool digits = is_ascii_digits(word); if (digits && is_ordinal_suffix()) { + if (!number_fits(to_utf8(word))) return {}; extend_words(number_ordinal(to_utf8(word)), true); } else if (result.empty() && word.size() == 4 && !(currency && (*currency == U"$" || *currency == U"£" || *currency == U"€")) && digits) { extend_words(number_year(to_utf8(word)), true); @@ -514,12 +522,13 @@ class Lexicon { words = U"point"; for (size_t i = 1; i < w2.size(); i++) words += U" " + number_cardinal(to_utf8(u32(1, w2[i]))); } else { + if (!number_fits(to_utf8(w2.substr(0, w2.find('.'))))) return {}; words = number_float(to_utf8(w2)); } } extend_words(words, true); } - if (result.empty()) return {}; + if (too_long || result.empty()) return {}; u32 joined; int rating = 99; for (size_t i = 0; i < result.size(); i++) { diff --git a/tools/misaki-g2p/src/lexicon.cpp b/tools/misaki-g2p/src/lexicon.cpp index 3922336e4..459f2f45d 100644 --- a/tools/misaki-g2p/src/lexicon.cpp +++ b/tools/misaki-g2p/src/lexicon.cpp @@ -2,6 +2,8 @@ // Reader for misaki-en-us.lex (written by convert_misaki_lexicon.py). #include "misaki-internal.h" +#include +#include #include #include @@ -123,7 +125,7 @@ std::unique_ptr LexFile::open(const std::string & path, std::string & e LexFile::~LexFile() { #if !defined(MISAKI_NO_MMAP) - if (mapped_ && base_) munmap((void *) base_, size_); + if (mapped_ && base_) munmap(const_cast(base_), size_); #endif } diff --git a/tools/misaki-g2p/src/misaki-internal.h b/tools/misaki-g2p/src/misaki-internal.h index 7d9d7d3d1..8324c3da5 100644 --- a/tools/misaki-g2p/src/misaki-internal.h +++ b/tools/misaki-g2p/src/misaki-internal.h @@ -156,7 +156,7 @@ class EspeakFallback { u32 phonemize_line(const u32 & text); // libespeak-ng is process-global: another handle may have re-initialised // it with a different data directory since this one was opened. - [[maybe_unused]] std::string data_dir_; + std::string data_dir_; }; bool espeak_linked(); diff --git a/tools/misaki-g2p/src/numbers.cpp b/tools/misaki-g2p/src/numbers.cpp index 5b264a207..bc960eb5f 100644 --- a/tools/misaki-g2p/src/numbers.cpp +++ b/tools/misaki-g2p/src/numbers.cpp @@ -49,6 +49,16 @@ std::string cardinal(const std::string & digits) { long start = end - 3 < 0 ? 0 : end - 3; groups.push_back(std::stoi(d.substr((size_t) start, (size_t) (end - start)))); } + if (groups.size() > kMaxGroups) { + // Beyond the named scales (num2words raises here). Callers gate on + // number_fits(); this keeps cardinal() itself in bounds regardless. + std::string out; + for (char c : d) { + if (!out.empty()) out += ' '; + out += kOnes[c - '0']; + } + return out; + } std::string out; int last = groups[0]; bool higher = false; diff --git a/tools/misaki-g2p/src/tagger.cpp b/tools/misaki-g2p/src/tagger.cpp index 9de6d6766..3110b8dba 100644 --- a/tools/misaki-g2p/src/tagger.cpp +++ b/tools/misaki-g2p/src/tagger.cpp @@ -9,6 +9,7 @@ // so the rules aim at those decisions, not at full tagging accuracy. #include "misaki-internal.h" +#include #include #include diff --git a/tools/misaki-g2p/tests/CMakeLists.txt b/tools/misaki-g2p/tests/CMakeLists.txt index 3f2d82bbf..d14e5add5 100644 --- a/tools/misaki-g2p/tests/CMakeLists.txt +++ b/tools/misaki-g2p/tests/CMakeLists.txt @@ -39,6 +39,9 @@ set_tests_properties(misaki_g2p_mini_lexicon PROPERTIES FIXTURES_SETUP misaki_mi add_test(NAME misaki_g2p_unit COMMAND test-misaki-g2p ${MINI_LEX} ${MISAKI_TEST_DATA}) set_tests_properties(misaki_g2p_unit PROPERTIES FIXTURES_REQUIRED misaki_mini_lex) +# "main" is the label CI runs (ctest -L main): both are offline (committed +# mini lexicon data) and need neither eSpeak nor model weights. +set_tests_properties(misaki_g2p_mini_lexicon misaki_g2p_unit PROPERTIES LABELS main) if (MISAKI_G2P_LEXICON) set(FULL_LEX ${MISAKI_G2P_LEXICON}) @@ -59,6 +62,23 @@ if (NOT GATE_ESPEAK_DATA AND MISAKI_ESPEAK_LIBRARY) endif() endif() +# misaki and Kokoro phonemizing concurrently through the shared libespeak-ng +# owner (tools/espeak-shared). Needs kokoro_lib, which is configured before +# this directory in the llama.cpp tree; skips (77) without libespeak-ng. +if (TARGET kokoro_lib) + find_package(Threads REQUIRED) + add_executable(test-espeak-shared test_espeak_shared.cpp) + target_link_libraries(test-espeak-shared PRIVATE misaki_g2p kokoro_lib Threads::Threads) + if (GATE_ESPEAK_DATA) + set(_espeak_shared_args ${GATE_ESPEAK_DATA}) + else() + set(_espeak_shared_args "") + endif() + add_test(NAME misaki_g2p_espeak_shared COMMAND test-espeak-shared ${MINI_LEX} ${_espeak_shared_args}) + set_tests_properties(misaki_g2p_espeak_shared PROPERTIES + FIXTURES_REQUIRED misaki_mini_lex SKIP_RETURN_CODE 77 TIMEOUT 300) +endif() + foreach(gate agreement agent) if (gate STREQUAL "agreement") set(fixture ${CMAKE_CURRENT_SOURCE_DIR}/fixtures/reference.txt) diff --git a/tools/misaki-g2p/tests/test_agreement.cpp b/tools/misaki-g2p/tests/test_agreement.cpp index 8a930d977..cbf052b4a 100644 --- a/tools/misaki-g2p/tests/test_agreement.cpp +++ b/tools/misaki-g2p/tests/test_agreement.cpp @@ -12,6 +12,7 @@ #include "misaki-g2p.h" #include "misaki-internal.h" +#include #include #include #include diff --git a/tools/misaki-g2p/tests/test_espeak_shared.cpp b/tools/misaki-g2p/tests/test_espeak_shared.cpp new file mode 100644 index 000000000..7aa43b5d2 --- /dev/null +++ b/tools/misaki-g2p/tests/test_espeak_shared.cpp @@ -0,0 +1,82 @@ +// SPDX-License-Identifier: Apache-2.0 +// misaki-g2p and Kokoro both drive the process-global libespeak-ng. They share +// one lock and init owner (tools/espeak-shared/eliza-espeak-shared.h); this +// runs both phonemizers concurrently from two threads and checks every result +// against the single-threaded one. Meant to be run under TSan as well. +// +// usage: test-espeak-shared MINI_LEX [ESPEAK_DATA_DIR] +// exit 77 (skip) when either side was built without libespeak-ng. +#include "kokoro-phonemes.h" +#include "misaki-g2p.h" + +#include +#include +#include +#include +#include +#include + +static std::string misaki_ps(misaki_g2p * g, const std::string & text) { + char * out = nullptr; + char * err = nullptr; + std::string s; + if (misaki_g2p_phonemize(g, text.data(), text.size(), &out, &err) == 0 && out) s = out; + else s = std::string("ERROR: ") + (err ? err : "?"); + misaki_g2p_free(out); + misaki_g2p_free(err); + return s; +} + +int main(int argc, char ** argv) { + if (argc < 2) { + std::fprintf(stderr, "usage: %s MINI_LEX [ESPEAK_DATA_DIR]\n", argv[0]); + return 2; + } + if (!misaki_g2p_espeak_linked() || !eliza_kokoro::espeak_available()) { + std::printf("SKIP: misaki-g2p or Kokoro built without libespeak-ng\n"); + return 77; + } + const char * data_dir = argc > 2 ? argv[2] : ""; + char * err = nullptr; + misaki_g2p * g = misaki_g2p_open(argv[1], data_dir, &err); + if (!g) { + std::printf("SKIP: eSpeak NG data not usable: %s\n", err ? err : "?"); + misaki_g2p_free(err); + return 77; + } + if (misaki_g2p_fallback_kind(g) != MISAKI_G2P_FALLBACK_ESPEAK) { + std::fprintf(stderr, "FAIL: misaki opened without the eSpeak fallback\n"); + misaki_g2p_close(g); + return 1; + } + + // Out-of-lexicon words, so every misaki call goes through eSpeak. + const std::string misaki_text = "Zyx quorbled the flimbit, glarf!"; + const std::string kokoro_text = "Hello world, this is a concurrency check."; + const std::string misaki_want = misaki_ps(g, misaki_text); + const std::vector kokoro_want = eliza_kokoro::phonemize_ipa(kokoro_text); + if (misaki_want.empty() || misaki_want.rfind("ERROR", 0) == 0 || kokoro_want.empty()) { + std::fprintf(stderr, "FAIL: baseline: misaki '%s', kokoro %zu ids\n", misaki_want.c_str(), kokoro_want.size()); + misaki_g2p_close(g); + return 1; + } + + const int iters = 200; + std::atomic misaki_bad{0}, kokoro_bad{0}; + std::atomic go{false}; + std::thread tm([&] { + while (!go.load()) std::this_thread::yield(); + for (int i = 0; i < iters; i++) if (misaki_ps(g, misaki_text) != misaki_want) misaki_bad++; + }); + std::thread tk([&] { + while (!go.load()) std::this_thread::yield(); + for (int i = 0; i < iters; i++) if (eliza_kokoro::phonemize_ipa(kokoro_text) != kokoro_want) kokoro_bad++; + }); + go = true; + tm.join(); + tk.join(); + misaki_g2p_close(g); + + std::printf("misaki: %d/%d differ, kokoro: %d/%d differ\n", misaki_bad.load(), iters, kokoro_bad.load(), iters); + return misaki_bad || kokoro_bad ? 1 : 0; +} diff --git a/tools/misaki-g2p/tests/test_misaki_g2p.cpp b/tools/misaki-g2p/tests/test_misaki_g2p.cpp index f2124d867..7fc542a05 100644 --- a/tools/misaki-g2p/tests/test_misaki_g2p.cpp +++ b/tools/misaki-g2p/tests/test_misaki_g2p.cpp @@ -100,6 +100,10 @@ static void test_numbers() { check_eq(to_utf8(number_float("10.05")), "ten point zero five", "float 10.05"); check_eq(to_utf8(number_float("3.0")), "three", "float 3.0"); CHECK(number_fits(std::string(42, '9')) && !number_fits(std::string(43, '9'))); + // past the named scales cardinal() reads digit by digit instead of indexing out of bounds + check_eq(to_utf8(number_cardinal(std::string(43, '1'))).substr(0, 11), "one one one", "cardinal 43 digits"); + check_eq(to_utf8(number_ordinal("1" + std::string(59, '0'))).substr(0, 9), "one zero ", "ordinal 60 digits"); + check_eq(to_utf8(number_float(std::string(60, '9') + ".5")).substr(0, 9), "nine nine", "float 60 digits"); } static void test_tables(const std::string & data) { @@ -211,7 +215,10 @@ static void test_api(const std::string & lex) { { std::string a(20000, 'a'), dots(5000, '.'), hy; for (int i = 0; i < 3000; i++) hy += "a-"; - for (const std::string & s : {a, dots, hy, std::string(3000, '9'), std::string(1000, '@') + "x.com"}) { + const std::string d60(60, '9'); + for (const std::string & s : {a, dots, hy, std::string(3000, '9'), std::string(1000, '@') + "x.com", + "It is " + d60 + ".5 units.", "Pay $" + d60 + " now.", "The " + d60 + "th one.", + "It is 1.1" + std::string(59, '0') + ".5 units.", "Pay $1." + d60 + " now."}) { int rc = 0; phonemize(g, s, &rc); CHECK(rc == 0); diff --git a/tools/omnivoice/CMakeLists.txt b/tools/omnivoice/CMakeLists.txt index ec56da599..593464cda 100644 --- a/tools/omnivoice/CMakeLists.txt +++ b/tools/omnivoice/CMakeLists.txt @@ -247,7 +247,22 @@ function(eliza_inference_fold_paradee target) endif() endfunction() +# Parakeet-TDT ASR (tools/parakeet, FermionResearch/Phonon-2 five-value +# weights). Folded in when parakeet_lib exists (added before this directory); +# it then serves eliza_inference_asr_* for bundles with a parakeet-tdt GGUF +# under asr/ (see src/eliza-asr-parakeet.inc). PRIVATE like kokoro/paradee. +function(eliza_inference_fold_parakeet target) + if(TARGET parakeet_lib) + target_link_libraries(${target} PRIVATE parakeet_lib) + target_include_directories(${target} PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/../parakeet/include) + target_compile_definitions(${target} PRIVATE ELIZA_ENABLE_PARAKEET) + endif() +endfunction() + if(TARGET mtmd) + if(TARGET parakeet_lib) + message(STATUS "elizainference: Parakeet-TDT ASR enabled (streaming via asr_stream_*)") + endif() if(TARGET paradee_lib AND TARGET misaki_g2p) message(STATUS "elizainference: Paradee TTS + misaki G2P enabled (ABI v17)") elseif(TARGET paradee_lib OR TARGET misaki_g2p) @@ -305,6 +320,7 @@ if(TARGET mtmd) target_compile_definitions(elizainference PRIVATE ELIZA_ENABLE_KOKORO) endif() eliza_inference_fold_paradee(elizainference) + eliza_inference_fold_parakeet(elizainference) set_target_properties(elizainference PROPERTIES OUTPUT_NAME elizainference POSITION_INDEPENDENT_CODE ON @@ -360,6 +376,7 @@ if(TARGET mtmd) target_compile_definitions(elizainference_static PRIVATE ELIZA_ENABLE_KOKORO) endif() eliza_inference_fold_paradee(elizainference_static) + eliza_inference_fold_parakeet(elizainference_static) set_target_properties(elizainference_static PROPERTIES OUTPUT_NAME elizainference POSITION_INDEPENDENT_CODE ON) @@ -402,6 +419,20 @@ if(TARGET mtmd) endif() endif() + if(TARGET parakeet_lib) + # Parakeet-TDT ASR through the real library. Skips (passes) unless + # ELIZA_TEST_PARAKEET_BUNDLE / ELIZA_TEST_PARAKEET_PCM are set. + add_executable(eliza-inference-test-parakeet-asr tests/ffi-parakeet-asr.cpp) + target_compile_features(eliza-inference-test-parakeet-asr PRIVATE cxx_std_17) + target_include_directories(eliza-inference-test-parakeet-asr PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/include) + find_package(Threads REQUIRED) + target_link_libraries(eliza-inference-test-parakeet-asr PRIVATE elizainference Threads::Threads) + if(BUILD_TESTING) + add_test(NAME eliza-inference-test-parakeet-asr COMMAND eliza-inference-test-parakeet-asr) + set_property(TEST eliza-inference-test-parakeet-asr PROPERTY LABELS main) + endif() + endif() + add_executable(eliza-inference-test-paradee-asr tests/ffi-paradee-asr.cpp tests/paradee-contract/fake-paradee.cpp) diff --git a/tools/omnivoice/include/eliza-inference-ffi.h b/tools/omnivoice/include/eliza-inference-ffi.h index a6e9ce9c7..e1083ceed 100644 --- a/tools/omnivoice/include/eliza-inference-ffi.h +++ b/tools/omnivoice/include/eliza-inference-ffi.h @@ -144,6 +144,28 @@ extern "C" { * non-positive capacity instead of creating a zero-sized context. * Additive symbol; no existing signature changes. A v17 library lacks * the symbol, so loaders must treat it as absent. + * Parakeet-TDT ASR backend (additive behaviour, no new symbols, no + * version bump): + * builds that fold tools/parakeet (`ELIZA_ENABLE_PARAKEET`) serve the + * ASR surface from a parakeet-tdt GGUF under `/asr/` (e.g. + * FermionResearch/Phonon-2, CC-BY-4.0) when `mmap_acquire("asr")` + * finds one; `ELIZA_ASR_BACKEND=qwen3|parakeet` overrides. With it: + * - `asr_transcribe_timed` returns TDT token-frame word times (80 ms + * resolution) instead of duration-proportional ones; + * - `asr_stream_supported()` returns 1 while at least one context has + * the parakeet backend loaded (0 before `mmap_acquire("asr")`, after + * `mmap_evict("asr")`, and for the Qwen3-ASR backend), and the + * `asr_stream_*` session decodes pause-delimited windows with local + * agreement. `stream_open` fails cleanly without the backend; + * partial/finish on a stream whose region was evicted return + * ELIZA_ERR_INVALID_ARG; + * - stream `out_tokens` are never written (`*io_n_tokens = 0`): + * Parakeet's SentencePiece vocabulary is not the text model's; + * - `asr_cancel` is polled between encoder layers and decode steps; + * - `asr_stream_feed` may run on another thread than partial/finish + * (it only buffers, and never waits for a decode); `close` must not + * race other calls on the same stream. Partial/finish truncate at a + * UTF-8 boundary when `max_text_bytes` is too small. * v17: Paradee-8M TTS with native misaki G2P, plus cancellation and * eviction for every TTS head and for ASR. All additions are new * symbols; no existing signature changes. @@ -293,7 +315,7 @@ const char * eliza_inference_abi_version(void); * ---------------------- -------------------------------------- ----- * batch ASR + cancel eliza_inference_asr_supported v17 * ASR word timestamps eliza_inference_asr_timestamps_supported v12 - * streaming ASR eliza_inference_asr_stream_supported v2 (0 today) + * streaming ASR eliza_inference_asr_stream_supported v2 (1 with parakeet loaded, v18) * OmniVoice stream+cancel eliza_inference_tts_stream_supported v2 * Kokoro TTS eliza_inference_kokoro_supported v10 * Kokoro G2P kind eliza_inference_kokoro_g2p_kind(ctx) v14 diff --git a/tools/omnivoice/src/eliza-asr-parakeet.inc b/tools/omnivoice/src/eliza-asr-parakeet.inc new file mode 100644 index 000000000..a56a7d08a --- /dev/null +++ b/tools/omnivoice/src/eliza-asr-parakeet.inc @@ -0,0 +1,201 @@ +/* tools/omnivoice/src/eliza-asr-parakeet.inc - Parakeet-TDT (Phonon-2) + * backend for the eliza_inference_asr_* surface. + * + * #included by eliza-inference-ffi.cpp inside `#ifdef ELIZA_ENABLE_PARAKEET` + * after EliInferenceContext is defined. Everything here runs on + * tools/parakeet (parakeet_lib): native FastConformer encoder + TDT decoder + * on the CPU, five-value int8 kernels, no ONNX / MLX / Core ML. + * + * Selection (mmap_acquire("asr")): + * ELIZA_ASR_BACKEND=parakeet require a parakeet-tdt GGUF under /asr/ + * ELIZA_ASR_BACKEND=qwen3 never use parakeet (Qwen3-ASR via mtmd) + * unset / auto parakeet when /asr/ holds a GGUF with + * general.architecture = "parakeet-tdt", + * otherwise the Qwen3-ASR path, unchanged. + * Threads: ELIZA_ASR_THREADS, else parakeet's default (fast cores, <= 4). + * + * Streaming: windowed re-decode with local agreement (tools/parakeet/src/ + * pk-stream.cpp). eliza_inference_asr_stream_supported() returns 1 while at + * least one context has the parakeet backend loaded. Token ids are not + * returned (*io_n_tokens = 0): Parakeet's 8192-piece SentencePiece vocabulary + * is not the text model's. + */ + +#include "parakeet.h" + +#include + +static std::atomic g_eliza_parakeet_loaded{ 0 }; + +struct EliParakeetAsr { + pk_model * model = nullptr; + std::string path; + uint64_t generation = 0; /* bumped on every load, so open streams notice an evict */ +}; + +/* Shared by every context (each bumps it under its own asr_mutex). */ +static std::atomic g_eliza_parakeet_generation{ 0 }; + +static int eliza_parakeet_cancel_cb(void * user) { + auto * ctx = (EliInferenceContext *) user; + return ctx && ctx->asr_cancel.load(std::memory_order_acquire) ? 1 : 0; +} + +static std::string eliza_parakeet_backend_pref() { + const char * env = std::getenv("ELIZA_ASR_BACKEND"); + std::string v = env ? env : ""; + std::transform(v.begin(), v.end(), v.begin(), [](unsigned char c) { return (char) std::tolower(c); }); + return v; +} + +static std::string eliza_parakeet_find(const std::filesystem::path & bundle_dir) { + const auto dir = bundle_dir / "asr"; + std::error_code ec; + if (!std::filesystem::is_directory(dir, ec)) return ""; + std::vector found; + for (const auto & e : std::filesystem::directory_iterator(dir, ec)) { + if (!e.is_regular_file(ec) || e.path().extension() != ".gguf") continue; + if (pk_gguf_is_parakeet(e.path().string().c_str())) found.push_back(e.path().string()); + } + std::sort(found.begin(), found.end()); + return found.empty() ? "" : found.front(); +} + +/* Caller holds asr_mutex. Returns 1 when the parakeet backend is (now) + * loaded, 0 when this bundle / preference does not use it, negative on a + * load error. */ +static int eliza_parakeet_load_locked(EliInferenceContext * ctx, char ** out_error) { + if (ctx->parakeet && ctx->parakeet->model) return 1; + const std::string pref = eliza_parakeet_backend_pref(); + if (pref == "qwen3" || pref == "mtmd") return 0; + const std::string path = eliza_parakeet_find(std::filesystem::path(ctx->bundle_dir)); + if (path.empty()) { + if (pref == "parakeet") { + eliza_set_error(out_error, std::string("[libelizainference] ELIZA_ASR_BACKEND=parakeet but no parakeet-tdt GGUF under ") + + (std::filesystem::path(ctx->bundle_dir) / "asr").string()); + return ELIZA_ERR_INVALID_ARG; + } + return 0; + } + pk_params pp = pk_default_params(); + /* same policy as the Qwen3-ASR path: Android's CPU profile keeps ASR + * weights in private memory (ELIZA_ASR_USE_MMAP overrides) */ + pp.use_mmap = eliza_asr_use_mmap() ? 1 : 0; + if (const char * env = std::getenv("ELIZA_ASR_THREADS")) { + const int n = std::atoi(env); + if (n > 0) pp.n_threads = n; + } + char err[512] = { 0 }; + eliza_asr_debug_log("loading parakeet-tdt ASR"); + pk_model * m = pk_model_load(path.c_str(), pp, err, sizeof(err)); + if (!m) { + eliza_set_error(out_error, std::string("[libelizainference] parakeet ASR load failed: ") + err); + return ELIZA_ERR_FFI_FAULT; + } + if (!ctx->parakeet) ctx->parakeet = new EliParakeetAsr(); + ctx->parakeet->model = m; + ctx->parakeet->path = path; + ctx->parakeet->generation = g_eliza_parakeet_generation.fetch_add(1) + 1; + g_eliza_parakeet_loaded.fetch_add(1); + eliza_asr_debug_log("parakeet-tdt ASR ready"); + return 1; +} + +/* Caller holds asr_mutex. */ +static void eliza_parakeet_free_locked(EliInferenceContext * ctx) { + if (!ctx->parakeet) return; + if (ctx->parakeet->model) { + pk_model_free(ctx->parakeet->model); + g_eliza_parakeet_loaded.fetch_sub(1); + } + delete ctx->parakeet; + ctx->parakeet = nullptr; +} + +static bool eliza_parakeet_active_locked(const EliInferenceContext * ctx) { + return ctx->parakeet && ctx->parakeet->model; +} + +/* Caller holds asr_mutex and has cleared asr_cancel. */ +static int eliza_parakeet_transcribe_locked( + EliInferenceContext * ctx, + const float * pcm, + size_t n_samples, + int sample_rate_hz, + std::string & out_text, + std::vector * word_start, + std::vector * word_end, + char ** out_error) { + pk_model * m = ctx->parakeet->model; + std::vector audio = eliza_resample_linear(pcm, n_samples, sample_rate_hz, pk_model_sample_rate(m)); + pk_result r{}; + const int rc = pk_transcribe(m, audio.data(), audio.size(), eliza_parakeet_cancel_cb, ctx, &r); + if (rc == PK_E_CANCELLED) { + pk_result_free(&r); + eliza_set_error(out_error, "[libelizainference] asr_transcribe: cancelled"); + return ELIZA_ERR_CANCELLED; + } + if (rc != PK_OK) { + pk_result_free(&r); + eliza_set_error(out_error, "[libelizainference] asr_transcribe: parakeet decode failed rc=" + std::to_string(rc)); + return ELIZA_ERR_FFI_FAULT; + } + out_text = r.text ? r.text : ""; + if (word_start && word_end) { + size_t n = r.n_tokens + 1; + word_start->assign(n, 0); + word_end->assign(n, 0); + pk_result_words(m, &r, word_start->data(), word_end->data(), &n); + word_start->resize(n); + word_end->resize(n); + } + pk_result_free(&r); + return 0; +} + +/* ---- streaming ---------------------------------------------------------- */ + +/* Threading: feed() may run on the audio thread while partial()/finish() + * decode on another. feed() only resamples into `pending` under `feed_mu` + * (never waits for a decode); partial()/finish() hold ctx->asr_mutex, move + * `pending` into the pk_stream under `feed_mu`, then decode without it. + * close() must not race any other call on the same stream. */ +struct EliAsrStream { + EliInferenceContext * ctx = nullptr; + pk_stream * st = nullptr; + uint64_t generation = 0; + int in_rate = 16000; + std::mutex feed_mu; + std::vector pending; /* 16 kHz, not yet handed to st */ + bool finished = false; /* under feed_mu */ + /* continuous linear resampler state (in_rate -> 16 kHz), under feed_mu */ + double phase = 0.0; + float last = 0.0f; + bool have_last = false; +}; + +/* Appends to `out`. Caller holds s->feed_mu. */ +static void eliza_parakeet_stream_resample(EliAsrStream * s, const float * pcm, size_t n, std::vector & out) { + if (s->in_rate == 16000) { + out.insert(out.end(), pcm, pcm + n); + return; + } + const double step = (double) s->in_rate / 16000.0; + /* input index space: -1 = s->last, 0..n-1 = pcm */ + double pos = s->phase; + for (;;) { + const double ip = std::floor(pos); + const long i0 = (long) ip; + if (i0 + 1 >= (long) n) break; + const float a = i0 < 0 ? (s->have_last ? s->last : pcm[0]) : pcm[i0]; + const float b = pcm[i0 + 1]; + const float t = (float) (pos - ip); + out.push_back(a + (b - a) * t); + pos += step; + } + if (n > 0) { + s->last = pcm[n - 1]; + s->have_last = true; + s->phase = pos - (double) n; /* relative to the next chunk; >= -1 */ + } +} diff --git a/tools/omnivoice/src/eliza-inference-ffi.cpp b/tools/omnivoice/src/eliza-inference-ffi.cpp index bd44a856d..ab42ad2a0 100644 --- a/tools/omnivoice/src/eliza-inference-ffi.cpp +++ b/tools/omnivoice/src/eliza-inference-ffi.cpp @@ -128,6 +128,10 @@ struct EliVoicePreset { bool empty_payload = true; }; +#ifdef ELIZA_ENABLE_PARAKEET +struct EliParakeetAsr; +#endif + struct EliInferenceContext { std::string bundle_dir; std::string tts_model_path; @@ -155,6 +159,11 @@ struct EliInferenceContext { std::atomic asr_cancel{false}; std::mutex tts_mutex; std::mutex asr_mutex; +#ifdef ELIZA_ENABLE_PARAKEET + /* Parakeet-TDT ASR backend (eliza-asr-parakeet.inc), owned under + * asr_mutex. When loaded it serves every asr_* call on this ctx. */ + EliParakeetAsr * parakeet = nullptr; +#endif /* Streaming-LLM text model, loaded lazily on the first * eliza_inference_llm_stream_open and shared (read-only weights) across * sessions. Each session owns its own llama_context (private KV) over @@ -923,8 +932,15 @@ static bool eliza_tts_stream_chunk(const float * samples, int n_samples, void * return !(state->ctx && state->ctx->tts_cancel.load(std::memory_order_acquire)); } +#ifdef ELIZA_ENABLE_PARAKEET +#include "eliza-asr-parakeet.inc" +#endif + static void eliza_free_asr(EliInferenceContext * ctx) { if (!ctx) return; +#ifdef ELIZA_ENABLE_PARAKEET + eliza_parakeet_free_locked(ctx); +#endif if (ctx->asr_sampler) { llama_sampler_free(ctx->asr_sampler); ctx->asr_sampler = nullptr; @@ -1000,6 +1016,13 @@ static int eliza_load_asr(EliInferenceContext * ctx, char ** out_error) { return ELIZA_ERR_INVALID_ARG; } std::lock_guard lock(ctx->asr_mutex); +#ifdef ELIZA_ENABLE_PARAKEET + { + const int pk = eliza_parakeet_load_locked(ctx, out_error); + if (pk < 0) return pk; + if (pk == 1) return ELIZA_OK; + } +#endif if (ctx->asr_model && ctx->asr_lctx && ctx->asr_mtmd && ctx->asr_sampler) { return ELIZA_OK; } @@ -2893,6 +2916,12 @@ static int eliza_asr_decode_core( /* Cleared only once this call owns the ASR region: a cancel issued while * this call waited for the mutex was aimed at the call that held it. */ ctx->asr_cancel.store(false, std::memory_order_release); +#ifdef ELIZA_ENABLE_PARAKEET + if (eliza_parakeet_active_locked(ctx)) { + (void) max_text_bytes; + return eliza_parakeet_transcribe_locked(ctx, pcm, n_samples, sample_rate_hz, out_transcript, nullptr, nullptr, out_error); + } +#endif if (!ctx->asr_model || !ctx->asr_lctx || !ctx->asr_mtmd || !ctx->asr_sampler) { eliza_set_error(out_error, "[libelizainference] asr_transcribe: ASR region is not acquired; call mmap_acquire(\"asr\") after arming voice input"); return ELIZA_ERR_INVALID_ARG; @@ -3036,6 +3065,36 @@ int eliza_inference_asr_transcribe_timed( const size_t word_cap = *io_n_words; *io_n_words = 0; +#ifdef ELIZA_ENABLE_PARAKEET + { + /* Parakeet: real TDT token frames (80 ms) give the word times. */ + std::unique_lock lock(ctx->asr_mutex); + if (eliza_parakeet_active_locked(ctx)) { + ctx->asr_cancel.store(false, std::memory_order_release); + std::string text; + std::vector ws, we; + if (n_samples > 0) { + const int prc = eliza_parakeet_transcribe_locked(ctx, pcm, n_samples, sample_rate_hz, text, &ws, &we, out_error); + if (prc < 0) return prc; + } + lock.unlock(); + if (text.size() + 1 > max_text_bytes) { + eliza_set_error(out_error, "[libelizainference] asr_transcribe_timed: output buffer too small"); + return ELIZA_ERR_INVALID_ARG; + } + std::memcpy(out_text, text.data(), text.size()); + out_text[text.size()] = '\0'; + const size_t n = std::min(word_cap, ws.size()); + for (size_t i = 0; i < n; ++i) { + out_word_start_ms[i] = ws[i]; + out_word_end_ms[i] = we[i]; + } + *io_n_words = n; + return (int) text.size(); + } + } +#endif + std::string transcript; int rc = eliza_asr_decode_core(ctx, pcm, n_samples, sample_rate_hz, max_text_bytes, transcript, out_error); if (rc < 0) { @@ -3098,6 +3157,137 @@ int eliza_inference_asr_transcribe_timed( * version check (ffi-bindings.ts expects v3) succeeds. */ +#ifdef ELIZA_ENABLE_PARAKEET +/* Parakeet backend (eliza-asr-parakeet.inc): 1 while some context has it + * loaded (mmap_acquire("asr") on a bundle with a parakeet-tdt GGUF). Before + * that, and for the Qwen3-ASR backend, it stays 0 so callers keep the batch + * adapter. */ +int eliza_inference_asr_stream_supported(void) { + return g_eliza_parakeet_loaded.load() > 0 ? 1 : 0; +} + +EliAsrStream * eliza_inference_asr_stream_open( + EliInferenceContext * ctx, + int sample_rate_hz, + char ** out_error) { + if (!ctx || sample_rate_hz <= 0) { + eliza_set_error(out_error, "[libelizainference] asr_stream_open: invalid arguments"); + return nullptr; + } + std::lock_guard lock(ctx->asr_mutex); + if (!eliza_parakeet_active_locked(ctx)) { + eliza_set_error(out_error, + "[libelizainference] asr_stream_open: streaming needs the parakeet-tdt ASR backend; " + "call mmap_acquire(\"asr\") on a bundle with asr/.gguf, or use the batch path"); + return nullptr; + } + auto * s = new (std::nothrow) EliAsrStream(); + if (!s) { + eliza_set_error(out_error, "[libelizainference] asr_stream_open: out of memory"); + return nullptr; + } + s->ctx = ctx; + s->in_rate = sample_rate_hz; + s->generation = ctx->parakeet->generation; + pk_stream_params sp = pk_stream_default_params(); + sp.right_context_ms = eliza_int_env_or_default("ELIZA_ASR_STREAM_RIGHT_MS", sp.right_context_ms); + sp.min_hop_ms = eliza_int_env_or_default("ELIZA_ASR_STREAM_HOP_MS", sp.min_hop_ms); + sp.max_window_ms = eliza_int_env_or_default("ELIZA_ASR_STREAM_MAX_WINDOW_MS", sp.max_window_ms); + s->st = pk_stream_open(ctx->parakeet->model, sp); + if (!s->st) { + delete s; + eliza_set_error(out_error, "[libelizainference] asr_stream_open: parakeet stream open failed"); + return nullptr; + } + return s; +} + +int eliza_inference_asr_stream_feed( + EliAsrStream * stream, + const float * pcm, + size_t n_samples, + char ** out_error) { + if (!stream || (!pcm && n_samples)) { + eliza_set_error(out_error, "[libelizainference] asr_stream_feed: invalid arguments"); + return ELIZA_ERR_INVALID_ARG; + } + if (n_samples > (size_t) INT_MAX) { + eliza_set_error(out_error, "[libelizainference] asr_stream_feed: chunk too large"); + return ELIZA_ERR_INVALID_ARG; + } + std::lock_guard lock(stream->feed_mu); + if (stream->finished) { + eliza_set_error(out_error, "[libelizainference] asr_stream_feed: stream already finished"); + return ELIZA_ERR_INVALID_ARG; + } + eliza_parakeet_stream_resample(stream, pcm, n_samples, stream->pending); + return (int) n_samples; +} + +static int eliza_parakeet_stream_text(EliAsrStream * stream, bool finish, char * out_text, size_t max_text_bytes, + size_t * io_n_tokens, char ** out_error) { + if (io_n_tokens) *io_n_tokens = 0; /* Parakeet's vocabulary is not the text model's */ + if (!stream || !out_text || max_text_bytes == 0) { + eliza_set_error(out_error, "[libelizainference] asr_stream: invalid arguments"); + return ELIZA_ERR_INVALID_ARG; + } + EliInferenceContext * ctx = stream->ctx; + std::lock_guard lock(ctx->asr_mutex); + if (!eliza_parakeet_active_locked(ctx) || ctx->parakeet->generation != stream->generation) { + eliza_set_error(out_error, "[libelizainference] asr_stream: the ASR region was evicted; open a new stream"); + return ELIZA_ERR_INVALID_ARG; + } + { + std::lock_guard flock(stream->feed_mu); + if (!stream->pending.empty()) { + pk_stream_feed(stream->st, stream->pending.data(), stream->pending.size()); + stream->pending.clear(); + } + if (finish) stream->finished = true; /* later feeds fail */ + } + ctx->asr_cancel.store(false, std::memory_order_release); + const int rc = finish + ? pk_stream_finish(stream->st, out_text, max_text_bytes, eliza_parakeet_cancel_cb, ctx) + : pk_stream_partial(stream->st, out_text, max_text_bytes, eliza_parakeet_cancel_cb, ctx); + if (rc == PK_E_CANCELLED) { + eliza_set_error(out_error, "[libelizainference] asr_stream: cancelled"); + return ELIZA_ERR_CANCELLED; + } + if (rc < 0) { + eliza_set_error(out_error, "[libelizainference] asr_stream: parakeet decode failed rc=" + std::to_string(rc)); + return ELIZA_ERR_FFI_FAULT; + } + return rc; +} + +int eliza_inference_asr_stream_partial( + EliAsrStream * stream, + char * out_text, + size_t max_text_bytes, + int * out_tokens, + size_t * io_n_tokens, + char ** out_error) { + (void) out_tokens; + return eliza_parakeet_stream_text(stream, false, out_text, max_text_bytes, io_n_tokens, out_error); +} + +int eliza_inference_asr_stream_finish( + EliAsrStream * stream, + char * out_text, + size_t max_text_bytes, + int * out_tokens, + size_t * io_n_tokens, + char ** out_error) { + (void) out_tokens; + return eliza_parakeet_stream_text(stream, true, out_text, max_text_bytes, io_n_tokens, out_error); +} + +void eliza_inference_asr_stream_close(EliAsrStream * stream) { + if (!stream) return; + pk_stream_close(stream->st); + delete stream; +} +#else int eliza_inference_asr_stream_supported(void) { return 0; } @@ -3161,6 +3351,7 @@ int eliza_inference_asr_stream_finish( void eliza_inference_asr_stream_close(EliAsrStream * stream) { (void) stream; } +#endif /* ---- Streaming TTS + DFlash verifier callback (ABI v2) ------------- * * diff --git a/tools/omnivoice/tests/ffi-parakeet-asr.cpp b/tools/omnivoice/tests/ffi-parakeet-asr.cpp new file mode 100644 index 000000000..edf93174f --- /dev/null +++ b/tools/omnivoice/tests/ffi-parakeet-asr.cpp @@ -0,0 +1,216 @@ +// tools/omnivoice/tests/ffi-parakeet-asr.cpp - the eliza_inference_asr_* +// surface on the Parakeet-TDT backend, through the real libelizainference. +// +// Needs a bundle with asr/.gguf and a 16 kHz fp32 mono clip: +// ELIZA_TEST_PARAKEET_BUNDLE= ELIZA_TEST_PARAKEET_PCM= +// [ELIZA_TEST_PARAKEET_TEXT=] +// Without them the test prints "skipped" and passes. +// +// Checks: stream probe is 0 before the region is acquired and 1 after; +// batch and timed transcripts agree (and match the expected text when +// given); word count equals the whitespace split and times are monotonic; +// a 48 kHz stream fed in 80 ms chunks yields partials and a final +// transcript; token ids are not returned; asr_cancel from another thread +// cancels a batch call; evicting the region turns the probe off and fails +// an open stream cleanly; re-acquiring restores it. +#include "eliza-inference-ffi.h" + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +static int fails = 0; +#define CHECK(c, ...) do { if (!(c)) { std::fprintf(stderr, "FAIL %s:%d: ", __FILE__, __LINE__); std::fprintf(stderr, __VA_ARGS__); std::fprintf(stderr, "\n"); ++fails; } } while (0) + +static std::vector read_f32(const char * path) { + std::ifstream f(path, std::ios::binary | std::ios::ate); + std::vector v; + if (!f) return v; + const size_t n = (size_t) f.tellg(); + f.seekg(0); + v.resize(n / 4); + f.read((char *) v.data(), (std::streamsize) n); + return v; +} + +static size_t count_words(const std::string & s) { + std::istringstream is(s); + std::string w; + size_t n = 0; + while (is >> w) ++n; + return n; +} + +int main() { + const char * bundle = std::getenv("ELIZA_TEST_PARAKEET_BUNDLE"); + const char * pcm_path = std::getenv("ELIZA_TEST_PARAKEET_PCM"); + const char * expect = std::getenv("ELIZA_TEST_PARAKEET_TEXT"); + if (!bundle || !pcm_path) { + std::printf("skipped: set ELIZA_TEST_PARAKEET_BUNDLE and ELIZA_TEST_PARAKEET_PCM\n"); + return 0; + } + std::vector pcm = read_f32(pcm_path); + CHECK(pcm.size() > 16000, "clip %s too short or unreadable", pcm_path); + char * err = nullptr; + EliInferenceContext * ctx = eliza_inference_create(bundle, &err); + CHECK(ctx, "create: %s", err ? err : "?"); + if (!ctx) return 1; + + CHECK(eliza_inference_asr_stream_supported() == 0, "stream probe must be 0 before acquire"); + int rc = eliza_inference_mmap_acquire(ctx, "asr", &err); + CHECK(rc == ELIZA_OK, "acquire asr: %d %s", rc, err ? err : ""); + CHECK(eliza_inference_asr_stream_supported() == 1, "stream probe must be 1 with parakeet loaded"); + + // batch + char text[4096]; + auto t0 = std::chrono::steady_clock::now(); + rc = eliza_inference_asr_transcribe(ctx, pcm.data(), pcm.size(), 16000, text, sizeof(text), &err); + const double ms = std::chrono::duration(std::chrono::steady_clock::now() - t0).count(); + CHECK(rc > 0, "transcribe rc %d %s", rc, err ? err : ""); + // lengths only: transcripts stay out of logs, CI logs included + std::printf("batch (%.0f ms for %.2f s): %zu words\n", ms, pcm.size() / 16000.0, count_words(text)); + if (expect) CHECK(std::strcmp(text, expect) == 0, "transcript differs from the expected text"); + + // timed + char text2[4096]; + int ws[512], we[512]; + size_t nw = 512; + rc = eliza_inference_asr_transcribe_timed(ctx, pcm.data(), pcm.size(), 16000, text2, sizeof(text2), ws, we, &nw, &err); + CHECK(rc > 0 && std::strcmp(text, text2) == 0, "timed transcript differs"); + CHECK(nw == count_words(text2), "word count %zu != whitespace split %zu", nw, count_words(text2)); + for (size_t i = 0; i < nw; ++i) { + CHECK(ws[i] < we[i], "word %zu: start %d >= end %d", i, ws[i], we[i]); + if (i) CHECK(ws[i] >= ws[i - 1], "word %zu starts before word %zu", i, i - 1); + CHECK(we[i] <= (int) (pcm.size() / 16) + 80, "word %zu ends past the audio", i); + } + + // stream at 48 kHz in 80 ms chunks + std::vector up(pcm.size() * 3); + for (size_t i = 0; i < up.size(); ++i) { + const double s = i / 3.0; + const size_t a = (size_t) s; + const size_t b = std::min(a + 1, pcm.size() - 1); + up[i] = (float) (pcm[a] + (pcm[b] - pcm[a]) * (s - a)); + } + EliAsrStream * st = eliza_inference_asr_stream_open(ctx, 48000, &err); + CHECK(st, "stream_open: %s", err ? err : ""); + int partials = 0; + std::string last; + for (size_t p = 0; st && p < up.size(); p += 3840) { + const size_t n = std::min(3840, up.size() - p); + CHECK(eliza_inference_asr_stream_feed(st, up.data() + p, n, &err) == (int) n, "feed"); + char part[4096]; + int toks[16]; + size_t nt = 16; + rc = eliza_inference_asr_stream_partial(st, part, sizeof(part), toks, &nt, &err); + CHECK(rc >= 0, "partial rc %d", rc); + CHECK(nt == 0, "token ids must not be returned"); + if (rc > 0 && last != part) { + last = part; + ++partials; + } + } + if (st) { + char fin[4096]; + rc = eliza_inference_asr_stream_finish(st, fin, sizeof(fin), nullptr, nullptr, &err); + CHECK(rc > 0, "finish rc %d", rc); + std::printf("stream (%d partial updates): %zu words\n", partials, count_words(fin)); + CHECK(partials >= 1, "no partial transcript before finish"); + CHECK(eliza_inference_asr_stream_feed(st, pcm.data(), 160, &err) < 0, "feed after finish must fail"); + eliza_inference_free_string(err); + err = nullptr; + eliza_inference_asr_stream_close(st); + } + + // feed from an audio thread while another thread polls partials (the + // JNI shape); run under TSan to check the stream's locking + { + EliAsrStream * cs = eliza_inference_asr_stream_open(ctx, 48000, &err); + CHECK(cs, "stream_open (concurrent): %s", err ? err : ""); + if (cs) { + std::atomic fed{ false }; + std::thread feeder([&] { + char * e = nullptr; + for (size_t p = 0; p < up.size(); p += 960) { + const size_t n = std::min(960, up.size() - p); + if (eliza_inference_asr_stream_feed(cs, up.data() + p, n, &e) != (int) n) { + std::fprintf(stderr, "FAIL concurrent feed\n"); + ++fails; + break; + } + std::this_thread::sleep_for(std::chrono::microseconds(200)); + } + eliza_inference_free_string(e); + fed = true; + }); + int polls = 0; + while (!fed.load()) { + char part[4096]; + char * e = nullptr; + rc = eliza_inference_asr_stream_partial(cs, part, sizeof(part), nullptr, nullptr, &e); + eliza_inference_free_string(e); + CHECK(rc >= 0, "concurrent partial rc %d", rc); + ++polls; + } + feeder.join(); + char fin[4096]; + rc = eliza_inference_asr_stream_finish(cs, fin, sizeof(fin), nullptr, nullptr, &err); + CHECK(rc > 0, "concurrent finish rc %d", rc); + CHECK(count_words(fin) + 3 >= count_words(text) && count_words(fin) <= count_words(text) + 3, + "concurrent stream word count %zu far from batch %zu", count_words(fin), count_words(text)); + std::printf("concurrent feed/partial: %d polls, %zu words\n", polls, count_words(fin)); + // a tiny buffer: truncated, NUL-terminated, length < capacity + char tiny[5]; + rc = eliza_inference_asr_stream_partial(cs, tiny, sizeof(tiny), nullptr, nullptr, &err); + CHECK(rc >= 0 && rc < (int) sizeof(tiny) && tiny[rc] == '\0', "tiny buffer rc %d", rc); + eliza_inference_asr_stream_close(cs); + } + } + + // cancel a batch call from another thread + { + std::vector longpcm; + for (int r = 0; r < 4; ++r) longpcm.insert(longpcm.end(), pcm.begin(), pcm.end()); + std::atomic crc{ 1 }; + std::thread th([&] { + char t[8192]; + char * e = nullptr; + crc = eliza_inference_asr_transcribe(ctx, longpcm.data(), longpcm.size(), 16000, t, sizeof(t), &e); + eliza_inference_free_string(e); + }); + std::this_thread::sleep_for(std::chrono::milliseconds(30)); + CHECK(eliza_inference_asr_cancel(ctx, &err) == ELIZA_OK, "cancel"); + th.join(); + CHECK(crc == ELIZA_ERR_CANCELLED || crc > 0, "cancelled transcribe rc %d", (int) crc); + std::printf("cancel: rc %d (%s)\n", (int) crc, crc == ELIZA_ERR_CANCELLED ? "cancelled" : "finished before the cancel landed"); + rc = eliza_inference_asr_transcribe(ctx, pcm.data(), pcm.size(), 16000, text2, sizeof(text2), &err); + CHECK(rc > 0 && std::strcmp(text, text2) == 0, "transcribe after cancel differs"); + } + + // evict / re-acquire + st = eliza_inference_asr_stream_open(ctx, 16000, &err); + CHECK(st, "stream_open before evict"); + CHECK(eliza_inference_mmap_evict(ctx, "asr", &err) == ELIZA_OK, "evict"); + CHECK(eliza_inference_asr_stream_supported() == 0, "probe must be 0 after evict"); + if (st) { + eliza_inference_asr_stream_feed(st, pcm.data(), 16000, &err); + char part[256]; + rc = eliza_inference_asr_stream_partial(st, part, sizeof(part), nullptr, nullptr, &err); + CHECK(rc < 0, "partial on an evicted region must fail, got %d", rc); + eliza_inference_asr_stream_close(st); + } + CHECK(eliza_inference_mmap_acquire(ctx, "asr", &err) == ELIZA_OK, "re-acquire"); + rc = eliza_inference_asr_transcribe(ctx, pcm.data(), pcm.size(), 16000, text2, sizeof(text2), &err); + CHECK(rc > 0 && std::strcmp(text, text2) == 0, "transcribe after re-acquire differs"); + + eliza_inference_destroy(ctx); + std::printf(fails ? "FAILED (%d)\n" : "parakeet FFI checks passed\n", fails); + return fails ? 1 : 0; +} diff --git a/tools/paradee/tests/CMakeLists.txt b/tools/paradee/tests/CMakeLists.txt index 3017d41df..3d83e853a 100644 --- a/tools/paradee/tests/CMakeLists.txt +++ b/tools/paradee/tests/CMakeLists.txt @@ -44,3 +44,7 @@ foreach(t test-paradee-api test-paradee-parity) endif() endforeach() set_tests_properties(test-paradee-parity PROPERTIES TIMEOUT 900) + +# "main" is the label CI runs (ctest -L main); unlabelled tests are built but +# never executed there. Only the self-contained, weight-free, offline tests. +set_tests_properties(test-paradee-phase-lock PROPERTIES LABELS main) diff --git a/tools/parakeet/CMakeLists.txt b/tools/parakeet/CMakeLists.txt new file mode 100644 index 000000000..d7a049fac --- /dev/null +++ b/tools/parakeet/CMakeLists.txt @@ -0,0 +1,100 @@ +# tools/parakeet/CMakeLists.txt +# +# Native Parakeet-TDT ASR (FastConformer encoder + TDT decoder) for +# FermionResearch/Phonon-2, the five-value QAT re-quantisation of +# nvidia/parakeet-tdt-0.6b-v3 (weights CC-BY-4.0). Built when +# LLAMA_BUILD_PARAKEET=ON (see the top-level CMakeLists.txt). +# +# Produces: +# parakeet_lib static library (GGUF loader, log-mel, encoder, TDT +# decoder, streaming). Uses ggml's gguf reader only; the +# compute is in pk-*.cpp with runtime-dispatched NEON +# i8mm / dotprod kernels, x86_64 AVX2 kernels and a +# scalar reference. +# parakeet-cli transcribe WAV files (host tools only). +# parakeet-bench RTFx / first-partial / RSS harness. + +add_library(parakeet_lib STATIC + include/parakeet.h + src/parakeet.cpp + src/pk-stream.cpp + src/pk-ops.cpp + src/pk-pool.cpp + src/pk-kernels.cpp + src/pk-kernels-dotprod.cpp + src/pk-kernels-i8mm.cpp + src/pk-kernels-avx2.cpp) +# hidden: folded PRIVATE into libelizainference, so pk_* must not be exported +set_target_properties(parakeet_lib PROPERTIES POSITION_INDEPENDENT_CODE ON + CXX_VISIBILITY_PRESET hidden + VISIBILITY_INLINES_HIDDEN ON) +target_include_directories(parakeet_lib PUBLIC ${CMAKE_CURRENT_SOURCE_DIR}/include + PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/src) +target_link_libraries(parakeet_lib PUBLIC ggml) +if (TARGET ggml-base) + target_link_libraries(parakeet_lib PUBLIC ggml-base) +endif() +target_compile_features(parakeet_lib PUBLIC cxx_std_17) +find_package(Threads REQUIRED) +target_link_libraries(parakeet_lib PUBLIC Threads::Threads) + +# The dotprod and i8mm tiers are separate translation units compiled for +# their extension and only entered after a runtime CPU check (HWCAP / +# HWCAP2 on Linux/Android, sysctl on Apple), so the baseline stays armv8-a. +set(PARAKEET_ARM64 OFF) +if (CMAKE_SYSTEM_PROCESSOR MATCHES "^(aarch64|arm64|ARM64)$" OR CMAKE_ANDROID_ARCH_ABI STREQUAL "arm64-v8a") + set(PARAKEET_ARM64 ON) +endif() +if (CMAKE_OSX_ARCHITECTURES STREQUAL "x86_64") # x86_64 build on an Apple silicon host + set(PARAKEET_ARM64 OFF) +endif() +list(LENGTH CMAKE_OSX_ARCHITECTURES PARAKEET_N_OSX_ARCHS) +if (PARAKEET_ARM64 AND NOT MSVC AND PARAKEET_N_OSX_ARCHS LESS_EQUAL 1) + set_source_files_properties(src/pk-kernels-dotprod.cpp PROPERTIES COMPILE_OPTIONS "-march=armv8.2-a+dotprod") + set_source_files_properties(src/pk-kernels-i8mm.cpp PROPERTIES COMPILE_OPTIONS "-march=armv8.2-a+dotprod+i8mm") +elseif (PARAKEET_ARM64) + message(WARNING "parakeet: dotprod/i8mm tiers need per-file -march; this configuration builds the scalar tier only") + target_compile_definitions(parakeet_lib PRIVATE PK_NO_ARM_EXT=1) +endif() + +# Likewise the AVX2 tier on x86_64 (Android x86_64 emulator images, desktop +# hosts): its own translation unit with -mavx2, entered after CPUID + XGETBV, +# so the baseline stays x86-64 (SSE2). +set(PARAKEET_X86_64 OFF) +if (CMAKE_SYSTEM_PROCESSOR MATCHES "^(x86_64|AMD64|amd64)$" OR CMAKE_ANDROID_ARCH_ABI STREQUAL "x86_64" OR CMAKE_OSX_ARCHITECTURES STREQUAL "x86_64") + set(PARAKEET_X86_64 ON) +endif() +if (CMAKE_OSX_ARCHITECTURES STREQUAL "arm64") + set(PARAKEET_X86_64 OFF) +endif() +if (PARAKEET_X86_64 AND MSVC) + set_source_files_properties(src/pk-kernels-avx2.cpp PROPERTIES COMPILE_OPTIONS "/arch:AVX2") +elseif (PARAKEET_X86_64 AND PARAKEET_N_OSX_ARCHS LESS_EQUAL 1) + set_source_files_properties(src/pk-kernels-avx2.cpp PROPERTIES COMPILE_OPTIONS "-mavx2") +else() + if (PARAKEET_X86_64) + message(WARNING "parakeet: the avx2 tier needs per-file -mavx2; this configuration builds the scalar tier only on x86_64") + endif() + target_compile_definitions(parakeet_lib PRIVATE PK_NO_X86_EXT=1) +endif() +if (NOT MSVC) + target_compile_options(parakeet_lib PRIVATE -O3) +endif() + +if (NOT CMAKE_SYSTEM_NAME STREQUAL "iOS" AND NOT CMAKE_SYSTEM_NAME STREQUAL "Android") + include(GNUInstallDirs) + add_executable(parakeet-cli tools/parakeet-cli.cpp) + target_link_libraries(parakeet-cli PRIVATE parakeet_lib) + target_include_directories(parakeet-cli PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/tools) + install(TARGETS parakeet-cli RUNTIME DESTINATION ${CMAKE_INSTALL_BINDIR}) + + if (NOT WIN32) # getrusage / getloadavg + add_executable(parakeet-bench tools/parakeet-bench.cpp) + target_link_libraries(parakeet-bench PRIVATE parakeet_lib) + target_include_directories(parakeet-bench PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/tools) + endif() +endif() + +if (LLAMA_BUILD_TESTS OR PARAKEET_BUILD_TESTS) + add_subdirectory(tests) +endif() diff --git a/tools/parakeet/README.md b/tools/parakeet/README.md new file mode 100644 index 000000000..8556c04f2 --- /dev/null +++ b/tools/parakeet/README.md @@ -0,0 +1,137 @@ +# tools/parakeet: Parakeet-TDT ASR for Phonon-2 + +[FermionResearch/Phonon-2](https://huggingface.co/FermionResearch/Phonon-2) is +[nvidia/parakeet-tdt-0.6b-v3](https://huggingface.co/nvidia/parakeet-tdt-0.6b-v3) (FastConformer encoder + +TDT decoder, 627M parameters, English) with every encoder linear retrained to a five-value alphabet +`{-hi, -lo, 0, +lo, +hi}` per row and everything else stored as 6-bit integer tables. This directory runs it +natively on the CPU inside libelizainference. It has no ONNX, MLX, Core ML or Python dependency at run time. + +**Licence (flagged, not a blocker).** The weights are CC-BY-4.0, derived from NVIDIA parakeet-tdt-0.6b-v3 +(CC-BY-4.0). The converter copies Fermion's `NOTICE` into the GGUF (`parakeet.notice`), and +`pk_model_notice()` returns it. Apps that ship the model must show the attribution. According to the NOTICE, +the training audio includes CHiME-6 (CC-BY-SA-4.0) and SPGISpeech (Kensho terms). Fermion's +`fermion_container.py` and `reference_transformers.py` are Apache-2.0. Fermion's fast CPU kernels ship only as +closed binaries in the `fermion-research` wheel, so none of their code is used here. Everything in this +directory is new code. + +## Pipeline + +| Stage | What runs | +|---|---| +| front end | pre-emphasis 0.97, 512-point STFT (400-sample periodic=False Hann, hop 160, centre zero-pad), slaney 128-mel, `log(x + 2^-24)`, per-feature normalisation over the utterance | +| subsampling | conv 3x3/2 (1→256) + ReLU, then [depthwise 3x3/2, pointwise 256→256 (int8), ReLU] x2, then a 4096→1024 linear (int8). 8x in time, 80 ms per frame. Computed in 1.28 s chunks, so memory stays flat | +| 24 conformer blocks | FF½ → rel-pos MHSA (Transformer-XL, 8 heads, full context) → conv module (pw1, GLU, depthwise k=9 with BatchNorm folded in, SiLU, pw2) → FF½ → LayerNorm. All 264 linears are five-value int8 GEMMs | +| TDT decoder | 2-layer LSTM prediction net (640), joint `ReLU(enc_proj + dec_proj)` → 8198 logits (8193 token/blank + 5 durations {0..4}), greedy, max 10 symbols per frame | + +## Numerics + +- **Five-value weights in int8.** `wmat_from_fv5` puts each row's `{lo, hi}` on an int8 grid `{L, H}` with + `H <= 127`. It picks the `H` in 64..127 that best keeps `lo/hi`, so the only weight error is + `|L*s - lo| <= 0.4% of hi` (reported as `fv5_grid_err`). `hi` is exact. +- **Activations.** Activations entering a linear are quantised to int8 per token (symmetric, + round-half-even). The products are exact int32, and the per-row and per-token scales are applied in fp32. + Fermion's own CPU engine takes the same approach. +- **Everything else is fp32:** attention, softmax, LayerNorm, convolutions, LSTM state and the joint + epilogue. +- **Position keys.** `W_pos · e(p)` depends only on the relative position `p`. It is computed once per layer, + stored as fp16 and grown on demand: 12 MB at 10 s windows, 37 MB at 30 s. +- **Decoder.** The LSTM, projection and joint-head matvecs use the exact int6 tables as int8 weights, with + int8 activations. + +## Kernels (`src/pk-kernels*.cpp`) + +The weights stay in the GGUF in kernel order (`fv5-pair16-v1`): 4-bit level codes, two rows by 16 k per +16-byte block. The model is mmapped and used in place. Per 8-row block, one `vqtbl1q_s8` per 16 weights +decodes the codes into an int8 panel that stays in L1. The panel is then multiplied against all tokens: + +| tier | instruction | runtime check | +|---|---|---| +| `i8mm` | `SMMLA` 2x8·8x2, 8 rows × 8 tokens per tile | `HWCAP2_I8MM` (Linux/Android), `hw.optional.arm.FEAT_I8MM` (Apple) | +| `dotprod` | `SDOT` on straight and half-swapped token pairs | `HWCAP_ASIMDDP`, `FEAT_DotProd` | +| `avx2` | `VPMADDWD` on sign-extended straight and half-swapped token pairs (x86_64) | CPUID AVX2 + `XGETBV` YMM state | +| `scalar` | portable C++ (old ARM, pre-AVX2 x86) | always | + +The dotprod and i8mm translation units are compiled with their own `-march`, and the avx2 one with `-mavx2` +(`/arch:AVX2` on MSVC), so the library baseline stays armv8-a / x86-64. The avx2 tier widens int8 to int16 +before `VPMADDWD`, so it is exact for every int8 input (no `VPMADDUBSW` saturation). All tiers produce +bit-identical int32 results (`tests/test_parakeet_kernels.cpp`, which also runs random full-range int8 +operands). `PK_KERNEL` set to `scalar`, `dotprod`, `i8mm` or `avx2` forces a tier. + +On an Apple silicon Mac the x86_64 build can be checked under Rosetta 2, which runs AVX2 but only reports it +when asked: `arch -x86_64 -e ROSETTA_ADVERTISE_AVX=1 build-x64/bin/test-parakeet-kernels` (configure with +`-DCMAKE_OSX_ARCHITECTURES=x86_64`). Rosetta timings are a sanity check, not x86 performance evidence. + +Threading uses a persistent pool (`src/pk-pool.cpp`) that spins briefly between jobs and then sleeps. By +default it starts one thread per fast core, at most 4. On Linux/Android those are the cores whose +`cpuinfo_max_freq` is within 20% of the fastest, and workers are pinned to them. On Apple it uses the +performance cores, with QoS user-interactive. + +## Streaming (`src/pk-stream.cpp`) + +The model attends over the whole input and has no cache-aware mode, so a stream re-decodes a window. A window +must not start in the middle of speech. When its input begins mid-word, the model drops or invents words: on +the same audio slices, the dense fp32 reference does exactly what this port does. Each decode therefore covers +a **pause-delimited segment**, from the segment start to the newest sample. + +- **Committing.** A token is committed when two consecutive decodes agree on it and it ends at least + `right_context_ms` before the newest sample. Commits happen on word boundaries only. Committed text never + changes. +- **Starting a new segment.** A segment ends when a pause of at least 150 ms lies between the last committed + token and the first uncommitted one. The next segment starts inside that pause, so windows begin in silence + and stay short. +- **Bounded windows.** Past `max_window_ms`, stable tokens are committed without agreement and the segment is + cut at the quietest point of that gap, which bounds every decode. +- **Finished words.** The newest word is committed once the agreed hypothesis has no token in the last + `right_context_ms`, so it does not wait for the next word. +- **Quiet audio.** With nothing uncommitted and only quiet audio since the last decode (every 10 ms frame below + -60 dBFS, or 30 dB under the speech level of the last window that held tokens), the decode is skipped and the + window keeps at most the newest 1.5 s of the pause. A long pause therefore costs no decodes and does not grow + the next window, and the model is not run on short windows of near-silence, which is where it invents words. +- **No speech at all.** A window past `max_window_ms` with no token keeps its newest half, which bounds noise + without speech. +- **Decode interval.** It is at least `min_hop_ms` and grows to a sixth of the window. +- **`finish()`** decodes the open segment with full right context. + +## Files + +- `convert_phonon2_to_gguf.py` reads the pinned `phonon-2.bps.tar.zst`, or an unpacked `model.fermion` and + `config.json`, and writes the GGUF. + - It checks every input against `pins.json` (HF revision `7e153e40…`, size and sha256). + - `--verify` re-reads the GGUF and checks that all 264 five-value records and 11 int tables expand + bit-identically to the container, that all 723 records are consumed, and that the mel filterbank equals + librosa's. +- `include/parakeet.h` is the C API: load, transcribe, word times, and stream open/feed/partial/finish. +- `tools/dump_parakeet_reference.py` writes fixtures from Fermion's `reference_transformers.py`, the dense + fp32 `transformers.ParakeetForTDT`. The fixtures hold the PCM, the normalised mel, the subsampling, layer-0 + and encoder outputs, and the token ids and text. +- `tests/test_parakeet_parity.cpp` compares the native pipeline with those fixtures stage by stage. +- `tools/parakeet-bench.cpp` and `tools/bench_parakeet.py` measure RTFx (wall and CPU time), first-window + partial latency, peak RSS and WER, plus stream WER and stream/batch agreement with `--stream`. +- libelizainference integration lives in `tools/omnivoice/src/eliza-asr-parakeet.inc`. + +## Build and use + +```sh +python tools/parakeet/convert_phonon2_to_gguf.py --archive phonon-2.bps.tar.zst --output phonon-2-fv5.gguf --verify +cmake -B build -DLLAMA_BUILD_PARAKEET=ON -DPARAKEET_BUILD_TESTS=ON ... +build/bin/parakeet-cli -m phonon-2-fv5.gguf [-t 4] [--stream] [--words] clip.wav +build/bin/test-parakeet-parity phonon-2-fv5.gguf fixtures/ +python tools/parakeet/tools/bench_parakeet.py --bench build/bin/parakeet-bench --model phonon-2-fv5.gguf \ + --rows rows.json --threads 1 4 --stream +``` + +In libelizainference, put the GGUF at `/asr/phonon-2-fv5.gguf`. On `mmap_acquire("asr")`, a bundle +whose `asr/` holds a GGUF with `general.architecture = parakeet-tdt` loads this backend. + +- `ELIZA_ASR_BACKEND=qwen3` keeps the Qwen3-ASR path; `=parakeet` requires this backend. +- With the backend loaded, `eliza_inference_asr_transcribe` and `_timed` run here. `_timed` returns real TDT + word times. +- `eliza_inference_asr_stream_supported()` returns 1 and the `asr_stream_*` session runs `pk-stream`. The + stream does not return token ids, because Parakeet's vocabulary is not the text model's. +- `ELIZA_ASR_THREADS` sets the thread count. `ELIZA_ASR_STREAM_{RIGHT,HOP,MAX_WINDOW}_MS` tune the stream. + +## Evidence boundary + +The measurements in the PR come from an Apple M4 Max used as an ARM CPU proxy, on a shared, heavily loaded +host. They are not phone evidence. The Android arm64-v8a and x86_64 builds are compile evidence only. Speed, +memory and thermals on a phone need a device run. diff --git a/tools/parakeet/convert_phonon2_to_gguf.py b/tools/parakeet/convert_phonon2_to_gguf.py new file mode 100644 index 000000000..91cb562bb --- /dev/null +++ b/tools/parakeet/convert_phonon2_to_gguf.py @@ -0,0 +1,577 @@ +#!/usr/bin/env python3 +# SPDX-License-Identifier: MIT +""" +convert_phonon2_to_gguf.py - convert FermionResearch/Phonon-2 (a five-value +QAT re-quantisation of nvidia/parakeet-tdt-0.6b-v3, CC-BY-4.0 weights) into the +GGUF that tools/parakeet loads (general.architecture = parakeet-tdt). + +Input: the pinned release archive `phonon-2.bps.tar.zst` (or the already +unpacked `model.fermion` + NeMo `config.json`). Every input is checked against +tools/parakeet/pins.json (size + sha256) before it is parsed. + +Container (`fermion-five-value-parakeet-v1`): u64 LE header length, JSON +header {format, index:[{n,k,shape,b}]}, then the records in index order. + five_value [O,I]: O rows of ceil(I/5) base-3 bytes (LSD first, sign=code-1), + then a little-endian bit per non-zero code (1 = hi), then + fp16 lo[O], fp16 hi[O]. Weight = sign * (is_hi ? hi : lo). + int6/int8 [O,...]: symmetric per-row integers (6-bit packed 4 per 3 bytes, + minus 32) followed by fp16 row scales. + fp16: raw half floats. + +GGUF layout written here (all lossless with respect to the container): + * five-value linears -> `.fv5` (I8 bytes [O/2, I], 4-bit level + indices in the `fv5-pair16-v1` kernel order, see pack_pair16; level 0..4 + = {-hi,-lo,0,+lo,+hi}) and `.lohi` (F32 [O,2]). q/k/v are + concatenated into one `attn.qkv`. The file is mmapped and used in place. + * int6 tables -> `.q` (I8, the exact integers) + `.scale` (F32). + * fp16 records -> F32 (exact widening). + * BatchNorm is folded into the depthwise conv (F32 weight + bias); this is + the only arithmetic transform and it is done in float64. + * The LSTM's two bias vectors are summed (float64 -> F32). + * The mel filterbank (librosa slaney, as ParakeetFeatureExtractor builds + it) is stored as `feat.mel_fb` [128, 257] F32. + +--verify re-reads the written GGUF and re-expands every tensor; five-value +and int tensors must equal the container's own expansion bit for bit (the +same check as Fermion's G-PLANE/G-LOAD gates), and all 723 records must be +consumed. + +The output is byte-reproducible for the same inputs and converter. + +usage: + convert_phonon2_to_gguf.py --archive phonon-2.bps.tar.zst --output phonon-2-fv5.gguf [--verify] + convert_phonon2_to_gguf.py --container model.fermion --nemo-config config.json \ + --notice NOTICE --output phonon-2-fv5.gguf [--verify] +""" + +from __future__ import annotations + +import argparse +import hashlib +import io +import json +import re +import subprocess +import sys +import tarfile +from pathlib import Path +from typing import Any, NoReturn + +import numpy as np + +_HERE = Path(__file__).resolve().parent +sys.path.insert(0, str(_HERE.parent.parent / "gguf-py")) + +import gguf # noqa: E402 + +PINS_PATH = _HERE / "pins.json" +ARCH = "parakeet-tdt" +NS = "parakeet" +FORMAT = "fermion-five-value-parakeet-v1" + +HP: dict[str, Any] = { + "d_model": 1024, + "n_layers": 24, + "n_heads": 8, + "ff_dim": 4096, + "conv_kernel": 9, + "sub_channels": 256, + "sub_factor": 8, + "n_mels": 128, + "sample_rate": 16000, + "n_fft": 512, + "win_length": 400, + "hop_length": 160, + "preemph": 0.97, + "log_guard": 2.0 ** -24, + "norm_eps": 1e-5, + "ln_eps": 1e-5, + "bn_eps": 1e-5, + "pos_max_len": 5000, + "pred_hidden": 640, + "pred_layers": 2, + "joint_hidden": 640, + "vocab": 8192, + "blank_id": 8192, + "durations": [0, 1, 2, 3, 4], + "max_symbols": 10, +} + + +def sha256_bytes(b: bytes) -> str: + return hashlib.sha256(b).hexdigest() + + +def die(msg: str) -> NoReturn: + print(f"convert_phonon2_to_gguf: {msg}", file=sys.stderr) + sys.exit(1) + + +def check_pin(name: str, data: bytes, pin: dict) -> None: + if len(data) != pin["size"] or sha256_bytes(data) != pin["sha256"]: + die(f"{name}: size/sha256 mismatch against pins.json " + f"(got {len(data)} bytes, sha256 {sha256_bytes(data)})") + + +# --------------------------------------------------------------- inputs +def read_archive(path: Path, pins: dict) -> dict[str, bytes]: + data = path.read_bytes() + check_pin(path.name, data, pins["files"]["phonon-2.bps.tar.zst"]) + try: + import zstandard # type: ignore + raw = zstandard.ZstdDecompressor().stream_reader(io.BytesIO(data)).read() + except ImportError: + raw = subprocess.run(["zstd", "-q", "-d", "-c"], input=data, check=True, + stdout=subprocess.PIPE).stdout + members: dict[str, bytes] = {} + with tarfile.open(fileobj=io.BytesIO(raw), mode="r:") as tar: + manifest = None + for m in tar: + fh = tar.extractfile(m) + if fh is None: + continue + blob = fh.read() + if m.name == "bps_manifest.json": + manifest = json.loads(blob) + continue + if manifest is None: + die("bps_manifest.json must be the first archive member") + rows = {r["path"]: r for r in manifest["files"]} + rel = m.name[:-4] if m.name.endswith(".bps") else m.name + row = rows[rel] + if "transform" in row: + die(f"{rel}: byte-plane transform present; the pinned release stores members verbatim") + if sha256_bytes(blob) != row["original_sha256"] or len(blob) != row["original_bytes"]: + die(f"{rel}: archive member checksum mismatch") + members[rel] = blob + for name, pin in pins["archive_members"].items(): + if name not in members: + die(f"archive is missing {name}") + check_pin(name, members[name], pin) + return members + + +# ------------------------------------------------------------ container +def trits(buf: np.ndarray, o: int, i: int) -> np.ndarray: + rb = (i + 4) // 5 + x = buf[: o * rb].reshape(o, rb).astype(np.uint16) + d = np.empty((o, rb, 5), dtype=np.uint8) + p = 1 + for k in range(5): + d[:, :, k] = (x // p) % 3 + p *= 3 + return d.reshape(o, rb * 5)[:, :i] + + +def parse_five_value(blob: bytes, shape) -> dict: + o, i = shape + rb = (i + 4) // 5 + u8 = np.frombuffer(blob, dtype=np.uint8) + codes = trits(u8, o, i) + if codes.max() > 2: + die("five_value: base-3 digit out of range") + nz = codes != 1 + nnz = int(nz.sum()) + off = o * rb + nbytes = (nnz + 7) // 8 + bits = np.unpackbits(u8[off: off + nbytes], bitorder="little")[:nnz].astype(bool) + off += nbytes + lo = np.frombuffer(blob[off: off + 2 * o], dtype=np.float16).copy() + hi = np.frombuffer(blob[off + 2 * o: off + 4 * o], dtype=np.float16).copy() + if off + 4 * o != len(blob): + die("five_value: record length mismatch") + is_hi = np.zeros((o, i), dtype=bool) + is_hi[nz] = bits + sign = codes.astype(np.int8) - 1 + return {"sign": sign, "is_hi": is_hi, "lo": lo, "hi": hi} + + +def parse_intn(blob: bytes, shape, bits: int) -> dict: + o = shape[0] + total = int(np.prod(shape)) + sc = np.frombuffer(blob[-2 * o:], dtype=np.float16).copy() + body = blob[: -2 * o] + if bits == 8: + q = np.frombuffer(body, dtype=np.int8)[:total].astype(np.int32) + elif bits == 6: + b = np.frombuffer(body, dtype=np.uint8).reshape(-1, 3).astype(np.uint32) + packed = b[:, 0] | (b[:, 1] << 8) | (b[:, 2] << 16) + u = np.stack([(packed >> s) & 0x3F for s in (0, 6, 12, 18)], axis=1).ravel() + q = u[:total].astype(np.int32) - 32 + else: + die(f"int{bits}: unsupported") + return {"q": q.reshape(o, total // o).astype(np.int8), "scale": sc, "shape": tuple(shape)} + + +def read_container(blob: bytes): + n = int.from_bytes(blob[:8], "little") + header = json.loads(blob[8: 8 + n]) + if header.get("format") != FORMAT: + die(f"container format {header.get('format')!r} != {FORMAT}") + off = 8 + n + recs = {} + for e in header["index"]: + b = blob[off: off + e["b"]] + off += e["b"] + k, shape = e["k"], tuple(e["shape"]) + if k == "five_value": + recs[e["n"]] = ("fv5", parse_five_value(b, shape)) + elif k.startswith("int"): + recs[e["n"]] = ("int", parse_intn(b, shape, int(k[3:]))) + elif k == "fp16": + recs[e["n"]] = ("f16", np.frombuffer(b, dtype=np.float16).reshape(shape).copy()) + else: + die(f"unknown record kind {k}") + if off != len(blob): + die("container has trailing bytes") + return recs, header["index"] + + +# ------------------------------------------------------------- features +def mel_filterbank(sr: int, n_fft: int, n_mels: int) -> np.ndarray: + """librosa.filters.mel(sr, n_fft, n_mels, fmin=0, fmax=sr/2, htk=False, norm='slaney').""" + f_sp = 200.0 / 3 + min_log_hz = 1000.0 + min_log_mel = min_log_hz / f_sp + logstep = np.log(6.4) / 27.0 + + def hz_to_mel(f): + f = np.asanyarray(f, dtype=np.float64) + m = f / f_sp + return np.where(f >= min_log_hz, min_log_mel + np.log(np.maximum(f, 1e-10) / min_log_hz) / logstep, m) + + def mel_to_hz(m): + m = np.asanyarray(m, dtype=np.float64) + f = f_sp * m + return np.where(m >= min_log_mel, min_log_hz * np.exp(logstep * (m - min_log_mel)), f) + + fftfreqs = np.fft.rfftfreq(n=n_fft, d=1.0 / sr) + mel_f = mel_to_hz(np.linspace(hz_to_mel(0.0), hz_to_mel(sr / 2.0), n_mels + 2)) + fdiff = np.diff(mel_f) + ramps = np.subtract.outer(mel_f, fftfreqs) + w = np.zeros((n_mels, 1 + n_fft // 2), dtype=np.float32) + for i in range(n_mels): + lower = -ramps[i] / fdiff[i] + upper = ramps[i + 2] / fdiff[i + 1] + w[i] = np.maximum(0, np.minimum(lower, upper)) + enorm = 2.0 / (mel_f[2: n_mels + 2] - mel_f[:n_mels]) + w *= enorm[:, np.newaxis] + return w + + +# ------------------------------------------------------------- packing +def fv5_levels(r: dict) -> np.ndarray: + """Level index 0..4 = {-hi,-lo,0,+lo,+hi}.""" + s = r["sign"].astype(np.int16) + mag = np.where(r["is_hi"], 2, 1).astype(np.int16) + return (2 + s * mag).astype(np.uint8) + + +def pack_pair16(levels: np.ndarray) -> np.ndarray: + """Kernel layout `fv5-pair16-v1`, [O/2, K] bytes. + + Rows are taken in pairs (2p, 2p+1) and K in groups of 16. For pair p and + group g, chunk A = [r0 k16g..16g+7, r1 k16g..16g+7] and chunk B = + [r0 k16g+8..16g+15, r1 k16g+8..16g+15], with 8 added to every r1 level. + Byte j of the 16-byte block = A[j] | B[j] << 4. A NEON kernel decodes 32 + weights with one load, an AND, a shift and two table lookups (one 16-byte + table per row pair holds both rows' five int8 levels), straight into the + 2x8 operand order of SMMLA / 8-wide SDOT pairs.""" + o, k = levels.shape + if o % 2 or k % 16: + die(f"five-value shape {levels.shape}: rows must be even and width a multiple of 16") + lv = levels.reshape(o // 2, 2, k // 16, 2, 8) # pair, row, group, half, 8 + lv = lv.astype(np.uint8).copy() + lv[:, 1] += 8 + seq = lv.transpose(0, 2, 3, 1, 4).reshape(o // 2, k // 16, 2, 16) # pair, group, half(A/B), 16 + out = seq[:, :, 0, :] | (seq[:, :, 1, :] << 4) + return out.reshape(o // 2, k).astype(np.uint8) + + +def unpack_pair16(packed: np.ndarray) -> np.ndarray: + p, k = packed.shape + b = packed.reshape(p, k // 16, 16) + seq = np.stack([b & 0x0F, b >> 4], axis=2) # pair, group, half, 16 + seq = seq.reshape(p, k // 16, 2, 2, 8).transpose(0, 3, 1, 2, 4) # pair, row, group, half, 8 + seq = seq.astype(np.int16) + seq[:, 1] -= 8 + if seq.min() < 0 or seq.max() > 4: + die("unpack_pair16: level out of range") + return seq.reshape(2 * p, k).astype(np.uint8) + + +def fv5_expand(levels: np.ndarray, lohi: np.ndarray) -> np.ndarray: + lo = lohi[:, 0:1].astype(np.float16) + hi = lohi[:, 1:2].astype(np.float16) + table = np.stack([-hi, -lo, np.zeros_like(lo), lo, hi], axis=-1)[:, 0, :] # [O,5] fp16 + return np.take_along_axis(table, levels.astype(np.int64), axis=1) + + +def container_fv5_expand(r: dict) -> np.ndarray: + mag = np.where(r["is_hi"], r["hi"][:, None], r["lo"][:, None]) + return (r["sign"].astype(np.float16) * mag).astype(np.float16) + + +# --------------------------------------------------------------- writer +class Builder: + def __init__(self, recs: dict): + self.recs = recs + self.used: set[str] = set() + self.tensors: list[tuple[str, np.ndarray]] = [] + self.fv5_src: dict[str, list[str]] = {} + self.int_src: dict[str, str] = {} + + def take(self, name: str, kind: str): + if name not in self.recs: + die(f"missing record {name}") + k, v = self.recs[name] + if k != kind: + die(f"record {name}: kind {k} != {kind}") + self.used.add(name) + return v + + def add(self, name: str, arr: np.ndarray): + self.tensors.append((name, np.ascontiguousarray(arr))) + + def fv5(self, out: str, names: list[str], o_expect: int, i_expect: int): + lv, lohi = [], [] + for n in names: + r = self.take(n, "fv5") + o, i = r["sign"].shape + if i != i_expect: + die(f"{n}: width {i} != {i_expect}") + lv.append(fv5_levels(r)) + lohi.append(np.stack([r["lo"].astype(np.float32), r["hi"].astype(np.float32)], axis=1)) + levels = np.concatenate(lv, axis=0) + if levels.shape[0] != o_expect: + die(f"{out}: rows {levels.shape[0]} != {o_expect}") + self.add(out + ".fv5", pack_pair16(levels).view(np.int8)) + self.add(out + ".lohi", np.concatenate(lohi, axis=0).astype(np.float32)) + self.fv5_src[out] = names + + def intq(self, out: str, name: str, shape2d: tuple[int, int] | None = None): + r = self.take(name, "int") + q = r["q"] + if shape2d is not None and q.shape != shape2d: + die(f"{name}: shape {q.shape} != {shape2d}") + self.add(out + ".q", q.astype(np.int8)) + self.add(out + ".scale", r["scale"].astype(np.float32)) + self.int_src[out] = name + + def f32(self, out: str, name: str, shape=None): + v = self.take(name, "f16").astype(np.float32) + if shape is not None: + v = v.reshape(shape) + self.add(out, v) + + +def build(recs: dict) -> Builder: + b = Builder(recs) + d, ff, L = HP["d_model"], HP["ff_dim"], HP["n_layers"] + C = HP["sub_channels"] + # subsampling (dw_striding): conv0 (1->C, 3x3 s2) relu, [dw 3x3 s2, pw 1x1, relu] x2, linear + b.f32("enc.sub.conv0.w", "encoder.subsampling.layers.0.weight", (C, 9)) + b.f32("enc.sub.conv0.b", "encoder.subsampling.layers.0.bias") + for src, dst in ((2, "dw1"), (5, "dw2")): + b.f32(f"enc.sub.{dst}.w", f"encoder.subsampling.layers.{src}.weight", (C, 9)) + b.f32(f"enc.sub.{dst}.b", f"encoder.subsampling.layers.{src}.bias") + for src, dst in ((3, "pw1"), (6, "pw2")): + b.intq(f"enc.sub.{dst}", f"encoder.subsampling.layers.{src}.weight", (C, C)) + b.f32(f"enc.sub.{dst}.b", f"encoder.subsampling.layers.{src}.bias") + b.intq("enc.sub.out", "encoder.subsampling.linear.weight", (d, C * (HP["n_mels"] // 8))) + b.f32("enc.sub.out.b", "encoder.subsampling.linear.bias") + + for l in range(L): + p = f"encoder.layers.{l}." + o = f"enc.l{l}." + b.fv5(o + "ff1.w1", [p + "feed_forward1.linear1"], ff, d) + b.fv5(o + "ff1.w2", [p + "feed_forward1.linear2"], d, ff) + b.fv5(o + "attn.qkv", [p + "self_attn.q_proj", p + "self_attn.k_proj", p + "self_attn.v_proj"], 3 * d, d) + b.fv5(o + "attn.out", [p + "self_attn.o_proj"], d, d) + b.fv5(o + "attn.pos", [p + "self_attn.relative_k_proj"], d, d) + b.fv5(o + "conv.pw1", [p + "conv.pointwise_conv1"], 2 * d, d) + b.fv5(o + "conv.pw2", [p + "conv.pointwise_conv2"], d, d) + b.fv5(o + "ff2.w1", [p + "feed_forward2.linear1"], ff, d) + b.fv5(o + "ff2.w2", [p + "feed_forward2.linear2"], d, ff) + b.f32(o + "attn.bias_u", p + "self_attn.bias_u") + b.f32(o + "attn.bias_v", p + "self_attn.bias_v") + for nm in ("norm_feed_forward1", "norm_self_att", "norm_conv", "norm_feed_forward2", "norm_out"): + b.f32(o + nm + ".w", p + nm + ".weight") + b.f32(o + nm + ".b", p + nm + ".bias") + # depthwise conv (int6, no bias) with BatchNorm folded, in float64 + r = b.take(p + "conv.depthwise_conv.weight", "int") + k = HP["conv_kernel"] + w = r["q"].astype(np.float64).reshape(d, k) * r["scale"].astype(np.float64)[:, None] + gamma = b.take(p + "conv.norm.weight", "f16").astype(np.float64) + beta = b.take(p + "conv.norm.bias", "f16").astype(np.float64) + mean = b.take(p + "conv.norm.running_mean", "f16").astype(np.float64) + var = b.take(p + "conv.norm.running_var", "f16").astype(np.float64) + b.take(p + "conv.norm.num_batches_tracked", "f16") # unused at inference (stored as fp16, may be inf) + inv = gamma / np.sqrt(var + HP["bn_eps"]) + b.add(o + "conv.dw.w", (w * inv[:, None]).astype(np.float32)) + b.add(o + "conv.dw.b", (beta - mean * inv).astype(np.float32)) + + H, V = HP["pred_hidden"], HP["vocab"] + b.intq("enc_proj", "encoder_projector.weight", (HP["joint_hidden"], d)) + b.f32("enc_proj.b", "encoder_projector.bias") + b.intq("dec.embed", "decoder.embedding.weight", (V + 1, H)) + for l in range(HP["pred_layers"]): + b.intq(f"dec.lstm{l}.ih", f"decoder.lstm.weight_ih_l{l}", (4 * H, H)) + b.intq(f"dec.lstm{l}.hh", f"decoder.lstm.weight_hh_l{l}", (4 * H, H)) + bi = b.take(f"decoder.lstm.bias_ih_l{l}", "f16").astype(np.float64) + bh = b.take(f"decoder.lstm.bias_hh_l{l}", "f16").astype(np.float64) + b.add(f"dec.lstm{l}.b", (bi + bh).astype(np.float32)) + b.intq("dec.proj", "decoder.decoder_projector.weight", (HP["joint_hidden"], H)) + b.f32("dec.proj.b", "decoder.decoder_projector.bias") + b.intq("joint.head", "joint.head.weight", (V + 1 + len(HP["durations"]), HP["joint_hidden"])) + b.f32("joint.head.b", "joint.head.bias") + + unused = sorted(set(recs) - b.used) + if unused: + die(f"{len(unused)} container records not consumed, e.g. {unused[:4]}") + b.add("feat.mel_fb", mel_filterbank(HP["sample_rate"], HP["n_fft"], HP["n_mels"])) + return b + + +def git_blob_sha(path: Path) -> str: + data = path.read_bytes() + return hashlib.sha1(b"blob %d\0" % len(data) + data).hexdigest() + + +def write_gguf(out: Path, b: Builder, labels: list[str], notice: str, pins: dict, src: dict): + w = gguf.GGUFWriter(str(out), ARCH) + w.add_name("Phonon-2") + w.add_string("general.license", "CC-BY-4.0") + w.add_string("general.license.link", "https://creativecommons.org/licenses/by/4.0/") + w.add_string("general.base_model", pins["model"]["base_model"]) + w.add_string(f"{NS}.source.repo", pins["model"]["repo"]) + w.add_string(f"{NS}.source.revision", pins["model"]["revision"]) + for k, v in src.items(): + w.add_string(f"{NS}.source.{k}", v) + w.add_string(f"{NS}.notice", notice) + w.add_string(f"{NS}.converter.blob_sha1", git_blob_sha(Path(__file__))) + w.add_string(f"{NS}.weights.format", "fv5-pair16-v1") + for key in ("d_model", "n_layers", "n_heads", "ff_dim", "conv_kernel", "sub_channels", "sub_factor", + "n_mels", "sample_rate", "n_fft", "win_length", "hop_length", "pos_max_len", + "pred_hidden", "pred_layers", "joint_hidden", "vocab", "blank_id", "max_symbols"): + w.add_uint32(f"{NS}.{key}", int(HP[key])) + for key in ("preemph", "log_guard", "norm_eps", "ln_eps"): + w.add_float32(f"{NS}.{key}", float(HP[key])) + w.add_array(f"{NS}.durations", [int(x) for x in HP["durations"]]) + w.add_string(f"{NS}.language", "en") + w.add_tokenizer_model("sentencepiece-bpe") + w.add_token_list(labels) + for name, arr in b.tensors: + if arr.dtype == np.float32: + w.add_tensor(name, arr, raw_dtype=gguf.GGMLQuantizationType.F32) + elif arr.dtype == np.int8: + w.add_tensor(name, arr, raw_dtype=gguf.GGMLQuantizationType.I8) + else: + die(f"{name}: unexpected dtype {arr.dtype}") + w.write_header_to_file() + w.write_kv_data_to_file() + w.write_tensors_to_file() + w.close() + + +def verify(out: Path, b: Builder, recs: dict): + r = gguf.GGUFReader(str(out)) + t = {x.name: np.array(x.data) for x in r.tensors} + nfv = nint = 0 + for name, srcs in b.fv5_src.items(): + lv = unpack_pair16(t[name + ".fv5"].view(np.uint8)) + lohi = t[name + ".lohi"] + got = fv5_expand(lv, lohi) + row = 0 + for s in srcs: + ref = container_fv5_expand(recs[s][1]) + o = ref.shape[0] + if not np.array_equal(got[row: row + o].view(np.uint16), ref.view(np.uint16)): + die(f"verify: {name} rows {row}:{row + o} differ from container record {s}") + row += o + nfv += 1 + for name, s in b.int_src.items(): + q = t[name + ".q"].astype(np.float32) + sc = t[name + ".scale"].astype(np.float32) + r0 = recs[s][1] + ref = r0["q"].astype(np.float32) * r0["scale"].astype(np.float32)[:, None] + if not np.array_equal((q * sc[:, None]).view(np.uint32), ref.view(np.uint32)): + die(f"verify: {name} differs from container record {s}") + nint += 1 + try: + import librosa # type: ignore + ref = librosa.filters.mel(sr=HP["sample_rate"], n_fft=HP["n_fft"], n_mels=HP["n_mels"], + fmin=0.0, fmax=HP["sample_rate"] / 2, norm="slaney") + err = float(np.max(np.abs(ref.astype(np.float32) - t["feat.mel_fb"]))) + if err > 1e-7: + die(f"verify: mel filterbank differs from librosa by {err}") + mel = f"librosa {getattr(librosa, '__version__', '?')} max abs diff {err:.1e}" + except ImportError: + mel = "librosa not installed; filterbank not cross-checked" + print(f"verify: OK, {nfv} five-value records and {nint} int records bit-identical; " + f"{len(recs)} records consumed; mel: {mel}") + + +def main(): + ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) + ap.add_argument("--archive", type=Path, help="phonon-2.bps.tar.zst at the pinned revision") + ap.add_argument("--container", type=Path, help="unpacked model.fermion") + ap.add_argument("--nemo-config", type=Path, help="unpacked config.json (with 'labels')") + ap.add_argument("--notice", type=Path, help="NOTICE from the model repo (default: next to --archive)") + ap.add_argument("--output", type=Path, required=True) + ap.add_argument("--verify", action="store_true") + a = ap.parse_args() + pins = json.loads(PINS_PATH.read_text()) + + if a.archive: + members = read_archive(a.archive, pins) + container, cfg_bytes = members["model.fermion"], members["config.json"] + notice_path = a.notice or a.archive.parent / "NOTICE" + elif a.container and a.nemo_config: + container = a.container.read_bytes() + check_pin("model.fermion", container, pins["archive_members"]["model.fermion"]) + cfg_bytes = a.nemo_config.read_bytes() + check_pin("config.json", cfg_bytes, pins["archive_members"]["config.json"]) + notice_path = a.notice or a.container.parent / "NOTICE" + else: + ap.error("give --archive, or --container and --nemo-config") + if not notice_path.exists(): + die(f"NOTICE not found at {notice_path} (pass --notice); CC-BY-4.0 attribution must ship with the GGUF") + notice = notice_path.read_bytes() + check_pin("NOTICE", notice, pins["files"]["NOTICE"]) + + cfg = json.loads(cfg_bytes) + labels = list(cfg["labels"]) + if len(labels) != HP["vocab"]: + die(f"labels: {len(labels)} != {HP['vocab']}") + pre = cfg["preprocessor"] + enc = cfg["encoder"] + checks = [ + (pre["sample_rate"], HP["sample_rate"]), (pre["features"], HP["n_mels"]), (pre["n_fft"], HP["n_fft"]), + (round(pre["window_size"] * pre["sample_rate"]), HP["win_length"]), + (round(pre["window_stride"] * pre["sample_rate"]), HP["hop_length"]), (pre["normalize"], "per_feature"), + (enc["d_model"], HP["d_model"]), (enc["n_layers"], HP["n_layers"]), (enc["n_heads"], HP["n_heads"]), + (enc["conv_kernel_size"], HP["conv_kernel"]), (enc["subsampling"], "dw_striding"), + (enc["subsampling_factor"], HP["sub_factor"]), (enc["subsampling_conv_channels"], HP["sub_channels"]), + (enc["self_attention_model"], "rel_pos"), (list(enc["att_context_size"]), [-1, -1]), + (enc["xscaling"], False), (enc["untie_biases"], True), (enc["use_bias"], False), + (cfg["model_defaults"]["tdt_durations"], HP["durations"]), + ] + for got, want in checks: + if got != want: + die(f"NeMo config mismatch: {got!r} != {want!r}") + + recs, index = read_container(container) + if len(index) != 723: + die(f"container has {len(index)} records, expected 723") + b = build(recs) + src = {"container_sha256": pins["archive_members"]["model.fermion"]["sha256"], + "archive_sha256": pins["files"]["phonon-2.bps.tar.zst"]["sha256"]} + write_gguf(a.output, b, labels, notice.decode("utf-8"), pins, src) + print(f"wrote {a.output} ({a.output.stat().st_size} bytes, {len(b.tensors)} tensors)") + if a.verify: + verify(a.output, b, recs) + + +if __name__ == "__main__": + main() diff --git a/tools/parakeet/include/parakeet.h b/tools/parakeet/include/parakeet.h new file mode 100644 index 000000000..8d4f86407 --- /dev/null +++ b/tools/parakeet/include/parakeet.h @@ -0,0 +1,119 @@ +/* + * tools/parakeet: native Parakeet-TDT (FastConformer encoder + TDT decoder) + * speech recognition for FermionResearch/Phonon-2, the five-value QAT + * re-quantisation of nvidia/parakeet-tdt-0.6b-v3 (weights CC-BY-4.0). + * + * Pure C ABI so the fused libelizainference FFI, the CLI tools and tests can + * share it. Everything runs on the CPU: the 264 encoder linears use the + * five-value int8 kernels in pk-kernels*.cpp (NEON i8mm / dotprod, x86_64 + * AVX2, scalar reference), the rest is fp32. There is no ONNX, MLX or Core ML path. + * + * Threading: a model owns a worker pool sized at load (pk_params.n_threads). + * Calls on one model are serialised by the caller (the FFI holds its ASR + * mutex); streams on one model must not decode concurrently. + */ +#pragma once + +#include +#include + +#ifdef __cplusplus +extern "C" { +#endif + +typedef struct pk_model pk_model; +typedef struct pk_stream pk_stream; + +enum pk_status { + PK_OK = 0, + PK_E_INVALID = -1, + PK_E_IO = -2, + PK_E_FORMAT = -3, + PK_E_CANCELLED = -4, + PK_E_NOMEM = -5, +}; + +typedef struct pk_params { + int n_threads; /* <= 0: one per big core (Android) / performance core (Apple), capped at 4 */ + int kernel; /* 0 = best available, 1 = scalar reference, 2 = dotprod, 3 = i8mm, 4 = avx2 */ + int pin_threads; /* 1 = pin workers to the fastest cores where the OS allows it (Linux/Android) */ + int use_mmap; /* 1 = mmap the GGUF (clean, reclaimable pages); 0 = read it into private memory */ +} pk_params; + +pk_params pk_default_params(void); + +/* Load a parakeet-tdt GGUF (convert_phonon2_to_gguf.py). The file is mmapped; + * the five-value weights are used in place. Returns NULL and writes a + * message into err (if err_len > 0) on failure. */ +pk_model * pk_model_load(const char * path, pk_params params, char * err, size_t err_len); +void pk_model_free(pk_model * model); + +/* 1 when the GGUF at `path` declares general.architecture = "parakeet-tdt". */ +int pk_gguf_is_parakeet(const char * path); + +/* Kernel tier actually in use: "i8mm", "dotprod", "avx2" or "scalar". */ +const char * pk_model_kernel_name(const pk_model * model); +int pk_model_n_threads(const pk_model * model); +int pk_model_sample_rate(const pk_model * model); /* 16000 */ +/* The NOTICE / attribution text carried in the GGUF (CC-BY-4.0). */ +const char * pk_model_notice(const pk_model * model); + +/* Cancellation hook polled between encoder layers and decode steps. */ +typedef int (*pk_cancel_cb)(void * user); + +typedef struct pk_token { + int32_t id; /* SentencePiece id (0..8191) */ + int32_t frame; /* encoder frame (80 ms) where the token was emitted */ + int32_t duration; /* predicted TDT duration in frames */ +} pk_token; + +typedef struct pk_result { + char * text; /* UTF-8, malloc'd */ + pk_token * tokens; /* malloc'd */ + size_t n_tokens; + int n_frames; /* encoder frames of the input */ + /* stage timings (ms, wall) of the call */ + double t_mel_ms, t_encoder_ms, t_decoder_ms; +} pk_result; + +/* Transcribe 16 kHz mono fp32 PCM (callers resample). Fills *out (free with + * pk_result_free). Long inputs are cut into windows of at most + * PK_MAX_WINDOW_S seconds at low-energy points and the transcripts joined. */ +int pk_transcribe(pk_model * model, const float * pcm16k, size_t n_samples, + pk_cancel_cb cancel, void * cancel_user, pk_result * out); +void pk_result_free(pk_result * res); + +/* Words with times: splits the transcript at SentencePiece word starts. The + * caller's arrays have capacity *io_n; *io_n receives the count written. */ +int pk_result_words(const pk_model * model, const pk_result * res, + int * word_start_ms, int * word_end_ms, size_t * io_n); + +/* ---- streaming (pause-delimited windowed re-decode, local agreement) -- */ + +typedef struct pk_stream_params { + int right_context_ms; /* tokens newer than this are never committed (default 800) */ + int min_hop_ms; /* minimum new audio between decodes (default 400); the + interval grows to a sixth of the window */ + int max_window_ms; /* longest decode window; past it stable tokens are committed + without agreement and the segment is cut (default 12000) */ +} pk_stream_params; + +pk_stream_params pk_stream_default_params(void); + +pk_stream * pk_stream_open(pk_model * model, pk_stream_params params); +/* Append 16 kHz mono fp32 PCM. Never decodes. */ +int pk_stream_feed(pk_stream * st, const float * pcm16k, size_t n_samples); +/* Decode if at least min_hop_ms of new audio arrived since the last decode, + * then write committed + tentative text. Returns bytes written (excluding + * NUL) or a negative pk_status. */ +int pk_stream_partial(pk_stream * st, char * out, size_t cap, pk_cancel_cb cancel, void * cancel_user); +/* Decode the remaining audio with full right context and write the final + * transcript. */ +int pk_stream_finish(pk_stream * st, char * out, size_t cap, pk_cancel_cb cancel, void * cancel_user); +/* Stats for benchmarks: number of decodes and total decode wall time. */ +void pk_stream_stats(const pk_stream * st, int * n_decodes, double * decode_ms, double * first_partial_ms); +void pk_stream_close(pk_stream * st); + +#ifdef __cplusplus +} +#endif diff --git a/tools/parakeet/pins.json b/tools/parakeet/pins.json new file mode 100644 index 000000000..f10f433c1 --- /dev/null +++ b/tools/parakeet/pins.json @@ -0,0 +1,71 @@ +{ + "model": { + "repo": "FermionResearch/Phonon-2", + "revision": "7e153e4054c0a10db6d47d6fff5c062d4c77c154", + "base_model": "nvidia/parakeet-tdt-0.6b-v3", + "license": "CC-BY-4.0", + "license_note": "Weights: CC-BY-4.0 (derived from NVIDIA parakeet-tdt-0.6b-v3, CC-BY-4.0). Attribution and the NOTICE text must ship with the model. NOTICE lists training audio from CHiME-6 (CC-BY-SA-4.0) and SPGISpeech (Kensho terms). Fermion code (fermion_container.py, reference_transformers.py) is Apache-2.0. Flagged for review, not a build blocker.", + "language": "en" + }, + "files": { + "phonon-2.bps.tar.zst": { + "size": 163515201, + "sha256": "98125795b6dda72f5c6eee9ba33d19815df65dcb18b50a357bf9f73c9935309e" + }, + "config.json": { + "size": 759, + "sha256": "422379e411f14dde97174554deead106ec1e4816e14dd78a1bf99517c6dcc663" + }, + "packed_manifest.json": { + "size": 225, + "sha256": "399e9397cf7415130a0d612b221bc8a37b91c29caafe5e0b20c515d317ffb308" + }, + "NOTICE": { + "size": 2741, + "sha256": "00624a5043e7ce74029317b024132ca5286116d6fbf3191d226374bc8273789f" + }, + "fermion_container.py": { + "size": 3915, + "sha256": "5cf172e8cf6b313e951c4ef588ca913c7285b0aa3b92de413c3d7ec2b939c0e4" + }, + "reference_transformers.py": { + "size": 6823, + "sha256": "8481cfc11537e6dc1e40a8296c93aa15f9b7437b1514564caf647f8773c0dcb3" + }, + "LICENSE-WEIGHTS-CC-BY-4.0.txt": { + "size": 18657, + "sha256": "9ba9550ad48438d0836ddab3da480b3b69ffa0aac7b7878b5a0039e7ab429411" + } + }, + "archive_members": { + "model.fermion": { + "size": 177438361, + "sha256": "4b6bfa3a12cc3c4e0a54f2ab3ec4ca7a842b09e5c7ecfc8e7ca0ac6cc8c11468", + "format": "fermion-five-value-parakeet-v1" + }, + "config.json": { + "size": 277493, + "sha256": "d0daad3b2a182893844f4abdc11e4f5b7083f7d42c8ad7e8203f71559785a31b", + "note": "full NeMo config; 'labels' holds the 8192 SentencePiece pieces" + }, + "packed_manifest.json": { + "size": 821, + "sha256": "3c1874501a1efb6ef569eab082aa68d91b93ec6e4d2fb00dbef7851edf2ae510" + } + }, + "reference_env": { + "note": "Environment used to produce tests/fixtures from reference_transformers.py (dense fp32 ParakeetForTDT) and the base repo's processor (nvidia/parakeet-tdt-0.6b-v3).", + "python": "3.12", + "torch": "2.14.1", + "transformers": "5.17.0", + "numpy": "2.5.4", + "librosa": "1.0.0" + }, + "gguf": { + "file": "phonon-2-fv5.gguf", + "size": 330632256, + "sha256": "dc552eb786486167dde3ac80bbe49a6088c8ef39c03c257d5a090fc57b2c47b3", + "format": "fv5-pair16-v1", + "note": "convert_phonon2_to_gguf.py --archive output; byte-reproducible. Changes whenever the converter changes, because its git blob sha is in the metadata." + } +} diff --git a/tools/parakeet/src/parakeet.cpp b/tools/parakeet/src/parakeet.cpp new file mode 100644 index 000000000..d63a257c0 --- /dev/null +++ b/tools/parakeet/src/parakeet.cpp @@ -0,0 +1,1301 @@ +// tools/parakeet/src/parakeet.cpp - Parakeet-TDT (Phonon-2) loader, log-mel +// front end, FastConformer encoder, greedy TDT decoder and the batch API. +// +// Numerics follow transformers' ParakeetFeatureExtractor / ParakeetEncoder / +// ParakeetForTDT (the reference Fermion ships). Differences from the dense +// fp32 reference: activations entering the five-value and int6 linears are +// quantised to int8 per token, the five-value levels are put on a per-row +// int8 grid (max relative error reported as fv5_grid_err), and the decoder's +// LSTM / joint matvecs use int8 activations. tests/ gates the effect. +#include "parakeet.h" + +#include "pk-model.h" + +#include "ggml.h" +#include "gguf.h" + +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#if !defined(_WIN32) +#include +#include +#include +#include +#endif + +#if (defined(__ARM_NEON) || defined(__aarch64__)) && !defined(_MSC_VER) +#include +#define PK_NEON 1 +#endif + +using clk = std::chrono::steady_clock; + +static constexpr double PK_PI = 3.14159265358979323846; + +// PK_PROFILE=1: per-stage wall time of the encoder, printed at exit +namespace { +struct prof_t { + bool on = std::getenv("PK_PROFILE") != nullptr; + double t[16] = {}; + const char * name[16] = { "subsample", "layernorm", "quantize", "ff gemm", "qkv gemm", "pos cache", "attention", + "out gemm", "pw1 gemm", "glu+dwconv", "pw2 gemm", "enc_proj", "", "", "", "" }; + ~prof_t() { + if (!on) return; + double tot = 0; + for (double v : t) tot += v; + for (int i = 0; i < 12; ++i) std::fprintf(stderr, "[pk-profile] %-12s %9.1f ms %5.1f%%\n", name[i], t[i], 100 * t[i] / (tot + 1e-9)); + } +} g_prof; +struct ptimer { + int k; + clk::time_point t0; + explicit ptimer(int k_) : k(k_), t0(g_prof.on ? clk::now() : clk::time_point()) {} + ~ptimer() { if (g_prof.on) g_prof.t[k] += std::chrono::duration(clk::now() - t0).count(); } +}; +} // namespace +#define PKP(k) ptimer pk_pt_##__LINE__(k) + +static double ms_since(clk::time_point t0) { + return std::chrono::duration(clk::now() - t0).count(); +} + +namespace pk { + +// ------------------------------------------------------------- file map + +mapped_file::~mapped_file() { +#if !defined(_WIN32) + if (addr) { + munmap(addr, size); + } +#endif +} + +bool mapped_file::open(const std::string & path, bool use_mmap, std::string & err) { +#if !defined(_WIN32) + if (use_mmap) { + const int fd = ::open(path.c_str(), O_RDONLY); + if (fd < 0) { + err = "cannot open " + path; + return false; + } + struct stat st {}; + if (fstat(fd, &st) != 0) { + ::close(fd); + err = "cannot stat " + path; + return false; + } + size = (size_t) st.st_size; + void * p = mmap(nullptr, size, PROT_READ, MAP_PRIVATE, fd, 0); + ::close(fd); + if (p == MAP_FAILED) { + err = "mmap failed for " + path; + size = 0; + return false; + } + addr = p; + return true; + } +#else + (void) use_mmap; +#endif + FILE * f = std::fopen(path.c_str(), "rb"); + if (!f) { + err = "cannot open " + path; + return false; + } + std::fseek(f, 0, SEEK_END); + size = (size_t) std::ftell(f); + std::fseek(f, 0, SEEK_SET); + heap.resize(size); + const bool ok = std::fread(heap.data(), 1, size, f) == size; + std::fclose(f); + if (!ok) { + err = "short read on " + path; + } + return ok; +} + +int subsampled_length(int n) { + for (int i = 0; i < 3; ++i) { + n = (n + 2 - 3) / 2 + 1; // kernel 3, stride 2, padding 1 + if (n < 0) n = 0; + } + return n; +} + +} // namespace pk + +using namespace pk; + +// --------------------------------------------------------------- loading + +namespace { + +struct loader { + pk_model * m; + std::string err; + + const ggml_tensor * meta_of(const std::string & name) { + return ggml_get_tensor(m->meta, name.c_str()); + } + + const void * data(const std::string & name, ggml_type type, std::initializer_list ne) { + if (!err.empty()) { + return nullptr; + } + const int64_t id = gguf_find_tensor(m->gg, name.c_str()); + const ggml_tensor * t = meta_of(name); + if (id < 0 || !t) { + err = "missing tensor " + name; + return nullptr; + } + if (t->type != type) { + err = "tensor " + name + ": unexpected type"; + return nullptr; + } + int i = 0; + for (int64_t d : ne) { + if (t->ne[i] != d) { + char buf[256]; + std::snprintf(buf, sizeof(buf), "tensor %s: ne[%d] = %lld, expected %lld", name.c_str(), i, + (long long) t->ne[i], (long long) d); + err = buf; + return nullptr; + } + ++i; + } + const size_t off = gguf_get_data_offset(m->gg) + gguf_get_tensor_offset(m->gg, id); + if (off + ggml_nbytes(t) > m->file.size) { + err = "tensor " + name + " runs past the end of the file"; + return nullptr; + } + return m->file.data() + off; + } + const float * f32(const std::string & n, std::initializer_list ne) { + return (const float *) data(n, GGML_TYPE_F32, ne); + } + const int8_t * i8(const std::string & n, std::initializer_list ne) { + return (const int8_t *) data(n, GGML_TYPE_I8, ne); + } + + uint32_t u32(const char * key, uint32_t def) { + const std::string k = std::string("parakeet.") + key; + const int64_t id = gguf_find_key(m->gg, k.c_str()); + if (id < 0) return def; + if (gguf_get_kv_type(m->gg, id) != GGUF_TYPE_UINT32) { + if (err.empty()) err = k + ": expected uint32"; + return def; + } + return gguf_get_val_u32(m->gg, id); + } + float f32kv(const char * key, float def) { + const std::string k = std::string("parakeet.") + key; + const int64_t id = gguf_find_key(m->gg, k.c_str()); + if (id < 0) return def; + if (gguf_get_kv_type(m->gg, id) != GGUF_TYPE_FLOAT32) { + if (err.empty()) err = k + ": expected float32"; + return def; + } + return gguf_get_val_f32(m->gg, id); + } + + void fv5(wmat & w, const std::string & name, int O, int K) { + const int8_t * codes = i8(name + ".fv5", { K, O / 2 }); + const float * lohi = f32(name + ".lohi", { 2, O }); + if (!err.empty()) { + return; + } + if (O % 8 || K % 16) { + err = name + ": rows must be a multiple of 8 and width a multiple of 16"; + return; + } + wmat_from_fv5(w, (const uint8_t *) codes, lohi, O, K); + m->fv5_grid_err = std::max(m->fv5_grid_err, w.grid_err); + } + void intw(wmat & w, const std::string & name, int O, int K, const float * bias, const int * perm = nullptr) { + const int8_t * q = i8(name + ".q", { K, O }); + const float * s = f32(name + ".scale", { O }); + if (!err.empty()) { + return; + } + wmat_from_int8(w, q, s, O, K, perm); + w.bias = bias; + } +}; + +void setup_frontend(pk_model * m) { + const auto & hp = m->hp; + const int nb = hp.n_fft / 2 + 1; + m->mel_lo.assign(hp.n_mels, 0); + m->mel_hi.assign(hp.n_mels, 0); + for (int i = 0; i < hp.n_mels; ++i) { + int lo = nb, hi = 0; + for (int k = 0; k < nb; ++k) { + if (m->mel_fb[(size_t) i * nb + k] != 0.0f) { + lo = std::min(lo, k); + hi = k + 1; + } + } + m->mel_lo[i] = lo < hi ? lo : 0; + m->mel_hi[i] = lo < hi ? hi : 0; + } + // torch.hann_window(win, periodic=False), zero-padded to n_fft and centred + m->window.assign(hp.n_fft, 0.0f); + const int off = (hp.n_fft - hp.win_length) / 2; + for (int i = 0; i < hp.win_length; ++i) { + m->window[off + i] = (float) (0.5 - 0.5 * std::cos(2.0 * PK_PI * i / (hp.win_length - 1))); + } +} + +} // namespace + +extern "C" pk_params pk_default_params(void) { + pk_params p; + p.n_threads = 0; + p.kernel = 0; + p.pin_threads = 1; + p.use_mmap = 1; + return p; +} + +extern "C" int pk_gguf_is_parakeet(const char * path) { + gguf_init_params ip = { /*no_alloc*/ true, /*ctx*/ nullptr }; + gguf_context * g = gguf_init_from_file(path, ip); + if (!g) { + return 0; + } + const int64_t id = gguf_find_key(g, "general.architecture"); + const int ok = id >= 0 && gguf_get_kv_type(g, id) == GGUF_TYPE_STRING && + std::strcmp(gguf_get_val_str(g, id), "parakeet-tdt") == 0; + gguf_free(g); + return ok; +} + +static void set_err(char * err, size_t n, const std::string & s) { + if (err && n > 0) { + std::snprintf(err, n, "%s", s.c_str()); + } +} + +extern "C" pk_model * pk_model_load(const char * path, pk_params params, char * err, size_t err_len) { + auto m = std::make_unique(); + std::string e; + if (!m->file.open(path, params.use_mmap != 0, e)) { + set_err(err, err_len, e); + return nullptr; + } + gguf_init_params ip = { /*no_alloc*/ true, /*ctx*/ &m->meta }; + m->gg = gguf_init_from_file(path, ip); + if (!m->gg) { + set_err(err, err_len, std::string("not a GGUF file: ") + path); + return nullptr; + } + { + // gguf_get_val_* abort on a type mismatch, so check the type first + auto str_kv = [&](const char * key) -> const char * { + const int64_t kid = gguf_find_key(m->gg, key); + return kid >= 0 && gguf_get_kv_type(m->gg, kid) == GGUF_TYPE_STRING ? gguf_get_val_str(m->gg, kid) : nullptr; + }; + const char * arch = str_kv("general.architecture"); + if (!arch || std::strcmp(arch, "parakeet-tdt") != 0) { + set_err(err, err_len, "general.architecture is not parakeet-tdt"); + return nullptr; + } + const char * fmt = str_kv("parakeet.weights.format"); + if (!fmt || std::strcmp(fmt, "fv5-pair16-v1") != 0) { + set_err(err, err_len, "parakeet.weights.format is not fv5-pair16-v1 (re-run convert_phonon2_to_gguf.py)"); + return nullptr; + } + if (const char * notice = str_kv("parakeet.notice")) { + m->notice = notice; + } + } + loader L{ m.get(), {} }; + auto & hp = m->hp; + hp.d_model = (int) L.u32("d_model", 1024); + hp.n_layers = (int) L.u32("n_layers", 24); + hp.n_heads = (int) L.u32("n_heads", 8); + hp.ff_dim = (int) L.u32("ff_dim", 4096); + hp.conv_kernel = (int) L.u32("conv_kernel", 9); + hp.sub_channels = (int) L.u32("sub_channels", 256); + hp.sub_factor = (int) L.u32("sub_factor", 8); + hp.n_mels = (int) L.u32("n_mels", 128); + hp.sample_rate = (int) L.u32("sample_rate", 16000); + hp.n_fft = (int) L.u32("n_fft", 512); + hp.win_length = (int) L.u32("win_length", 400); + hp.hop_length = (int) L.u32("hop_length", 160); + hp.pred_hidden = (int) L.u32("pred_hidden", 640); + hp.pred_layers = (int) L.u32("pred_layers", 2); + hp.joint_hidden = (int) L.u32("joint_hidden", 640); + hp.vocab = (int) L.u32("vocab", 8192); + hp.blank_id = (int) L.u32("blank_id", 8192); + hp.max_symbols = (int) L.u32("max_symbols", 10); + hp.preemph = L.f32kv("preemph", 0.97f); + hp.log_guard = L.f32kv("log_guard", 5.9604644775390625e-8f); + hp.norm_eps = L.f32kv("norm_eps", 1e-5f); + hp.ln_eps = L.f32kv("ln_eps", 1e-5f); + if (!L.err.empty()) { + set_err(err, err_len, L.err); + return nullptr; + } + if (hp.sub_factor != 8 || hp.pred_layers != 2 || hp.n_heads <= 0 || hp.d_model % hp.n_heads || + hp.d_model / hp.n_heads != 128 || hp.n_mels != 128 || hp.n_fft != 512 || hp.n_layers <= 0 || + hp.vocab <= 0 || hp.blank_id != hp.vocab || hp.max_symbols <= 0 || hp.hop_length <= 0 || + hp.win_length <= 1 || hp.win_length > hp.n_fft || hp.sample_rate <= 0 || hp.conv_kernel <= 0 || + hp.conv_kernel % 2 == 0) { + set_err(err, err_len, "unsupported parakeet-tdt hyper-parameters"); + return nullptr; + } + { + const int64_t id = gguf_find_key(m->gg, "parakeet.durations"); + if (id >= 0) { + if (gguf_get_kv_type(m->gg, id) != GGUF_TYPE_ARRAY || gguf_get_arr_type(m->gg, id) != GGUF_TYPE_INT32) { + set_err(err, err_len, "parakeet.durations: expected an int32 array"); + return nullptr; + } + const size_t n = gguf_get_arr_n(m->gg, id); + hp.durations.clear(); + for (size_t i = 0; i < n; ++i) { + hp.durations.push_back(((const int32_t *) gguf_get_arr_data(m->gg, id))[i]); + } + } + // the greedy decoder indexes durations and advances by them: a + // negative duration would loop forever, an empty table read past it + bool ok = !hp.durations.empty() && hp.durations.size() <= 64; + for (int d : hp.durations) ok = ok && d >= 0 && d <= 64; + if (!ok) { + set_err(err, err_len, "parakeet.durations missing or out of range"); + return nullptr; + } + } + { + const int64_t id = gguf_find_key(m->gg, "tokenizer.ggml.tokens"); + if (id < 0 || gguf_get_kv_type(m->gg, id) != GGUF_TYPE_ARRAY || + gguf_get_arr_type(m->gg, id) != GGUF_TYPE_STRING || (int) gguf_get_arr_n(m->gg, id) != hp.vocab) { + set_err(err, err_len, "tokenizer.ggml.tokens missing or wrong size"); + return nullptr; + } + m->pieces.resize(hp.vocab); + m->is_control.assign(hp.vocab, 0); + m->word_start.assign(hp.vocab, 0); + for (int i = 0; i < hp.vocab; ++i) { + std::string s = gguf_get_arr_str(m->gg, id, i); + const bool ctrl = s.size() >= 2 && s.front() == '<' && s.back() == '>' && s.find("\xe2\x96\x81") == std::string::npos; + m->is_control[i] = ctrl; + std::string o; + for (size_t k = 0; k < s.size();) { + if (s.compare(k, 3, "\xe2\x96\x81") == 0) { + o += ' '; + k += 3; + } else { + o += s[k++]; + } + } + m->word_start[i] = !o.empty() && o[0] == ' '; + m->pieces[i] = o; + } + } + + const int D = hp.d_model, C = hp.sub_channels, F = hp.ff_dim, H = hp.pred_hidden, J = hp.joint_hidden; + const int nb = hp.n_fft / 2 + 1; + m->mel_fb = L.f32("feat.mel_fb", { nb, hp.n_mels }); + + auto transpose9 = [&](const float * w, std::vector & out) { + out.assign((size_t) 9 * C, 0.0f); + if (!w) return; + for (int c = 0; c < C; ++c) { + for (int k = 0; k < 9; ++k) { + out[(size_t) k * C + c] = w[(size_t) c * 9 + k]; + } + } + }; + transpose9(L.f32("enc.sub.conv0.w", { 9, C }), m->c0_wT); + m->c0_b = L.f32("enc.sub.conv0.b", { C }); + transpose9(L.f32("enc.sub.dw1.w", { 9, C }), m->dw1_wT); + m->dw1_b = L.f32("enc.sub.dw1.b", { C }); + transpose9(L.f32("enc.sub.dw2.w", { 9, C }), m->dw2_wT); + m->dw2_b = L.f32("enc.sub.dw2.b", { C }); + L.intw(m->sub_pw1, "enc.sub.pw1", C, C, L.f32("enc.sub.pw1.b", { C })); + L.intw(m->sub_pw2, "enc.sub.pw2", C, C, L.f32("enc.sub.pw2.b", { C })); + { + // torch flattens (C, F') -> c * F' + f; the encoder keeps channel-last + // rows (f * C + c), so the linear's columns are permuted at load. + const int Fq = hp.n_mels / 8; + std::vector perm((size_t) Fq * C); + for (int f = 0; f < Fq; ++f) { + for (int c = 0; c < C; ++c) { + perm[(size_t) f * C + c] = c * Fq + f; + } + } + L.intw(m->sub_out, "enc.sub.out", D, C * Fq, L.f32("enc.sub.out.b", { D }), perm.data()); + } + + m->layers.resize(hp.n_layers); + static const char * ln_names[5] = { "norm_feed_forward1", "norm_self_att", "norm_conv", "norm_feed_forward2", + "norm_out" }; + for (int l = 0; l < hp.n_layers && L.err.empty(); ++l) { + auto & ly = m->layers[l]; + const std::string p = "enc.l" + std::to_string(l) + "."; + L.fv5(ly.ff1_w1, p + "ff1.w1", F, D); + L.fv5(ly.ff1_w2, p + "ff1.w2", D, F); + L.fv5(ly.qkv, p + "attn.qkv", 3 * D, D); + L.fv5(ly.out, p + "attn.out", D, D); + L.fv5(ly.pos, p + "attn.pos", D, D); + L.fv5(ly.pw1, p + "conv.pw1", 2 * D, D); + L.fv5(ly.pw2, p + "conv.pw2", D, D); + L.fv5(ly.ff2_w1, p + "ff2.w1", F, D); + L.fv5(ly.ff2_w2, p + "ff2.w2", D, F); + ly.bias_u = L.f32(p + "attn.bias_u", { 128, hp.n_heads }); + ly.bias_v = L.f32(p + "attn.bias_v", { 128, hp.n_heads }); + for (int i = 0; i < 5; ++i) { + ly.ln_w[i] = L.f32(p + ln_names[i] + ".w", { D }); + ly.ln_b[i] = L.f32(p + ln_names[i] + ".b", { D }); + } + ly.dw_w = L.f32(p + "conv.dw.w", { hp.conv_kernel, D }); + ly.dw_b = L.f32(p + "conv.dw.b", { D }); + if (ly.dw_w) { + ly.dw_wT.resize((size_t) hp.conv_kernel * D); + for (int c = 0; c < D; ++c) { + for (int k = 0; k < hp.conv_kernel; ++k) { + ly.dw_wT[(size_t) k * D + c] = ly.dw_w[(size_t) c * hp.conv_kernel + k]; + } + } + } + } + L.intw(m->enc_proj, "enc_proj", J, D, L.f32("enc_proj.b", { J })); + m->embed_q = L.i8("dec.embed.q", { H, hp.vocab + 1 }); + m->embed_s = L.f32("dec.embed.scale", { hp.vocab + 1 }); + for (int l = 0; l < 2; ++l) { + const std::string p = "dec.lstm" + std::to_string(l); + m->lstm_ih[l] = L.i8(p + ".ih.q", { H, 4 * H }); + m->lstm_ih_s[l] = L.f32(p + ".ih.scale", { 4 * H }); + m->lstm_hh[l] = L.i8(p + ".hh.q", { H, 4 * H }); + m->lstm_hh_s[l] = L.f32(p + ".hh.scale", { 4 * H }); + m->lstm_b[l] = L.f32(p + ".b", { 4 * H }); + } + m->proj_q = L.i8("dec.proj.q", { H, J }); + m->proj_s = L.f32("dec.proj.scale", { J }); + m->proj_b = L.f32("dec.proj.b", { J }); + m->head_rows = hp.vocab + 1 + (int) hp.durations.size(); + m->head_q = L.i8("joint.head.q", { J, m->head_rows }); + m->head_s = L.f32("joint.head.scale", { m->head_rows }); + m->head_b = L.f32("joint.head.b", { m->head_rows }); + if (!L.err.empty()) { + set_err(err, err_len, L.err); + return nullptr; + } + if (H % 16 || J % 16 || D % 16) { + set_err(err, err_len, "decoder widths must be multiples of 16"); + return nullptr; + } + setup_frontend(m.get()); + + m->kt = &kernels_for(params.kernel); + const int nt = params.n_threads > 0 ? params.n_threads : default_threads(); + m->pool = std::make_unique(nt, params.pin_threads != 0); + return m.release(); +} + +extern "C" void pk_model_free(pk_model * m) { + if (!m) { + return; + } + if (m->gg) { + gguf_free(m->gg); + } + if (m->meta) { + ggml_free(m->meta); + } + delete m; +} + +extern "C" const char * pk_model_kernel_name(const pk_model * m) { return m ? m->kt->name : ""; } +extern "C" int pk_model_n_threads(const pk_model * m) { return m ? m->pool->size() : 0; } +extern "C" int pk_model_sample_rate(const pk_model * m) { return m ? m->hp.sample_rate : 0; } +extern "C" const char * pk_model_notice(const pk_model * m) { return m ? m->notice.c_str() : ""; } + +// ------------------------------------------------------------ front end + +namespace { + +// in-place radix-2 FFT of size n (power of two), double precision +void fft(std::vector> & a, const std::vector> & tw) { + const size_t n = a.size(); + for (size_t i = 1, j = 0; i < n; ++i) { + size_t bit = n >> 1; + for (; j & bit; bit >>= 1) { + j ^= bit; + } + j ^= bit; + if (i < j) { + std::swap(a[i], a[j]); + } + } + for (size_t len = 2; len <= n; len <<= 1) { + const size_t step = n / len; + for (size_t i = 0; i < n; i += len) { + for (size_t k = 0; k < len / 2; ++k) { + const std::complex u = a[i + k]; + const std::complex v = a[i + k + len / 2] * tw[k * step]; + a[i + k] = u + v; + a[i + k + len / 2] = u - v; + } + } + } +} + +} // namespace + +bool pk::compute_mel(pk_model * m, const float * pcm, size_t n, std::vector & mel, int & n_frames) { + const auto & hp = m->hp; + const int hop = hp.hop_length, nfft = hp.n_fft, nm = hp.n_mels, nb = nfft / 2 + 1; + n_frames = (int) (n / (size_t) hop); + if (n_frames < 2) { + mel.clear(); + n_frames = 0; + return false; + } + std::vector y(n); + y[0] = pcm[0]; + for (size_t i = 1; i < n; ++i) { + y[i] = pcm[i] - hp.preemph * pcm[i - 1]; + } + mel.assign((size_t) n_frames * nm, 0.0f); + std::vector> tw(nfft / 2); + for (int k = 0; k < nfft / 2; ++k) { + tw[k] = std::polar(1.0, -2.0 * PK_PI * k / nfft); + } + auto frames = [&](int b, int e) { + std::vector> buf(nfft); + std::vector pw(nb); + for (int f = b; f < e; ++f) { + const long c = (long) f * hop - nfft / 2; + for (int i = 0; i < nfft; ++i) { + const long s = c + i; + const float v = (s >= 0 && s < (long) n) ? y[s] : 0.0f; + buf[i] = std::complex((double) v * (double) m->window[i], 0.0); + } + fft(buf, tw); + for (int k = 0; k < nb; ++k) { + const float re = (float) buf[k].real(), im = (float) buf[k].imag(); + pw[k] = re * re + im * im; + } + float * out = mel.data() + (size_t) f * nm; + for (int i = 0; i < nm; ++i) { + const float * fb = m->mel_fb + (size_t) i * nb; + float s = 0.0f; + for (int k = m->mel_lo[i]; k < m->mel_hi[i]; ++k) { + s += fb[k] * pw[k]; + } + out[i] = std::log(s + hp.log_guard); + } + } + }; + m->pool->parallel_for(n_frames, 64, frames); + // per-feature normalisation over the valid frames (unbiased variance) + for (int i = 0; i < nm; ++i) { + double mean = 0.0; + for (int f = 0; f < n_frames; ++f) { + mean += mel[(size_t) f * nm + i]; + } + mean /= n_frames; + double var = 0.0; + for (int f = 0; f < n_frames; ++f) { + const double d = mel[(size_t) f * nm + i] - mean; + var += d * d; + } + var /= (n_frames - 1); + const float sd = (float) std::sqrt(var) + hp.norm_eps; + const float mf = (float) mean; + for (int f = 0; f < n_frames; ++f) { + float & v = mel[(size_t) f * nm + i]; + v = (v - mf) / sd; + } + } + return true; +} + +// -------------------------------------------------------------- encoder + +namespace { + +// conv0 (1 -> C, 3x3, stride 2, pad 1) + ReLU for output time rows [t0, t1), +// channel-last output rows [t][F1][C] +void conv0_rows(const pk_model * m, const float * mel, int Tm, int t0, int t1, float * out) { + const int C = m->hp.sub_channels, Fin = m->hp.n_mels, F1 = Fin / 2; + for (int t = t0; t < t1; ++t) { + for (int f = 0; f < F1; ++f) { + float * o = out + ((size_t) (t - t0) * F1 + f) * C; + std::memcpy(o, m->c0_b, sizeof(float) * C); + for (int dt = 0; dt < 3; ++dt) { + const int ti = 2 * t + dt - 1; + if (ti < 0 || ti >= Tm) continue; + for (int df = 0; df < 3; ++df) { + const int fi = 2 * f + df - 1; + if (fi < 0 || fi >= Fin) continue; + const float v = mel[(size_t) ti * Fin + fi]; + const float * w = m->c0_wT.data() + (size_t) (dt * 3 + df) * C; + int c = 0; +#if PK_NEON + const float32x4_t v4 = vdupq_n_f32(v); + for (; c + 4 <= C; c += 4) { + vst1q_f32(o + c, vfmaq_f32(vld1q_f32(o + c), v4, vld1q_f32(w + c))); + } +#endif + for (; c < C; ++c) o[c] += v * w[c]; + } + } + int c = 0; +#if PK_NEON + for (; c + 4 <= C; c += 4) vst1q_f32(o + c, vmaxq_f32(vld1q_f32(o + c), vdupq_n_f32(0.0f))); +#endif + for (; c < C; ++c) o[c] = o[c] > 0.0f ? o[c] : 0.0f; + } + } +} + +// depthwise 3x3 stride 2 pad 1, channel-last. `in` holds rows [in_t0, ...) +// of a [Tin][Fin][C] map; rows outside [0, Tin) are zero padding. +void dw_rows(const float * in, int in_t0, int in_t1, int Tin, int Fin, int C, const float * wT, const float * b, + int t0, int t1, float * out) { + const int Fo = (Fin + 2 - 3) / 2 + 1; + for (int t = t0; t < t1; ++t) { + for (int f = 0; f < Fo; ++f) { + float * o = out + ((size_t) (t - t0) * Fo + f) * C; + std::memcpy(o, b, sizeof(float) * C); + for (int dt = 0; dt < 3; ++dt) { + const int ti = 2 * t + dt - 1; + if (ti < 0 || ti >= Tin) continue; + if (ti < in_t0 || ti >= in_t1) continue; // never hit: callers pass the full halo + for (int df = 0; df < 3; ++df) { + const int fi = 2 * f + df - 1; + if (fi < 0 || fi >= Fin) continue; + const float * x = in + ((size_t) (ti - in_t0) * Fin + fi) * C; + const float * w = wT + (size_t) (dt * 3 + df) * C; + int c = 0; +#if PK_NEON + for (; c + 4 <= C; c += 4) { + vst1q_f32(o + c, vfmaq_f32(vld1q_f32(o + c), vld1q_f32(x + c), vld1q_f32(w + c))); + } +#endif + for (; c < C; ++c) o[c] += x[c] * w[c]; + } + } + } + } +} + +void subsample(pk_model * m, const float * mel, int Tm, std::vector & x, int & T3) { + const int C = m->hp.sub_channels; + const int T1 = (Tm - 1) / 2 + 1, T2 = (T1 - 1) / 2 + 1; + T3 = (T2 - 1) / 2 + 1; + const int F1 = m->hp.n_mels / 2, F2 = F1 / 2, F3 = F2 / 2; + const int D = m->hp.d_model; + auto & P = *m->pool; + auto & ws = m->ws; + std::vector s3((size_t) T3 * F3 * C); + const int CH = 16; // T3 rows per chunk (1.28 s) + for (int a3 = 0; a3 < T3; a3 += CH) { + const int b3 = std::min(T3, a3 + CH); + const int a2 = std::max(0, 2 * a3 - 1), b2 = std::min(T2, 2 * b3); + const int a1 = std::max(0, 2 * a2 - 1), b1 = std::min(T1, 2 * b2); + ws.sub_a.resize((size_t) (b1 - a1) * F1 * C); + P.parallel_for(b1 - a1, 4, [&](int b, int e) { conv0_rows(m, mel, Tm, a1 + b, a1 + e, ws.sub_a.data() + (size_t) b * F1 * C); }); + ws.sub_b.resize((size_t) (b2 - a2) * F2 * C); + P.parallel_for(b2 - a2, 2, [&](int b, int e) { + dw_rows(ws.sub_a.data(), a1, b1, T1, F1, C, m->dw1_wT.data(), m->dw1_b, a2 + b, a2 + e, + ws.sub_b.data() + (size_t) b * F2 * C); + }); + // pointwise 1x1 + ReLU over every (t, f) position + const int n2 = (b2 - a2) * F2; + quantize_acts(P, ws.sub_b.data(), n2, C, C, ws.qa); + ws.sub_c.resize((size_t) n2 * C); + gemm_q8(P, *m->kt, m->sub_pw1, ws.qa, ws.sub_c.data(), C, epi::relu); + float * o3 = s3.data() + (size_t) a3 * F3 * C; + P.parallel_for(b3 - a3, 2, [&](int b, int e) { + dw_rows(ws.sub_c.data(), a2, b2, T2, F2, C, m->dw2_wT.data(), m->dw2_b, a3 + b, a3 + e, + o3 + (size_t) b * F3 * C); + }); + const int n3 = (b3 - a3) * F3; + quantize_acts(P, o3, n3, C, C, ws.qa); + gemm_q8(P, *m->kt, m->sub_pw2, ws.qa, o3, C, epi::relu); + } + x.resize((size_t) T3 * D); + quantize_acts(P, s3.data(), T3, F3 * C, (size_t) F3 * C, ws.qa); + gemm_q8(P, *m->kt, m->sub_out, ws.qa, x.data(), D, epi::store); +} + +inline void f16_to_f32_row(const uint16_t * src, float * dst, int n) { + int i = 0; +#if PK_NEON + for (; i + 8 <= n; i += 8) { + const float16x8_t h = vreinterpretq_f16_u16(vld1q_u16(src + i)); + vst1q_f32(dst + i, vcvt_f32_f16(vget_low_f16(h))); + vst1q_f32(dst + i + 4, vcvt_high_f32_f16(h)); + } +#endif + for (; i < n; ++i) dst[i] = ggml_fp16_to_fp32((ggml_fp16_t) src[i]); +} + +// sinusoidal relative-position embedding rows for positions p (the +// ParakeetEncoderRelPositionalEncoding values) +void rel_pos_rows(const std::vector & pos, int D, std::vector & pe) { + const int n = (int) pos.size(); + pe.resize((size_t) n * D); + for (int i = 0; i < D / 2; ++i) { + const float inv = (float) (1.0 / std::pow(10000.0, (2.0 * i) / D)); + for (int r = 0; r < n; ++r) { + const float a = inv * (float) pos[r]; + pe[(size_t) r * D + 2 * i] = std::sin(a); + pe[(size_t) r * D + 2 * i + 1] = std::cos(a); + } + } +} + +// make sure the per-layer position-key cache covers |p| <= need. +// Cache rows c = 0..2P hold position p = P - c. Every row is computed on its +// own (per-row activation scales), so growing the cache only computes the +// new rows at both ends and keeps the old ones: the result is bit-identical +// to a full rebuild, and a stream whose window grows pays for each row once. +void ensure_rpos_cache(pk_model * m, int need) { + if (need <= m->rc_P) return; + const int P = ((std::max(need, 1) + 63) / 64) * 64; + const int D = m->hp.d_model, n = 2 * P + 1; + const int oldP = m->rc_P; // -1 when empty + const int n_top = oldP < 0 ? n : P - oldP; // positions P .. oldP+1 (or everything) + const int n_bot = oldP < 0 ? 0 : P - oldP; // positions -oldP-1 .. -P + std::vector pos; + pos.reserve((size_t) n_top + n_bot); + for (int i = 0; i < n_top; ++i) pos.push_back(P - i); + for (int i = 0; i < n_bot; ++i) pos.push_back(-oldP - 1 - i); + const int nn = (int) pos.size(); + std::vector pe, r((size_t) nn * D); + rel_pos_rows(pos, D, pe); + // Two-pass (residual) int8 quantisation: pe ~ s1*q1 + s2*q2, so the + // cached keys carry ~14-bit activations. The keys feed every attention + // logit, and each row is built once, so the second pass is cheap. + qact q, q2; + quantize_acts(*m->pool, pe.data(), nn, D, D, q); + std::vector res(pe.size()); + for (int t = 0; t < nn; ++t) { + for (int k = 0; k < D; ++k) { + const int c = k / 8, j = k % 8; + const int8_t v = q.buf[(size_t) (t / 2) * q.pair_stride() + (size_t) c * 16 + (t & 1) * 8 + j]; + res[(size_t) t * D + k] = pe[(size_t) t * D + k] - (float) v * q.scale[t]; + } + } + quantize_acts(*m->pool, res.data(), nn, D, D, q2); + // PK_DEBUG_F32_ACTS runs the first pass on the fp32 source already + const bool f32_debug = std::getenv("PK_DEBUG_F32_ACTS") != nullptr; + m->rcache.resize(m->hp.n_layers); + for (int l = 0; l < m->hp.n_layers; ++l) { + gemm_q8(*m->pool, *m->kt, m->layers[l].pos, q, r.data(), D, epi::store); + if (!f32_debug) { + gemm_q8(*m->pool, *m->kt, m->layers[l].pos, q2, r.data(), D, epi::add, 1.0f); + } + std::vector c((size_t) n * D); + // new rows: [0, n_top) are cache rows 0..n_top-1; [n_top, nn) are the + // bottom rows n - n_bot .. n - 1; the old cache sits in between + auto put = [&](int src_row, int dst_row, int count) { + m->pool->parallel_for(count, 16, [&](int b, int e) { + for (int i = b; i < e; ++i) { + const float * sr = r.data() + (size_t) (src_row + i) * D; + uint16_t * dr = c.data() + (size_t) (dst_row + i) * D; + for (int k = 0; k < D; ++k) dr[k] = (uint16_t) ggml_fp32_to_fp16(sr[k]); + } + }); + }; + put(0, 0, n_top); + if (oldP >= 0) { + const auto & old = m->rcache[l]; + std::memcpy(c.data() + (size_t) n_top * D, old.data(), sizeof(uint16_t) * old.size()); + put(n_top, n - n_bot, n_bot); + } + m->rcache[l].swap(c); + } + m->rc_P = P; +} + +// rel-pos multi-head self-attention (Transformer-XL), full context. +void attention(pk_model * m, const layer & ly, int l, int T) { + auto & P = *m->pool; + auto & ws = m->ws; + const int D = m->hp.d_model, NH = m->hp.n_heads, HD = D / NH; + const int Tp = (T + 3) & ~3; + const float scale = 1.0f / std::sqrt((float) HD); + // V^T per head: [NH][HD][Tp] + ws.v.assign((size_t) NH * HD * Tp, 0.0f); + P.parallel_for(NH, 1, [&](int b, int e) { + for (int h = b; h < e; ++h) { + float * vt = ws.v.data() + (size_t) h * HD * Tp; + for (int t = 0; t < T; ++t) { + const float * src = ws.qkv.data() + (size_t) t * 3 * D + 2 * D + h * HD; + for (int d = 0; d < HD; ++d) { + vt[(size_t) d * Tp + t] = src[d]; + } + } + } + }); + const int QB = 16; + const int nqb = (T + QB - 1) / QB; + if ((int) ws.att_tls.size() < P.size()) { + ws.att_tls.resize(P.size()); + } + ws.att.resize((size_t) T * D); + std::atomic next{ 0 }; + P.run([&](int tid) { + auto & s = ws.att_tls[tid]; + const size_t need = (size_t) QB * HD * 2 + (size_t) QB * Tp + (size_t) QB * (T + QB) + (size_t) (T + QB) * HD; + if (s.size() < need) s.resize(need); + float * qu = s.data(); + float * qv = qu + QB * HD; + float * S = qv + QB * HD; + float * band = S + (size_t) QB * Tp; + float * rk = band + (size_t) QB * (T + QB); + const uint16_t * rc = m->rcache[l].data(); + const int rP = m->rc_P; + for (;;) { + const int task = next.fetch_add(1); + if (task >= NH * nqb) break; + const int h = task / nqb, qb = task % nqb; + const int i0 = qb * QB, B = std::min(QB, T - i0); + const float * bu = ly.bias_u + h * HD; + const float * bv = ly.bias_v + h * HD; + for (int i = 0; i < B; ++i) { + const float * q = ws.qkv.data() + (size_t) (i0 + i) * 3 * D + h * HD; + for (int d = 0; d < HD; ++d) { + qu[i * HD + d] = q[d] + bu[d]; + qv[i * HD + d] = q[d] + bv[d]; + } + } + // content term (q + u) . k_j + gemm_f32_nt(qu, HD, ws.qkv.data() + D + h * HD, (size_t) 3 * D, S, Tp, B, T, HD, 1.0f); + // position term (q + v) . r_{i-j}: rows T-1-i+j of R, a band + const int m0 = T - 1 - (i0 + B - 1); + const int W = T + B - 1; + // R row m (position T-1-m) is cache row m + rP - T + 1 + for (int w = 0; w < W; ++w) { + f16_to_f32_row(rc + (size_t) (m0 + w + rP - T + 1) * D + h * HD, rk + (size_t) w * HD, HD); + } + gemm_f32_nt(qv, HD, rk, HD, band, W, B, W, HD, 1.0f); + for (int i = 0; i < B; ++i) { + float * srow = S + (size_t) i * Tp; + const float * brow = band + (size_t) i * W + (B - 1 - i); + float mx = -INFINITY; + for (int j = 0; j < T; ++j) { + srow[j] = (srow[j] + brow[j]) * scale; + mx = std::max(mx, srow[j]); + } + for (int j = 0; j < T; ++j) srow[j] -= mx; + exp_inplace(srow, T); + float sum = 0.0f; + for (int j = 0; j < T; ++j) sum += srow[j]; + const float inv = 1.0f / sum; + for (int j = 0; j < T; ++j) srow[j] *= inv; + for (int j = T; j < Tp; ++j) srow[j] = 0.0f; + } + gemm_f32_nt(S, Tp, ws.v.data() + (size_t) h * HD * Tp, Tp, ws.att.data() + (size_t) i0 * D + h * HD, D, B, + HD, Tp, 1.0f); + } + }); +} + +void conv_module(pk_model * m, const layer & ly, int T) { + auto & P = *m->pool; + auto & ws = m->ws; + const int D = m->hp.d_model, K = m->hp.conv_kernel, pad = (K - 1) / 2; + // GLU: u = a * sigmoid(b) + ws.u.resize((size_t) T * D); + P.parallel_for(T, 8, [&](int b, int e) { + std::vector tmp(D); + for (int t = b; t < e; ++t) { + const float * g = ws.g.data() + (size_t) t * 2 * D; + for (int c = 0; c < D; ++c) tmp[c] = -g[D + c]; + exp_inplace(tmp.data(), D); + float * u = ws.u.data() + (size_t) t * D; + for (int c = 0; c < D; ++c) u[c] = g[c] / (1.0f + tmp[c]); + } + }); + // depthwise conv over time (BatchNorm folded) + SiLU + ws.xn.resize((size_t) T * D); + P.parallel_for(T, 8, [&](int b, int e) { + for (int t = b; t < e; ++t) { + float * o = ws.xn.data() + (size_t) t * D; + std::memcpy(o, ly.dw_b, sizeof(float) * D); + for (int k = 0; k < K; ++k) { + const int ti = t + k - pad; + if (ti < 0 || ti >= T) continue; + const float * x = ws.u.data() + (size_t) ti * D; + const float * w = ly.dw_wT.data() + (size_t) k * D; + int c = 0; +#if PK_NEON + for (; c + 4 <= D; c += 4) { + vst1q_f32(o + c, vfmaq_f32(vld1q_f32(o + c), vld1q_f32(w + c), vld1q_f32(x + c))); + } +#endif + for (; c < D; ++c) o[c] += w[c] * x[c]; + } + silu_inplace(o, D); + } + }); +} + +} // namespace + +int pk::encode(pk_model * m, const float * mel, int n_mel, std::vector & encp, pk_cancel_cb cancel, void * cu, + std::vector * enc_out, std::vector * sub_out, std::vector * l0_out) { + auto & P = *m->pool; + auto & ws = m->ws; + const auto & kt = *m->kt; + const int D = m->hp.d_model, F = m->hp.ff_dim; + int T = 0; + { PKP(0); subsample(m, mel, n_mel, ws.x, T); } + if (sub_out) *sub_out = ws.x; + if (T <= 0) return 0; + { PKP(5); ensure_rpos_cache(m, T - 1); } + ws.xn.resize((size_t) T * D); + ws.h.resize((size_t) T * F); + ws.qkv.resize((size_t) T * 3 * D); + ws.g.resize((size_t) T * 2 * D); + float * X = ws.x.data(); + const float eps = m->hp.ln_eps; + auto ln = [&](const float * w, const float * b, float * out) { + PKP(1); + P.parallel_for(T, 8, [&](int bb, int ee) { + layernorm(X + (size_t) bb * D, out + (size_t) bb * D, ee - bb, D, w, b, eps); + }); + }; + auto quant = [&](const float * src, int K, qact & q) { + PKP(2); + quantize_acts(P, src, T, K, K, q); + }; + for (int l = 0; l < m->hp.n_layers; ++l) { + if (cancel && cancel(cu)) return -1; + const layer & ly = m->layers[l]; + // FF1 (macaron, x0.5) + ln(ly.ln_w[0], ly.ln_b[0], ws.xn.data()); + quant(ws.xn.data(), D, ws.qa); + { PKP(3); gemm_q8(P, kt, ly.ff1_w1, ws.qa, ws.h.data(), F, epi::silu); } + quant(ws.h.data(), F, ws.qh); + { PKP(3); gemm_q8(P, kt, ly.ff1_w2, ws.qh, X, D, epi::add, 0.5f); } + // self-attention + ln(ly.ln_w[1], ly.ln_b[1], ws.xn.data()); + quant(ws.xn.data(), D, ws.qa); + { PKP(4); gemm_q8(P, kt, ly.qkv, ws.qa, ws.qkv.data(), 3 * D, epi::store); } + { PKP(6); attention(m, ly, l, T); } + quant(ws.att.data(), D, ws.qa); + { PKP(7); gemm_q8(P, kt, ly.out, ws.qa, X, D, epi::add, 1.0f); } + // convolution module + ln(ly.ln_w[2], ly.ln_b[2], ws.xn.data()); + quant(ws.xn.data(), D, ws.qa); + { PKP(8); gemm_q8(P, kt, ly.pw1, ws.qa, ws.g.data(), 2 * D, epi::store); } + { PKP(9); conv_module(m, ly, T); } + quant(ws.xn.data(), D, ws.qa); + { PKP(10); gemm_q8(P, kt, ly.pw2, ws.qa, X, D, epi::add, 1.0f); } + // FF2 (x0.5) + ln(ly.ln_w[3], ly.ln_b[3], ws.xn.data()); + quant(ws.xn.data(), D, ws.qa); + { PKP(3); gemm_q8(P, kt, ly.ff2_w1, ws.qa, ws.h.data(), F, epi::silu); } + quant(ws.h.data(), F, ws.qh); + { PKP(3); gemm_q8(P, kt, ly.ff2_w2, ws.qh, X, D, epi::add, 0.5f); } + // final layer norm, in place + ln(ly.ln_w[4], ly.ln_b[4], X); + if (l == 0 && l0_out) *l0_out = ws.x; + } + if (enc_out) *enc_out = ws.x; + const int J = m->hp.joint_hidden; + encp.resize((size_t) T * J); + PKP(11); + quantize_acts(P, X, T, D, D, ws.qa); + gemm_q8(P, kt, m->enc_proj, ws.qa, encp.data(), J, epi::store); + return T; +} + +// -------------------------------------------------------------- decoder + +namespace { + +struct qvec { + std::vector q; + float s = 0.0f; +}; + +void quant_vec(const float * x, int n, qvec & out) { + out.q.resize(n); + float amax = 0.0f; + for (int i = 0; i < n; ++i) amax = std::max(amax, std::fabs(x[i])); + if (!(amax > 0.0f)) { + std::fill(out.q.begin(), out.q.end(), 0); + out.s = 0.0f; + return; + } + const float inv = 127.0f / amax; + for (int i = 0; i < n; ++i) { + out.q[i] = (int8_t) std::max(-127.0f, std::min(127.0f, std::nearbyint(x[i] * inv))); + } + out.s = amax / 127.0f; +} + +struct dec_state { + std::vector h[2], c[2], out; +}; + +inline float sigm(float x) { return 1.0f / (1.0f + std::exp(-x)); } + +void lstm_step(const pk_model * m, int token, dec_state & st, std::vector & a1, std::vector & a2, + std::vector & gates, qvec & xq, qvec & hq) { + const int H = m->hp.pred_hidden; + const auto & kt = *m->kt; + // layer 0 input: the int8 embedding row is its own quantisation + const int8_t * x8 = m->embed_q + (size_t) token * H; + float xs = m->embed_s[token]; + for (int l = 0; l < 2; ++l) { + if (l == 1) { + quant_vec(st.h[0].data(), H, xq); + x8 = xq.q.data(); + xs = xq.s; + } + quant_vec(st.h[l].data(), H, hq); + kt.matvec(m->lstm_ih[l], 4 * H, H, x8, a1.data()); + kt.matvec(m->lstm_hh[l], 4 * H, H, hq.q.data(), a2.data()); + for (int o = 0; o < 4 * H; ++o) { + gates[o] = (float) a1[o] * m->lstm_ih_s[l][o] * xs + (float) a2[o] * m->lstm_hh_s[l][o] * hq.s + + m->lstm_b[l][o]; + } + float * c = st.c[l].data(); + float * h = st.h[l].data(); + for (int j = 0; j < H; ++j) { + const float ig = sigm(gates[j]); + const float fg = sigm(gates[H + j]); + const float gg = std::tanh(gates[2 * H + j]); + const float og = sigm(gates[3 * H + j]); + c[j] = fg * c[j] + ig * gg; + h[j] = og * std::tanh(c[j]); + } + } + const int J = m->hp.joint_hidden; + quant_vec(st.h[1].data(), H, hq); + kt.matvec(m->proj_q, J, H, hq.q.data(), a1.data()); + for (int o = 0; o < J; ++o) { + st.out[o] = (float) a1[o] * m->proj_s[o] * hq.s + m->proj_b[o]; + } +} + +} // namespace + +int pk::tdt_decode(pk_model * m, const float * encp, int T, std::vector & out, pk_cancel_cb cancel, + void * cu) { + const auto & hp = m->hp; + const int H = hp.pred_hidden, J = hp.joint_hidden, NR = m->head_rows; + const int nv = hp.vocab + 1; // tokens + blank + const int nd = (int) hp.durations.size(); + dec_state st; + for (int l = 0; l < 2; ++l) { + st.h[l].assign(H, 0.0f); + st.c[l].assign(H, 0.0f); + } + st.out.assign(J, 0.0f); + std::vector a1(4 * H), a2(4 * H), acc(NR); + std::vector gates(4 * H), z(J); + qvec xq, hq, zq; + lstm_step(m, hp.blank_id, st, a1, a2, gates, xq, hq); + out.clear(); + int t = 0, symbols = 0, steps = 0; + auto & P = *m->pool; + const auto & kt = *m->kt; + while (t < T) { + if (cancel && (++steps & 15) == 0 && cancel(cu)) return -1; + const float * e = encp + (size_t) t * J; + for (int j = 0; j < J; ++j) { + const float v = e[j] + st.out[j]; + z[j] = v > 0.0f ? v : 0.0f; + } + quant_vec(z.data(), J, zq); + const int8_t * zq8 = zq.q.data(); + if (P.size() > 1) { + P.parallel_for(NR, 1056, [&](int b, int en) { + kt.matvec(m->head_q + (size_t) b * J, en - b, J, zq8, acc.data() + b); + }); + } else { + kt.matvec(m->head_q, NR, J, zq8, acc.data()); + } + int best = 0; + float bv = -INFINITY; + for (int o = 0; o < nv; ++o) { + const float v = (float) acc[o] * m->head_s[o] * zq.s + m->head_b[o]; + if (v > bv) { + bv = v; + best = o; + } + } + int bd = 0; + float bdv = -INFINITY; + for (int d = 0; d < nd; ++d) { + const int o = nv + d; + const float v = (float) acc[o] * m->head_s[o] * zq.s + m->head_b[o]; + if (v > bdv) { + bdv = v; + bd = d; + } + } + int dur = hp.durations[bd]; + if (best == hp.blank_id) { + if (dur == 0) dur = 1; + symbols = 0; + } else { + out.push_back(pk_token{ best, t, dur }); + lstm_step(m, best, st, a1, a2, gates, xq, hq); + if (dur == 0) { + if (++symbols >= hp.max_symbols) { + dur = 1; + symbols = 0; + } + } else { + symbols = 0; + } + } + t += dur; + } + return 0; +} + +std::string pk::detokenize(const pk_model * m, const pk_token * toks, size_t n) { + std::string s; + for (size_t i = 0; i < n; ++i) { + const int id = toks[i].id; + if (id < 0 || id >= (int) m->pieces.size() || m->is_control[id]) continue; + s += m->pieces[id]; + } + size_t b = 0; + while (b < s.size() && s[b] == ' ') ++b; + size_t e = s.size(); + while (e > b && s[e - 1] == ' ') --e; + return s.substr(b, e - b); +} + +// ------------------------------------------------------------------ API + +static int pk_max_window_samples(const pk_model * m) { + const char * env = std::getenv("PK_MAX_WINDOW_S"); + double s = env && *env ? std::atof(env) : 30.0; + if (s < 4.0) s = 4.0; + return (int) (s * m->hp.sample_rate); +} + +// choose a cut in [lo, hi) at the quietest 100 ms +static size_t quiet_cut(const float * pcm, size_t lo, size_t hi) { + const size_t win = 1600; + size_t best = hi; + double be = 1e300; + for (size_t s = lo; s + win <= hi; s += 160) { + double e = 0.0; + for (size_t i = 0; i < win; ++i) e += (double) pcm[s + i] * pcm[s + i]; + if (e < be) { + be = e; + best = s + win / 2; + } + } + return best; +} + +extern "C" int pk_transcribe(pk_model * m, const float * pcm, size_t n, pk_cancel_cb cancel, void * cu, + pk_result * out) { + if (!m || !out || (!pcm && n)) return PK_E_INVALID; + std::memset(out, 0, sizeof(*out)); + std::vector all; + const size_t maxw = (size_t) pk_max_window_samples(m); + size_t pos = 0; + while (pos < n) { + size_t end = n; + if (n - pos > maxw) { + end = quiet_cut(pcm, pos + maxw - 5 * (size_t) m->hp.sample_rate, pos + maxw); + } + auto t0 = clk::now(); + std::vector mel; + int nm = 0; + if (!compute_mel(m, pcm + pos, end - pos, mel, nm)) { + pos = end; + continue; + } + out->t_mel_ms += ms_since(t0); + t0 = clk::now(); + std::vector encp; + const int T = encode(m, mel.data(), nm, encp, cancel, cu); + if (T < 0) return PK_E_CANCELLED; + out->t_encoder_ms += ms_since(t0); + t0 = clk::now(); + std::vector toks; + if (tdt_decode(m, encp.data(), T, toks, cancel, cu) < 0) return PK_E_CANCELLED; + out->t_decoder_ms += ms_since(t0); + // encoder frames are 80 ms; window starts are rounded to a frame + const int base = (int) std::lround((double) pos / (m->hp.hop_length * 8)); + for (auto & tk : toks) { + tk.frame += base; + all.push_back(tk); + } + out->n_frames = base + T; + pos = end; + } + const std::string text = detokenize(m, all.data(), all.size()); + out->text = (char *) std::malloc(text.size() + 1); + out->tokens = (pk_token *) std::malloc(sizeof(pk_token) * std::max(1, all.size())); + if (!out->text || !out->tokens) return PK_E_NOMEM; + std::memcpy(out->text, text.c_str(), text.size() + 1); + if (!all.empty()) std::memcpy(out->tokens, all.data(), sizeof(pk_token) * all.size()); + out->n_tokens = all.size(); + return PK_OK; +} + +extern "C" void pk_result_free(pk_result * r) { + if (!r) return; + std::free(r->text); + std::free(r->tokens); + r->text = nullptr; + r->tokens = nullptr; + r->n_tokens = 0; +} + +extern "C" int pk_result_words(const pk_model * m, const pk_result * r, int * ws, int * we, size_t * io_n) { + if (!m || !r || !io_n) return PK_E_INVALID; + const size_t cap = *io_n; + size_t n = 0; + const int fms = m->hp.hop_length * 8 * 1000 / m->hp.sample_rate; // 80 ms + int cur_start = -1, cur_end = -1; + bool has_chars = false; + auto flush = [&]() { + if (cur_start >= 0 && has_chars) { + if (n < cap) { + if (ws) ws[n] = cur_start; + if (we) we[n] = cur_end; + } + ++n; + } + cur_start = cur_end = -1; + has_chars = false; + }; + for (size_t i = 0; i < r->n_tokens; ++i) { + const pk_token & t = r->tokens[i]; + if (t.id < 0 || t.id >= (int) m->pieces.size() || m->is_control[t.id]) continue; + const std::string & p = m->pieces[t.id]; + // a word starts at a piece with a leading space, unless the current + // word has no characters yet (a lone space piece) + if (cur_start < 0 || (m->word_start[t.id] && has_chars)) { + flush(); + cur_start = t.frame * fms; + } + cur_end = (t.frame + std::max(1, t.duration)) * fms; + if (p.find_first_not_of(' ') != std::string::npos) has_chars = true; + } + flush(); + *io_n = std::min(n, cap); + return PK_OK; +} diff --git a/tools/parakeet/src/pk-kernels-avx2.cpp b/tools/parakeet/src/pk-kernels-avx2.cpp new file mode 100644 index 000000000..c92e94b7d --- /dev/null +++ b/tools/parakeet/src/pk-kernels-avx2.cpp @@ -0,0 +1,137 @@ +// tools/parakeet/src/pk-kernels-avx2.cpp - x86_64 AVX2 tier. +// Compiled with -mavx2 (/arch:AVX2 on MSVC); only called after a runtime +// CPUID + XGETBV check (cpu_has_avx2 in pk-kernels.cpp). +// +// The int8 operands are sign-extended to int16 and multiplied with VPMADDWD, +// which is exact for every int8 input (no VPMADDUBSW saturation and no +// -128 corner case), so this tier is bit-identical to the scalar reference. +#include "pk-kernels.h" + +#if PK_HAVE_X86_EXT +#include + +namespace pk { +namespace { + +// Same nibble -> level lookup as the NEON tiers: PSHUFB with a 16-byte table +// (the nibble indices are 0..15, so the high-bit zeroing never applies). +void decode_fv5_avx2(const uint8_t * codes, size_t pair_stride, const int8_t * luts, int K, int8_t * panel) { + const int G = K / 16; + const __m128i m = _mm_set1_epi8(0x0F); + for (int p = 0; p < 4; ++p) { + const __m128i lut = _mm_loadu_si128((const __m128i *) (luts + p * 16)); + const uint8_t * src = codes + (size_t) p * pair_stride; + int8_t * dst = panel + p * 16; + for (int g = 0; g < G; ++g) { + const __m128i b = _mm_loadu_si128((const __m128i *) (src + g * 16)); + const __m128i lo = _mm_and_si128(b, m); + const __m128i hi = _mm_and_si128(_mm_srli_epi16(b, 4), m); + _mm_storeu_si128((__m128i *) (dst + (size_t) (2 * g) * 64), _mm_shuffle_epi8(lut, lo)); + _mm_storeu_si128((__m128i *) (dst + (size_t) (2 * g + 1) * 64), _mm_shuffle_epi8(lut, hi)); + } + } +} + +inline __m256i load_s8_as_s16(const int8_t * p) { + return _mm256_cvtepi8_epi16(_mm_loadu_si128((const __m128i *) p)); +} + +// Sum the 4 int32 lanes of each 128-bit half: returns {sum(lo half), sum(hi half)}. +inline void hsum_halves(__m256i v, int32_t & lo, int32_t & hi) { + const __m128i a = _mm256_castsi256_si128(v); + const __m128i b = _mm256_extracti128_si256(v, 1); + // [a0+a1, a2+a3, b0+b1, b2+b3] -> [a, b, a, b] + __m128i s = _mm_hadd_epi32(a, b); + s = _mm_hadd_epi32(s, s); + lo = _mm_cvtsi128_si32(s); + hi = _mm_extract_epi32(s, 1); +} + +// One token pair against the 8-row panel. Per k chunk, row pair p is one +// 256-bit int16 vector [row 2p | row 2p+1] and the token pair is +// [token 2q | token 2q+1]; VPMADDWD against it (straight) gives +// [r0 t0 | r1 t1] and against the half-swapped tokens [r0 t1 | r1 t0], +// mirroring the dotprod tier. 8 accumulators + 4 weights + 2 tokens fit in +// the 16 YMM registers. +inline void pair_tile(const int8_t * panel, const int8_t * x, int C, int32_t * out, int t0) { + __m256i s0 = _mm256_setzero_si256(), s1 = s0, s2 = s0, s3 = s0; + __m256i w0 = s0, w1 = s0, w2 = s0, w3 = s0; + for (int c = 0; c < C; ++c) { + const int8_t * wp = panel + (size_t) c * 64; + const __m256i v = load_s8_as_s16(x + (size_t) c * 16); + const __m256i u = _mm256_permute4x64_epi64(v, 0x4E); + const __m256i p0 = load_s8_as_s16(wp); + const __m256i p1 = load_s8_as_s16(wp + 16); + const __m256i p2 = load_s8_as_s16(wp + 32); + const __m256i p3 = load_s8_as_s16(wp + 48); + s0 = _mm256_add_epi32(s0, _mm256_madd_epi16(p0, v)); w0 = _mm256_add_epi32(w0, _mm256_madd_epi16(p0, u)); + s1 = _mm256_add_epi32(s1, _mm256_madd_epi16(p1, v)); w1 = _mm256_add_epi32(w1, _mm256_madd_epi16(p1, u)); + s2 = _mm256_add_epi32(s2, _mm256_madd_epi16(p2, v)); w2 = _mm256_add_epi32(w2, _mm256_madd_epi16(p2, u)); + s3 = _mm256_add_epi32(s3, _mm256_madd_epi16(p3, v)); w3 = _mm256_add_epi32(w3, _mm256_madd_epi16(p3, u)); + } + const __m256i S[4] = { s0, s1, s2, s3 }; + const __m256i W[4] = { w0, w1, w2, w3 }; + for (int p = 0; p < 4; ++p) { + int32_t r0t0, r1t1, r0t1, r1t0; + hsum_halves(S[p], r0t0, r1t1); + hsum_halves(W[p], r0t1, r1t0); + out[t0 * 8 + 2 * p] = r0t0; + out[(t0 + 1) * 8 + 2 * p + 1] = r1t1; + out[(t0 + 1) * 8 + 2 * p] = r0t1; + out[t0 * 8 + 2 * p + 1] = r1t0; + } +} + +void tile_avx2(const int8_t * panel, const int8_t * x, size_t x_stride, int K, int np, int32_t * out) { + const int C = K / 8; + for (int q = 0; q < np; ++q) { + pair_tile(panel, x + (size_t) q * x_stride, C, out, 2 * q); + } +} + +inline int32_t hsum8(__m256i v) { + __m128i s = _mm_add_epi32(_mm256_castsi256_si128(v), _mm256_extracti128_si256(v, 1)); + s = _mm_hadd_epi32(s, s); + s = _mm_hadd_epi32(s, s); + return _mm_cvtsi128_si32(s); +} + +void matvec_avx2(const int8_t * w, int O, int K, const int8_t * x, int32_t * y) { + int o = 0; + for (; o + 4 <= O; o += 4) { + const int8_t * r0 = w + (size_t) o * K; + const int8_t * r1 = r0 + K; + const int8_t * r2 = r1 + K; + const int8_t * r3 = r2 + K; + __m256i a0 = _mm256_setzero_si256(), a1 = a0, a2 = a0, a3 = a0; + for (int k = 0; k < K; k += 16) { + const __m256i xv = load_s8_as_s16(x + k); + a0 = _mm256_add_epi32(a0, _mm256_madd_epi16(load_s8_as_s16(r0 + k), xv)); + a1 = _mm256_add_epi32(a1, _mm256_madd_epi16(load_s8_as_s16(r1 + k), xv)); + a2 = _mm256_add_epi32(a2, _mm256_madd_epi16(load_s8_as_s16(r2 + k), xv)); + a3 = _mm256_add_epi32(a3, _mm256_madd_epi16(load_s8_as_s16(r3 + k), xv)); + } + y[o] = hsum8(a0); + y[o + 1] = hsum8(a1); + y[o + 2] = hsum8(a2); + y[o + 3] = hsum8(a3); + } + for (; o < O; ++o) { + const int8_t * r0 = w + (size_t) o * K; + __m256i a0 = _mm256_setzero_si256(); + for (int k = 0; k < K; k += 16) { + a0 = _mm256_add_epi32(a0, _mm256_madd_epi16(load_s8_as_s16(r0 + k), load_s8_as_s16(x + k))); + } + y[o] = hsum8(a0); + } +} + +} // namespace + +const kernel_table & kernels_avx2() { + static const kernel_table t = { tier::avx2, "avx2", decode_fv5_avx2, tile_avx2, matvec_avx2 }; + return t; +} + +} // namespace pk +#endif diff --git a/tools/parakeet/src/pk-kernels-dotprod.cpp b/tools/parakeet/src/pk-kernels-dotprod.cpp new file mode 100644 index 000000000..ca631117d --- /dev/null +++ b/tools/parakeet/src/pk-kernels-dotprod.cpp @@ -0,0 +1,63 @@ +// tools/parakeet/src/pk-kernels-dotprod.cpp - ARMv8.2 dotprod tier. +// Compiled with -march=armv8.2-a+dotprod; only called after a runtime check. +#include "pk-kernels.h" + +#if PK_HAVE_ARM_EXT +#include "pk-kernels-neon.inc" + +namespace pk { +namespace { + +// NQ (1 or 2) token pairs at a time: 4 row pairs x NQ token pairs x +// {straight, swapped} accumulators. Straight lanes: r0t0, r0t0, r1t1, r1t1 +// (two k halves); swapped (token halves exchanged): r0t1, r0t1, r1t0, r1t0. +template +inline void part_tile(const int8_t * panel, const int8_t * x, size_t x_stride, int C, int32_t * out, int q0) { + int32x4_t s[4][NQ], w[4][NQ]; + for (int p = 0; p < 4; ++p) { + for (int q = 0; q < NQ; ++q) { + s[p][q] = vdupq_n_s32(0); + w[p][q] = vdupq_n_s32(0); + } + } + for (int c = 0; c < C; ++c) { + const int8_t * wp = panel + (size_t) c * 64; + const int8x16_t p0 = vld1q_s8(wp), p1 = vld1q_s8(wp + 16), p2 = vld1q_s8(wp + 32), p3 = vld1q_s8(wp + 48); + for (int q = 0; q < NQ; ++q) { + const int8x16_t v = vld1q_s8(x + (size_t) (q0 + q) * x_stride + (size_t) c * 16); + const int8x16_t u = vextq_s8(v, v, 8); + s[0][q] = vdotq_s32(s[0][q], p0, v); w[0][q] = vdotq_s32(w[0][q], p0, u); + s[1][q] = vdotq_s32(s[1][q], p1, v); w[1][q] = vdotq_s32(w[1][q], p1, u); + s[2][q] = vdotq_s32(s[2][q], p2, v); w[2][q] = vdotq_s32(w[2][q], p2, u); + s[3][q] = vdotq_s32(s[3][q], p3, v); w[3][q] = vdotq_s32(w[3][q], p3, u); + } + } + for (int p = 0; p < 4; ++p) { + for (int q = 0; q < NQ; ++q) { + // [r0t0, r1t1, r0t1, r1t0] + const int32x4_t r = vpaddq_s32(s[p][q], w[p][q]); + const int t0 = 2 * (q0 + q); + out[t0 * 8 + 2 * p] = vgetq_lane_s32(r, 0); + out[(t0 + 1) * 8 + 2 * p + 1] = vgetq_lane_s32(r, 1); + out[(t0 + 1) * 8 + 2 * p] = vgetq_lane_s32(r, 2); + out[t0 * 8 + 2 * p + 1] = vgetq_lane_s32(r, 3); + } + } +} + +void tile_dotprod(const int8_t * panel, const int8_t * x, size_t x_stride, int K, int np, int32_t * out) { + const int C = K / 8; + int q = 0; + for (; q + 2 <= np; q += 2) part_tile<2>(panel, x, x_stride, C, out, q); + if (q < np) part_tile<1>(panel, x, x_stride, C, out, q); +} + +} // namespace + +const kernel_table & kernels_dotprod() { + static const kernel_table t = { tier::dotprod, "dotprod", decode_fv5_neon, tile_dotprod, matvec_dotprod }; + return t; +} + +} // namespace pk +#endif diff --git a/tools/parakeet/src/pk-kernels-i8mm.cpp b/tools/parakeet/src/pk-kernels-i8mm.cpp new file mode 100644 index 000000000..6ecff5dc8 --- /dev/null +++ b/tools/parakeet/src/pk-kernels-i8mm.cpp @@ -0,0 +1,59 @@ +// tools/parakeet/src/pk-kernels-i8mm.cpp - ARMv8.6 int8 matrix multiply tier +// (SMMLA). Compiled with -march=armv8.2-a+dotprod+i8mm; only called after a +// runtime check (HWCAP2_I8MM / hw.optional.arm.FEAT_I8MM). +#include "pk-kernels.h" + +#if PK_HAVE_ARM_EXT +#include "pk-kernels-neon.inc" + +namespace pk { +namespace { + +template +inline void tile_np(const int8_t * panel, const int8_t * x, size_t x_stride, int K, int32_t * out) { + const int C = K / 8; + int32x4_t acc[4][NP]; + for (int p = 0; p < 4; ++p) { + for (int q = 0; q < NP; ++q) { + acc[p][q] = vdupq_n_s32(0); + } + } + for (int c = 0; c < C; ++c) { + const int8_t * w = panel + (size_t) c * 64; + const int8x16_t p0 = vld1q_s8(w), p1 = vld1q_s8(w + 16), p2 = vld1q_s8(w + 32), p3 = vld1q_s8(w + 48); + for (int q = 0; q < NP; ++q) { + const int8x16_t v = vld1q_s8(x + (size_t) q * x_stride + (size_t) c * 16); + acc[0][q] = vmmlaq_s32(acc[0][q], p0, v); + acc[1][q] = vmmlaq_s32(acc[1][q], p1, v); + acc[2][q] = vmmlaq_s32(acc[2][q], p2, v); + acc[3][q] = vmmlaq_s32(acc[3][q], p3, v); + } + } + for (int p = 0; p < 4; ++p) { + for (int q = 0; q < NP; ++q) { + // SMMLA result order: [r0t0, r0t1, r1t0, r1t1] (r0 = 2p, t0 = 2q) + const int32x4_t v = acc[p][q]; + vst1_s32(out + (2 * q) * 8 + 2 * p, vget_low_s32(vuzp1q_s32(v, v))); + vst1_s32(out + (2 * q + 1) * 8 + 2 * p, vget_low_s32(vuzp2q_s32(v, v))); + } + } +} + +void tile_i8mm(const int8_t * panel, const int8_t * x, size_t x_stride, int K, int np, int32_t * out) { + switch (np) { + case 1: tile_np<1>(panel, x, x_stride, K, out); break; + case 2: tile_np<2>(panel, x, x_stride, K, out); break; + case 3: tile_np<3>(panel, x, x_stride, K, out); break; + default: tile_np<4>(panel, x, x_stride, K, out); break; + } +} + +} // namespace + +const kernel_table & kernels_i8mm() { + static const kernel_table t = { tier::i8mm, "i8mm", decode_fv5_neon, tile_i8mm, matvec_dotprod }; + return t; +} + +} // namespace pk +#endif diff --git a/tools/parakeet/src/pk-kernels-neon.inc b/tools/parakeet/src/pk-kernels-neon.inc new file mode 100644 index 000000000..2ecbf4703 --- /dev/null +++ b/tools/parakeet/src/pk-kernels-neon.inc @@ -0,0 +1,63 @@ +// tools/parakeet/src/pk-kernels-neon.inc - NEON pieces shared by the dotprod +// and i8mm translation units (each is compiled with its own -march). +#include + +namespace pk { +namespace { + +inline void decode_fv5_neon(const uint8_t * codes, size_t pair_stride, const int8_t * luts, int K, int8_t * panel) { + const int G = K / 16; + const uint8x16_t m = vdupq_n_u8(0x0F); + for (int p = 0; p < 4; ++p) { + const int8x16_t lut = vld1q_s8(luts + p * 16); + const uint8_t * src = codes + (size_t) p * pair_stride; + int8_t * dst = panel + p * 16; + int g = 0; + for (; g + 2 <= G; g += 2) { + const uint8x16_t b0 = vld1q_u8(src + g * 16); + const uint8x16_t b1 = vld1q_u8(src + g * 16 + 16); + vst1q_s8(dst + (size_t) (2 * g) * 64, vqtbl1q_s8(lut, vandq_u8(b0, m))); + vst1q_s8(dst + (size_t) (2 * g + 1) * 64, vqtbl1q_s8(lut, vshrq_n_u8(b0, 4))); + vst1q_s8(dst + (size_t) (2 * g + 2) * 64, vqtbl1q_s8(lut, vandq_u8(b1, m))); + vst1q_s8(dst + (size_t) (2 * g + 3) * 64, vqtbl1q_s8(lut, vshrq_n_u8(b1, 4))); + } + for (; g < G; ++g) { + const uint8x16_t b0 = vld1q_u8(src + g * 16); + vst1q_s8(dst + (size_t) (2 * g) * 64, vqtbl1q_s8(lut, vandq_u8(b0, m))); + vst1q_s8(dst + (size_t) (2 * g + 1) * 64, vqtbl1q_s8(lut, vshrq_n_u8(b0, 4))); + } + } +} + +inline void matvec_dotprod(const int8_t * w, int O, int K, const int8_t * x, int32_t * y) { + int o = 0; + for (; o + 4 <= O; o += 4) { + const int8_t * r0 = w + (size_t) o * K; + const int8_t * r1 = r0 + K; + const int8_t * r2 = r1 + K; + const int8_t * r3 = r2 + K; + int32x4_t a0 = vdupq_n_s32(0), a1 = a0, a2 = a0, a3 = a0; + for (int k = 0; k < K; k += 16) { + const int8x16_t xv = vld1q_s8(x + k); + a0 = vdotq_s32(a0, vld1q_s8(r0 + k), xv); + a1 = vdotq_s32(a1, vld1q_s8(r1 + k), xv); + a2 = vdotq_s32(a2, vld1q_s8(r2 + k), xv); + a3 = vdotq_s32(a3, vld1q_s8(r3 + k), xv); + } + y[o] = vaddvq_s32(a0); + y[o + 1] = vaddvq_s32(a1); + y[o + 2] = vaddvq_s32(a2); + y[o + 3] = vaddvq_s32(a3); + } + for (; o < O; ++o) { + const int8_t * r0 = w + (size_t) o * K; + int32x4_t a0 = vdupq_n_s32(0); + for (int k = 0; k < K; k += 16) { + a0 = vdotq_s32(a0, vld1q_s8(r0 + k), vld1q_s8(x + k)); + } + y[o] = vaddvq_s32(a0); + } +} + +} // namespace +} // namespace pk diff --git a/tools/parakeet/src/pk-kernels.cpp b/tools/parakeet/src/pk-kernels.cpp new file mode 100644 index 000000000..3f83a8ff0 --- /dev/null +++ b/tools/parakeet/src/pk-kernels.cpp @@ -0,0 +1,178 @@ +// tools/parakeet/src/pk-kernels.cpp - scalar reference kernels and the +// runtime CPU-feature dispatch. +#include "pk-kernels.h" + +#include +#include + +#if defined(__aarch64__) && defined(__linux__) +#include +#ifndef HWCAP_ASIMDDP +#define HWCAP_ASIMDDP (1 << 20) +#endif +#ifndef HWCAP2_I8MM +#define HWCAP2_I8MM (1 << 13) +#endif +#endif +#if defined(__APPLE__) +#include +#endif +#if PK_HAVE_X86_EXT +#if defined(_MSC_VER) +#include +#else +#include +#endif +#endif + +namespace pk { + +static void decode_fv5_scalar(const uint8_t * codes, size_t pair_stride, const int8_t * luts, int K, int8_t * panel) { + const int G = K / 16; + for (int p = 0; p < 4; ++p) { + const uint8_t * src = codes + (size_t) p * pair_stride; + const int8_t * lut = luts + p * 16; + for (int g = 0; g < G; ++g) { + int8_t * a = panel + ((size_t) (2 * g) * 4 + p) * 16; + int8_t * b = panel + ((size_t) (2 * g + 1) * 4 + p) * 16; + for (int j = 0; j < 16; ++j) { + const uint8_t v = src[g * 16 + j]; + a[j] = lut[v & 0x0F]; + b[j] = lut[v >> 4]; + } + } + } +} + +static void tile_scalar(const int8_t * panel, const int8_t * x, size_t x_stride, int K, int np, int32_t * out) { + int32_t acc[8][8]; + std::memset(acc, 0, sizeof(acc)); + const int C = K / 8; + for (int c = 0; c < C; ++c) { + const int8_t * w = panel + (size_t) c * 64; + for (int q = 0; q < np; ++q) { + const int8_t * xv = x + (size_t) q * x_stride + (size_t) c * 16; + for (int p = 0; p < 4; ++p) { + const int8_t * wv = w + p * 16; + for (int rr = 0; rr < 2; ++rr) { + for (int tt = 0; tt < 2; ++tt) { + int32_t s = 0; + for (int k = 0; k < 8; ++k) { + s += (int32_t) wv[rr * 8 + k] * (int32_t) xv[tt * 8 + k]; + } + acc[2 * p + rr][2 * q + tt] += s; + } + } + } + } + } + for (int r = 0; r < 8; ++r) { + for (int t = 0; t < 2 * np; ++t) { + out[t * 8 + r] = acc[r][t]; + } + } +} + +static void matvec_scalar(const int8_t * w, int O, int K, const int8_t * x, int32_t * y) { + for (int o = 0; o < O; ++o) { + const int8_t * row = w + (size_t) o * K; + int32_t s = 0; + for (int k = 0; k < K; ++k) { + s += (int32_t) row[k] * (int32_t) x[k]; + } + y[o] = s; + } +} + +const kernel_table & kernels_scalar() { + static const kernel_table t = { tier::scalar, "scalar", decode_fv5_scalar, tile_scalar, matvec_scalar }; + return t; +} + +bool cpu_has_dotprod() { +#if defined(__aarch64__) && defined(__APPLE__) + int v = 0; + size_t n = sizeof(v); + return sysctlbyname("hw.optional.arm.FEAT_DotProd", &v, &n, nullptr, 0) == 0 && v != 0; +#elif defined(__aarch64__) && defined(__linux__) + return (getauxval(AT_HWCAP) & HWCAP_ASIMDDP) != 0; +#else + return false; +#endif +} + +bool cpu_has_i8mm() { +#if defined(__aarch64__) && defined(__APPLE__) + int v = 0; + size_t n = sizeof(v); + return sysctlbyname("hw.optional.arm.FEAT_I8MM", &v, &n, nullptr, 0) == 0 && v != 0; +#elif defined(__aarch64__) && defined(__linux__) + return (getauxval(AT_HWCAP2) & HWCAP2_I8MM) != 0; +#else + return false; +#endif +} + +bool cpu_has_avx2() { +#if PK_HAVE_X86_EXT + unsigned a = 0, b = 0, c = 0, d = 0; +#if defined(_MSC_VER) + int r[4]; + __cpuid(r, 0); + if (r[0] < 7) return false; + __cpuid(r, 1); + c = (unsigned) r[2]; +#else + if (__get_cpuid_max(0, nullptr) < 7) return false; + __cpuid(1, a, b, c, d); +#endif + // OSXSAVE (ecx bit 27) and AVX (bit 28), then XCR0 must enable the XMM and + // YMM state (bits 1 and 2) or the OS does not preserve the upper halves. + if (!(c & (1u << 27)) || !(c & (1u << 28))) return false; +#if defined(_MSC_VER) + const unsigned long long xcr0 = _xgetbv(0); +#else + unsigned lo = 0, hi = 0; + __asm__ volatile("xgetbv" : "=a"(lo), "=d"(hi) : "c"(0)); + const unsigned long long xcr0 = ((unsigned long long) hi << 32) | lo; +#endif + if ((xcr0 & 0x6) != 0x6) return false; +#if defined(_MSC_VER) + __cpuidex(r, 7, 0); + b = (unsigned) r[1]; +#else + __cpuid_count(7, 0, a, b, c, d); +#endif + (void) a; (void) d; + return (b & (1u << 5)) != 0; // AVX2 +#else + return false; +#endif +} + +const kernel_table & kernels_for(int want) { + const char * env = std::getenv("PK_KERNEL"); + if (env && *env) { + if (!std::strcmp(env, "scalar")) want = 1; + else if (!std::strcmp(env, "dotprod")) want = 2; + else if (!std::strcmp(env, "i8mm")) want = 3; + else if (!std::strcmp(env, "avx2")) want = 4; + } +#if PK_HAVE_ARM_EXT + const bool dp = cpu_has_dotprod(); + const bool mm = dp && cpu_has_i8mm(); + if (want == 1) return kernels_scalar(); + if (want == 2 && dp) return kernels_dotprod(); + if (want == 3 && mm) return kernels_i8mm(); + if (mm) return kernels_i8mm(); + if (dp) return kernels_dotprod(); +#elif PK_HAVE_X86_EXT + if (want == 1) return kernels_scalar(); + if (cpu_has_avx2()) return kernels_avx2(); +#else + (void) want; +#endif + return kernels_scalar(); +} + +} // namespace pk diff --git a/tools/parakeet/src/pk-kernels.h b/tools/parakeet/src/pk-kernels.h new file mode 100644 index 000000000..43cc5fa3e --- /dev/null +++ b/tools/parakeet/src/pk-kernels.h @@ -0,0 +1,68 @@ +// tools/parakeet/src/pk-kernels.h - int8 GEMM kernels for the five-value +// (fv5) encoder linears and the int6 tables, plus the dispatch table. +// +// Tiers: scalar (portable), dotprod and i8mm (aarch64), avx2 (x86_64). +// +// Shared layouts (K is always a multiple of 16, rows a multiple of 8): +// +// activations [N/2 token pairs][K/8 chunks][16 bytes] +// 16 bytes = token 2q at k 8c..8c+7, then token 2q+1 at the +// same k. Odd N is padded with a zero token. One fp32 scale +// per token (symmetric, round-half-even). +// +// weight panel [K/8 chunks][4 row pairs][16 bytes] per block of 8 rows +// 16 bytes = row 2p at k 8c..8c+7, then row 2p+1. This is the +// operand order of SMMLA (2x8 by 2x8) and of a pair of SDOTs. +// +// fv5 codes `fv5-pair16-v1` (see convert_phonon2_to_gguf.py): per row +// pair, per 16 k, 16 bytes; low nibbles give chunk 2g, high +// nibbles chunk 2g+1, each already in panel order. A 16-byte +// table per row pair maps the nibble to the int8 level: +// 0..4 -> row 2p {-H,-L,0,L,H}, 8..12 -> row 2p+1. +// +// Every tier computes exact int32 dot products over the same int8 operands, +// so all tiers are bit-identical (tests/test_parakeet_kernels.cpp). +#pragma once + +#include +#include + +namespace pk { + +enum class tier : int { scalar = 1, dotprod = 2, i8mm = 3, avx2 = 4 }; + +struct kernel_table { + tier t; + const char * name; + // decode the fv5 codes of one 8-row block (4 row pairs starting at + // `codes`, `pair_stride` bytes apart) into a weight panel. + void (*decode_fv5)(const uint8_t * codes, size_t pair_stride, const int8_t * luts, int K, int8_t * panel); + // int32 tile: 8 rows x 2*np tokens (np = 1..4 token pairs). `panel` is + // the 8-row weight panel, `x` the first token pair, `x_stride` the bytes + // between token pairs. out[t * 8 + r] (token-major, so the epilogue + // reads 8 consecutive output rows of one token). + void (*tile)(const int8_t * panel, const int8_t * x, size_t x_stride, int K, int np, int32_t * out); + // int8 matvec for the decoder: y[o] = sum_k w[o][k] * x[k] for o in + // [0, O) over row-major int8 weights (K % 16 == 0). + void (*matvec)(const int8_t * w, int O, int K, const int8_t * x, int32_t * y); +}; + +// Best tier this CPU supports, or `want` when it is available. +const kernel_table & kernels_for(int want /*0 = best*/); +bool cpu_has_dotprod(); +bool cpu_has_i8mm(); +bool cpu_has_avx2(); // CPUID AVX2 and the OS saving YMM state (XGETBV) + +// per-tier tables (defined in their own translation units) +const kernel_table & kernels_scalar(); +#if defined(__aarch64__) && !defined(PK_NO_ARM_EXT) && !defined(_MSC_VER) +#define PK_HAVE_ARM_EXT 1 +const kernel_table & kernels_dotprod(); +const kernel_table & kernels_i8mm(); +#endif +#if (defined(__x86_64__) || defined(_M_X64)) && !defined(PK_NO_X86_EXT) +#define PK_HAVE_X86_EXT 1 +const kernel_table & kernels_avx2(); +#endif + +} // namespace pk diff --git a/tools/parakeet/src/pk-model.h b/tools/parakeet/src/pk-model.h new file mode 100644 index 000000000..50444aa27 --- /dev/null +++ b/tools/parakeet/src/pk-model.h @@ -0,0 +1,121 @@ +// tools/parakeet/src/pk-model.h - internal model structures shared by +// parakeet.cpp (load, encoder, decoder) and pk-stream.cpp. +#pragma once + +#include "parakeet.h" +#include "pk-kernels.h" +#include "pk-ops.h" +#include "pk-pool.h" + +#include +#include +#include + +struct gguf_context; +struct ggml_context; + +namespace pk { + +struct hparams { + int d_model = 1024, n_layers = 24, n_heads = 8, ff_dim = 4096, conv_kernel = 9; + int sub_channels = 256, sub_factor = 8, n_mels = 128; + int sample_rate = 16000, n_fft = 512, win_length = 400, hop_length = 160; + int pred_hidden = 640, pred_layers = 2, joint_hidden = 640, vocab = 8192, blank_id = 8192, max_symbols = 10; + float preemph = 0.97f, log_guard = 5.9604644775390625e-8f, norm_eps = 1e-5f, ln_eps = 1e-5f; + std::vector durations{ 0, 1, 2, 3, 4 }; +}; + +struct layer { + wmat ff1_w1, ff1_w2, qkv, out, pos, pw1, pw2, ff2_w1, ff2_w2; + const float * bias_u = nullptr; + const float * bias_v = nullptr; + const float * ln_w[5] = {}; + const float * ln_b[5] = {}; // ff1, att, conv, ff2, out + const float * dw_w = nullptr; // [d][k] (BatchNorm folded) + std::vector dw_wT; // [k][d] + const float * dw_b = nullptr; +}; + +struct mapped_file { + void * addr = nullptr; + size_t size = 0; + std::vector heap; // fallback when mmap is unavailable + ~mapped_file(); + bool open(const std::string & path, bool use_mmap, std::string & err); + const uint8_t * data() const { return addr ? (const uint8_t *) addr : heap.data(); } +}; + +// scratch buffers reused across calls on one model (calls are serialised) +struct workspace { + std::vector x, xn, h, qkv, att, g, u, v, pos, rpos, encp; + qact qa, qh, qp; + std::vector sub_a, sub_b, sub_c; + std::vector> att_tls; // per-thread attention scratch +}; + +} // namespace pk + +struct pk_model { + pk::hparams hp; + pk::mapped_file file; + gguf_context * gg = nullptr; + ggml_context * meta = nullptr; + std::unique_ptr pool; + const pk::kernel_table * kt = nullptr; + std::string notice; + + std::vector pieces; // with U+2581 already turned into ' ' + std::vector is_control; // <...> control pieces are never printed + std::vector word_start; // piece starts a word + + const float * mel_fb = nullptr; // [n_mels][n_fft/2+1] + std::vector mel_lo, mel_hi; // non-zero bin range per mel + std::vector window; // [n_fft], hann(win) centred + + std::vector c0_wT, dw1_wT, dw2_wT; // [9][C] + const float *c0_b = nullptr, *dw1_b = nullptr, *dw2_b = nullptr; + pk::wmat sub_pw1, sub_pw2, sub_out, enc_proj; + std::vector layers; + + const int8_t * embed_q = nullptr; + const float * embed_s = nullptr; + const int8_t * lstm_ih[2] = {}; + const int8_t * lstm_hh[2] = {}; + const float * lstm_ih_s[2] = {}; + const float * lstm_hh_s[2] = {}; + const float * lstm_b[2] = {}; + const int8_t * proj_q = nullptr; + const float * proj_s = nullptr; + const float * proj_b = nullptr; + const int8_t * head_q = nullptr; + const float * head_s = nullptr; + const float * head_b = nullptr; + int head_rows = 0; + + pk::workspace ws; + double fv5_grid_err = 0.0; + + // relative-position keys W_pos . e(p) per layer for p in [-rc_P, rc_P], + // row c <-> p = rc_P - c, fp16. They depend only on p, so they are + // computed once and grown on demand (24 x (2P+1) x 1024 x 2 bytes: + // 12 MB for 10 s windows, 37 MB for 30 s). + int rc_P = -1; + std::vector> rcache; +}; + +namespace pk { + +// log-mel features [n_frames][n_mels] (n_frames = n_samples / hop); returns +// false when the input is too short to normalise (< 2 frames). +bool compute_mel(pk_model * m, const float * pcm, size_t n, std::vector & mel, int & n_frames); +// encoder output projected to the joint width: [T][joint_hidden]. Optionally +// keeps the raw encoder output [T][d_model] in `enc_out`. +int encode(pk_model * m, const float * mel, int n_mel, std::vector & encp, pk_cancel_cb cancel, void * cu, + std::vector * enc_out = nullptr, std::vector * sub_out = nullptr, + std::vector * l0_out = nullptr); +// greedy TDT decode over encp [T][joint_hidden] +int tdt_decode(pk_model * m, const float * encp, int T, std::vector & out, pk_cancel_cb cancel, void * cu); +std::string detokenize(const pk_model * m, const pk_token * toks, size_t n); +int subsampled_length(int n_mel); + +} // namespace pk diff --git a/tools/parakeet/src/pk-ops.cpp b/tools/parakeet/src/pk-ops.cpp new file mode 100644 index 000000000..748aa5802 --- /dev/null +++ b/tools/parakeet/src/pk-ops.cpp @@ -0,0 +1,473 @@ +// tools/parakeet/src/pk-ops.cpp +#include "pk-ops.h" + +#include +#include +#include +#include + +#if (defined(__ARM_NEON) || defined(__aarch64__)) && !defined(_MSC_VER) +#include +#define PK_NEON 1 +#endif + +namespace pk { + +// ---------------------------------------------------------------- exp + +#if PK_NEON +// Cephes-style expf: range reduction to [-ln2/2, ln2/2], degree-5 polynomial, +// 2^n through the exponent field. Max rel error ~2 ulp on [-87, 88]. +static inline float32x4_t exp_ps(float32x4_t x) { + const float32x4_t hi = vdupq_n_f32(88.3762626647949f); + const float32x4_t lo = vdupq_n_f32(-87.3365447505531f); + x = vminq_f32(vmaxq_f32(x, lo), hi); + const float32x4_t log2e = vdupq_n_f32(1.44269504088896341f); + const float32x4_t n = vrndnq_f32(vmulq_f32(x, log2e)); + float32x4_t r = vfmsq_f32(x, n, vdupq_n_f32(0.693359375f)); + r = vfmsq_f32(r, n, vdupq_n_f32(-2.12194440e-4f)); + float32x4_t p = vdupq_n_f32(1.9875691500E-4f); + p = vfmaq_f32(vdupq_n_f32(1.3981999507E-3f), p, r); + p = vfmaq_f32(vdupq_n_f32(8.3334519073E-3f), p, r); + p = vfmaq_f32(vdupq_n_f32(4.1665795894E-2f), p, r); + p = vfmaq_f32(vdupq_n_f32(1.6666665459E-1f), p, r); + p = vfmaq_f32(vdupq_n_f32(5.0000001201E-1f), p, r); + const float32x4_t r2 = vmulq_f32(r, r); + p = vfmaq_f32(vaddq_f32(r, vdupq_n_f32(1.0f)), p, r2); + const int32x4_t e = vshlq_n_s32(vaddq_s32(vcvtq_s32_f32(n), vdupq_n_s32(127)), 23); + return vmulq_f32(p, vreinterpretq_f32_s32(e)); +} + +static inline float32x4_t silu_ps(float32x4_t x) { + const float32x4_t e = exp_ps(vnegq_f32(x)); + return vdivq_f32(x, vaddq_f32(vdupq_n_f32(1.0f), e)); +} +#endif + +float sigmoidf_(float x) { return 1.0f / (1.0f + std::exp(-x)); } + +void exp_inplace(float * x, size_t n) { + size_t i = 0; +#if PK_NEON + for (; i + 4 <= n; i += 4) { + vst1q_f32(x + i, exp_ps(vld1q_f32(x + i))); + } +#endif + for (; i < n; ++i) { + x[i] = std::exp(x[i]); + } +} + +void silu_inplace(float * x, size_t n) { + size_t i = 0; +#if PK_NEON + for (; i + 4 <= n; i += 4) { + vst1q_f32(x + i, silu_ps(vld1q_f32(x + i))); + } +#endif + for (; i < n; ++i) { + x[i] = x[i] / (1.0f + std::exp(-x[i])); + } +} + +// ---------------------------------------------------------- layernorm + +void layernorm(const float * x, float * y, int N, int D, const float * g, const float * b, float eps) { + for (int t = 0; t < N; ++t) { + const float * xr = x + (size_t) t * D; + float * yr = y + (size_t) t * D; + int i = 0; + float sum = 0.0f; +#if PK_NEON + float32x4_t s4 = vdupq_n_f32(0.0f); + for (; i + 4 <= D; i += 4) { + s4 = vaddq_f32(s4, vld1q_f32(xr + i)); + } + sum = vaddvq_f32(s4); +#endif + for (; i < D; ++i) { + sum += xr[i]; + } + const float mean = sum / (float) D; + float var = 0.0f; + i = 0; +#if PK_NEON + const float32x4_t m4 = vdupq_n_f32(mean); + float32x4_t v4 = vdupq_n_f32(0.0f); + for (; i + 4 <= D; i += 4) { + const float32x4_t d = vsubq_f32(vld1q_f32(xr + i), m4); + v4 = vfmaq_f32(v4, d, d); + } + var = vaddvq_f32(v4); +#endif + for (; i < D; ++i) { + const float d = xr[i] - mean; + var += d * d; + } + const float inv = 1.0f / std::sqrt(var / (float) D + eps); + i = 0; +#if PK_NEON + const float32x4_t inv4 = vdupq_n_f32(inv); + for (; i + 4 <= D; i += 4) { + const float32x4_t d = vmulq_f32(vsubq_f32(vld1q_f32(xr + i), m4), inv4); + vst1q_f32(yr + i, vfmaq_f32(vld1q_f32(b + i), d, vld1q_f32(g + i))); + } +#endif + for (; i < D; ++i) { + yr[i] = (xr[i] - mean) * inv * g[i] + b[i]; + } + } +} + +// ------------------------------------------------------ fp32 NT GEMM + +void gemm_f32_nt(const float * A, size_t lda, const float * B, size_t ldb, float * C, size_t ldc, int M, int N, int K, + float alpha) { +#if PK_NEON + const int K4 = K & ~3; + for (int m = 0; m < M; m += 4) { + const float * a[4]; + for (int i = 0; i < 4; ++i) { + a[i] = A + (size_t) std::min(m + i, M - 1) * lda; + } + for (int n = 0; n < N; n += 4) { + const float * bp[4]; + for (int j = 0; j < 4; ++j) { + bp[j] = B + (size_t) std::min(n + j, N - 1) * ldb; + } + float32x4_t c00 = vdupq_n_f32(0), c01 = c00, c02 = c00, c03 = c00, c10 = c00, c11 = c00, c12 = c00, + c13 = c00, c20 = c00, c21 = c00, c22 = c00, c23 = c00, c30 = c00, c31 = c00, c32 = c00, + c33 = c00; + for (int k = 0; k < K4; k += 4) { + const float32x4_t a0 = vld1q_f32(a[0] + k), a1 = vld1q_f32(a[1] + k), a2 = vld1q_f32(a[2] + k), + a3 = vld1q_f32(a[3] + k); + const float32x4_t b0 = vld1q_f32(bp[0] + k), b1 = vld1q_f32(bp[1] + k), b2 = vld1q_f32(bp[2] + k), + b3 = vld1q_f32(bp[3] + k); + c00 = vfmaq_f32(c00, a0, b0); c01 = vfmaq_f32(c01, a0, b1); c02 = vfmaq_f32(c02, a0, b2); c03 = vfmaq_f32(c03, a0, b3); + c10 = vfmaq_f32(c10, a1, b0); c11 = vfmaq_f32(c11, a1, b1); c12 = vfmaq_f32(c12, a1, b2); c13 = vfmaq_f32(c13, a1, b3); + c20 = vfmaq_f32(c20, a2, b0); c21 = vfmaq_f32(c21, a2, b1); c22 = vfmaq_f32(c22, a2, b2); c23 = vfmaq_f32(c23, a2, b3); + c30 = vfmaq_f32(c30, a3, b0); c31 = vfmaq_f32(c31, a3, b1); c32 = vfmaq_f32(c32, a3, b2); c33 = vfmaq_f32(c33, a3, b3); + } + float r[4][4] = { + { vaddvq_f32(c00), vaddvq_f32(c01), vaddvq_f32(c02), vaddvq_f32(c03) }, + { vaddvq_f32(c10), vaddvq_f32(c11), vaddvq_f32(c12), vaddvq_f32(c13) }, + { vaddvq_f32(c20), vaddvq_f32(c21), vaddvq_f32(c22), vaddvq_f32(c23) }, + { vaddvq_f32(c30), vaddvq_f32(c31), vaddvq_f32(c32), vaddvq_f32(c33) }, + }; + for (int k = K4; k < K; ++k) { + for (int i = 0; i < 4; ++i) { + for (int j = 0; j < 4; ++j) { + r[i][j] += a[i][k] * bp[j][k]; + } + } + } + for (int i = 0; i < 4 && m + i < M; ++i) { + for (int j = 0; j < 4 && n + j < N; ++j) { + C[(size_t) (m + i) * ldc + n + j] = alpha * r[i][j]; + } + } + } + } +#else + for (int m = 0; m < M; ++m) { + for (int n = 0; n < N; ++n) { + const float * a = A + (size_t) m * lda; + const float * b = B + (size_t) n * ldb; + float s = 0.0f; + for (int k = 0; k < K; ++k) { + s += a[k] * b[k]; + } + C[(size_t) m * ldc + n] = alpha * s; + } + } +#endif +} + +// -------------------------------------------------------- activations + +static void quantize_row(const float * x, int K, int8_t * dst, float * scale_out) { + // dst: first chunk of this token; chunks are 16 bytes apart, 8 bytes each + float amax = 0.0f; + int k = 0; +#if PK_NEON + float32x4_t m4 = vdupq_n_f32(0.0f); + for (; k + 4 <= K; k += 4) { + m4 = vmaxq_f32(m4, vabsq_f32(vld1q_f32(x + k))); + } + amax = vmaxvq_f32(m4); +#endif + for (; k < K; ++k) { + amax = std::max(amax, std::fabs(x[k])); + } + if (!(amax > 0.0f) || !std::isfinite(amax)) { + *scale_out = 0.0f; + for (int c = 0; c < K / 8; ++c) { + std::memset(dst + c * 16, 0, 8); + } + return; + } + const float inv = 127.0f / amax; + *scale_out = amax / 127.0f; +#if PK_NEON + const float32x4_t iv = vdupq_n_f32(inv); + for (int c = 0; c < K / 8; ++c) { + const int32x4_t q0 = vcvtnq_s32_f32(vmulq_f32(vld1q_f32(x + 8 * c), iv)); + const int32x4_t q1 = vcvtnq_s32_f32(vmulq_f32(vld1q_f32(x + 8 * c + 4), iv)); + const int16x8_t h = vcombine_s16(vqmovn_s32(q0), vqmovn_s32(q1)); + vst1_s8(dst + c * 16, vqmovn_s16(h)); + } +#else + for (int c = 0; c < K / 8; ++c) { + for (int j = 0; j < 8; ++j) { + float v = std::nearbyint(x[8 * c + j] * inv); + v = std::min(127.0f, std::max(-127.0f, v)); + dst[c * 16 + j] = (int8_t) v; + } + } +#endif +} + +void quantize_acts(pool & P, const float * X, int N, int K, size_t ldx, qact & out) { + out.N = N; + out.K = K; + out.src = X; + out.ldsrc = ldx; + out.Npad = (N + 1) & ~1; + const size_t need = (size_t) out.Npad * K; + if (out.buf.size() < need) { + out.buf.resize(need); + } + if (out.scale.size() < (size_t) out.Npad) { + out.scale.resize(out.Npad); + } + const size_t ps = out.pair_stride(); + auto body = [&](int b, int e) { + for (int t = b; t < e; ++t) { + int8_t * dst = out.buf.data() + (size_t) (t / 2) * ps + (t & 1) * 8; + if (t < N) { + quantize_row(X + (size_t) t * ldx, K, dst, &out.scale[t]); + } else { + out.scale[t] = 0.0f; + for (int c = 0; c < K / 8; ++c) { + std::memset(dst + c * 16, 0, 8); + } + } + } + }; + if ((size_t) N * K >= 32768) { + P.parallel_for(out.Npad, 8, body); + } else { + body(0, out.Npad); + } +} + +// ------------------------------------------------------------ weights + +void wmat_from_fv5(wmat & w, const uint8_t * codes, const float * lohi, int O, int K) { + w.O = O; + w.K = K; + w.fv5 = codes; + w.lut.assign((size_t) (O / 2) * 16, 0); + w.ws.assign(O, 0.0f); + w.grid_err = 0.0; + for (int r = 0; r < O; ++r) { + const double lo = lohi[2 * r], hi = lohi[2 * r + 1]; + int H = 0, L = 0; + double s = 0.0; + if (hi > 0.0 && lo >= 0.0 && lo <= hi) { + const double ratio = lo / hi; + double best = 1e30; + for (int h = 127; h >= 64; --h) { + const int l = (int) std::lround(h * ratio); + const double err = ratio > 0 ? std::fabs(l - h * ratio) / (h * ratio) : 0.0; + if (err < best - 1e-12) { + best = err; + H = h; + L = l; + } + } + s = hi / H; + // error of the lo level relative to the row's largest weight + w.grid_err = std::max(w.grid_err, std::fabs(L * s - lo) / hi); + } else if (hi != 0.0 || lo != 0.0) { + // not seen in Phonon-2; keep both magnitudes on one grid + const double m = std::max(std::fabs(lo), std::fabs(hi)); + s = m / 127.0; + H = (int) std::lround(hi / s); + L = (int) std::lround(lo / s); + } + w.ws[r] = (float) s; + int8_t * lut = w.lut.data() + (size_t) (r / 2) * 16 + (r & 1) * 8; + lut[0] = (int8_t) -H; + lut[1] = (int8_t) -L; + lut[2] = 0; + lut[3] = (int8_t) L; + lut[4] = (int8_t) H; + } +} + +void wmat_from_int8(wmat & w, const int8_t * q, const float * scale, int O, int K, const int * perm) { + w.O = O; + w.K = K; + w.fv5 = nullptr; + w.ws.assign(scale, scale + O); + w.panels.assign((size_t) O * K, 0); + const int C = K / 8; + for (int rb = 0; rb < O / 8; ++rb) { + int8_t * panel = w.panels.data() + (size_t) rb * 8 * K; + for (int c = 0; c < C; ++c) { + for (int p = 0; p < 4; ++p) { + for (int h = 0; h < 2; ++h) { + const int8_t * row = q + (size_t) (rb * 8 + 2 * p + h) * K; + for (int j = 0; j < 8; ++j) { + const int col = 8 * c + j; + panel[((size_t) c * 4 + p) * 16 + h * 8 + j] = row[perm ? perm[col] : col]; + } + } + } + } + } +} + +// --------------------------------------------------------------- GEMM + +static void gemm_f32_debug(pool & P, const kernel_table & kt, const wmat & W, const qact & X, float * Y, size_t ldy, + epi op, float alpha) { + // same weights (decoded int8 levels x row scale), fp32 activations; the + // source is copied first because Y may alias it + const int K = W.K; + std::vector xs((size_t) X.N * K); + for (int t = 0; t < X.N; ++t) std::memcpy(xs.data() + (size_t) t * K, X.src + (size_t) t * X.ldsrc, sizeof(float) * K); + P.parallel_for(W.O / 8, 1, [&](int b, int e) { + std::vector scratch((size_t) 8 * K); + for (int rb = b; rb < e; ++rb) { + const int8_t * panel; + if (W.fv5) { + kt.decode_fv5(W.fv5 + (size_t) rb * 4 * K, (size_t) K, W.lut.data() + (size_t) rb * 64, K, scratch.data()); + panel = scratch.data(); + } else { + panel = W.panels.data() + (size_t) rb * 8 * K; + } + for (int t = 0; t < X.N; ++t) { + const float * x = xs.data() + (size_t) t * K; + for (int r = 0; r < 8; ++r) { + double s = 0; + for (int k = 0; k < K; ++k) { + s += (double) panel[((size_t) (k / 8) * 4 + r / 2) * 16 + (r & 1) * 8 + (k & 7)] * x[k]; + } + const int o = rb * 8 + r; + float v = (float) s * W.ws[o] + (W.bias ? W.bias[o] : 0.0f); + float & y = Y[(size_t) t * ldy + o]; + switch (op) { + case epi::store: y = v; break; + case epi::silu: y = v / (1.0f + std::exp(-v)); break; + case epi::relu: y = v > 0 ? v : 0; break; + case epi::add: y += alpha * v; break; + } + } + } + } + }); +} + +void gemm_q8(pool & P, const kernel_table & kt, const wmat & W, const qact & X, float * Y, size_t ldy, epi op, + float alpha) { + static const bool dbg = std::getenv("PK_DEBUG_F32_ACTS") != nullptr; + if (dbg && X.src) { + gemm_f32_debug(P, kt, W, X, Y, ldy, op, alpha); + return; + } + const int K = W.K; + const int nb = W.O / 8; + const int npairs = (X.N + 1) / 2; + const int ntt = (npairs + 3) / 4; + const size_t ps = X.pair_stride(); + const int N = X.N; + // Row blocks are processed in groups of G: the group's panels are decoded + // once, then every token tile is multiplied against all G panels while it + // is hot in L1, which cuts the activation traffic from L2 by G. + int G = 1; + if (ntt > 1) { + G = std::max(1, std::min(4, nb / std::max(1, 4 * P.size()))); + } + const int ng = (nb + G - 1) / G; + P.parallel_for(ng, 1, [&](int b, int e) { + thread_local std::vector scratch; + if (W.fv5 && scratch.size() < (size_t) G * 8 * K) { + scratch.resize((size_t) G * 8 * K); + } + alignas(16) int32_t acc[64]; + const int8_t * panels[4]; + for (int g = b; g < e; ++g) { + const int rb0 = g * G, rb1 = std::min(nb, rb0 + G); + for (int rb = rb0; rb < rb1; ++rb) { + if (W.fv5) { + int8_t * dst = scratch.data() + (size_t) (rb - rb0) * 8 * K; + kt.decode_fv5(W.fv5 + (size_t) rb * 4 * K, (size_t) K, W.lut.data() + (size_t) rb * 64, K, dst); + panels[rb - rb0] = dst; + } else { + panels[rb - rb0] = W.panels.data() + (size_t) rb * 8 * K; + } + } + for (int tt = 0; tt < ntt; ++tt) { + for (int rb = rb0; rb < rb1; ++rb) { + const int8_t * panel = panels[rb - rb0]; + const float * ws = W.ws.data() + rb * 8; + const float * bias = W.bias ? W.bias + rb * 8 : nullptr; + kt.tile(panel, X.buf.data() + (size_t) tt * 4 * ps, ps, K, std::min(4, npairs - tt * 4), acc); + const int tn = std::min(8, N - tt * 8); + for (int t = 0; t < tn; ++t) { + const int tok = tt * 8 + t; + const float xs = X.scale[tok]; + float * y = Y + (size_t) tok * ldy + rb * 8; + const int32_t * a = acc + t * 8; +#if PK_NEON + const float32x4_t xs4 = vdupq_n_f32(xs); + float32x4_t v0 = vmulq_f32(vmulq_f32(vcvtq_f32_s32(vld1q_s32(a)), vld1q_f32(ws)), xs4); + float32x4_t v1 = vmulq_f32(vmulq_f32(vcvtq_f32_s32(vld1q_s32(a + 4)), vld1q_f32(ws + 4)), xs4); + if (bias) { + v0 = vaddq_f32(v0, vld1q_f32(bias)); + v1 = vaddq_f32(v1, vld1q_f32(bias + 4)); + } + switch (op) { + case epi::store: + break; + case epi::silu: + v0 = silu_ps(v0); + v1 = silu_ps(v1); + break; + case epi::relu: + v0 = vmaxq_f32(v0, vdupq_n_f32(0.0f)); + v1 = vmaxq_f32(v1, vdupq_n_f32(0.0f)); + break; + case epi::add: { + const float32x4_t al = vdupq_n_f32(alpha); + v0 = vfmaq_f32(vld1q_f32(y), v0, al); + v1 = vfmaq_f32(vld1q_f32(y + 4), v1, al); + break; + } + } + vst1q_f32(y, v0); + vst1q_f32(y + 4, v1); +#else + for (int r = 0; r < 8; ++r) { + float v = (float) a[r] * ws[r] * xs; + if (bias) { + v += bias[r]; + } + switch (op) { + case epi::store: y[r] = v; break; + case epi::silu: y[r] = v / (1.0f + std::exp(-v)); break; + case epi::relu: y[r] = v > 0.0f ? v : 0.0f; break; + case epi::add: y[r] += alpha * v; break; + } + } +#endif + } + } + } + } + }); +} + +} // namespace pk diff --git a/tools/parakeet/src/pk-ops.h b/tools/parakeet/src/pk-ops.h new file mode 100644 index 000000000..f08b381b4 --- /dev/null +++ b/tools/parakeet/src/pk-ops.h @@ -0,0 +1,69 @@ +// tools/parakeet/src/pk-ops.h - fp32 helpers (NEON with a scalar fallback) +// and the int8 GEMM driver used by the encoder. +#pragma once + +#include "pk-kernels.h" +#include "pk-pool.h" + +#include +#include +#include + +namespace pk { + +// ---- int8 activations ------------------------------------------------------ + +// N tokens x K in the token-pair layout of pk-kernels.h, padded to an even +// number of tokens. +struct qact { + std::vector buf; + std::vector scale; + int N = 0, K = 0, Npad = 0; + // the fp32 source (diagnostics only: PK_DEBUG_F32_ACTS=1 runs the GEMMs + // on it instead of the int8 copy; it must still be alive at GEMM time) + const float * src = nullptr; + size_t ldsrc = 0; + size_t pair_stride() const { return (size_t) K * 2; } +}; + +// Symmetric per-token int8 quantisation, round half to even. +void quantize_acts(pool & P, const float * X, int N, int K, size_t ldx, qact & out); + +// ---- int8 weights ---------------------------------------------------------- + +struct wmat { + int O = 0, K = 0; + // five-value: codes in `fv5-pair16-v1`, one 16-byte LUT per row pair + const uint8_t * fv5 = nullptr; + std::vector lut; + // int tables: weight panels prepared at load ([O/8][K/8][4][16]) + std::vector panels; + std::vector ws; // per-row dequant scale + const float * bias = nullptr; // optional + // max |grid lo - lo| / hi over the rows (fv5); hi is exact + double grid_err = 0.0; +}; + +void wmat_from_fv5(wmat & w, const uint8_t * codes, const float * lohi, int O, int K); +// row-major int8 [O][K] with per-row scales; `perm` (optional) permutes the +// input columns: new column j reads old column perm[j]. +void wmat_from_int8(wmat & w, const int8_t * q, const float * scale, int O, int K, const int * perm = nullptr); + +enum class epi : int { store, add, silu, relu }; + +// Y[t][o] (ldy) <- op(acc * ws[o] * xs[t] + bias[o]); for epi::add, +// Y += alpha * value. +void gemm_q8(pool & P, const kernel_table & kt, const wmat & W, const qact & X, float * Y, size_t ldy, epi op, + float alpha = 1.0f); + +// ---- fp32 ------------------------------------------------------------------ + +void layernorm(const float * x, float * y, int N, int D, const float * g, const float * b, float eps); +void silu_inplace(float * x, size_t n); +float sigmoidf_(float x); +void exp_inplace(float * x, size_t n); // vectorised expf +// C[m][n] = alpha * sum_k A[m][k] * B[n][k] (A: M x K, lda; B: N x K, ldb) +void gemm_f32_nt(const float * A, size_t lda, const float * B, size_t ldb, float * C, size_t ldc, int M, int N, int K, + float alpha); + +} // namespace pk diff --git a/tools/parakeet/src/pk-pool.cpp b/tools/parakeet/src/pk-pool.cpp new file mode 100644 index 000000000..c6012975d --- /dev/null +++ b/tools/parakeet/src/pk-pool.cpp @@ -0,0 +1,175 @@ +// tools/parakeet/src/pk-pool.cpp +#include "pk-pool.h" + +#include +#include +#include + +#if defined(__linux__) +#include +#include +#include +#endif +#if defined(__APPLE__) +#include +#include +#include +#endif +#if defined(_MSC_VER) +#include +#if defined(_M_ARM64) +#define PK_RELAX() __yield() +#else +#define PK_RELAX() _mm_pause() +#endif +#elif defined(__aarch64__) +#define PK_RELAX() __asm__ __volatile__("yield" ::: "memory") +#elif defined(__x86_64__) || defined(__i386__) +#define PK_RELAX() __builtin_ia32_pause() +#else +#define PK_RELAX() ((void) 0) +#endif + +namespace pk { + +std::vector fast_cpus() { + std::vector out; +#if defined(__linux__) + std::vector> f; + for (int c = 0; c < 64; ++c) { + const std::string p = "/sys/devices/system/cpu/cpu" + std::to_string(c) + "/cpufreq/cpuinfo_max_freq"; + FILE * fp = std::fopen(p.c_str(), "r"); + if (!fp) { + continue; + } + long v = 0; + if (std::fscanf(fp, "%ld", &v) == 1) { + f.emplace_back(c, v); + } + std::fclose(fp); + } + long mx = 0; + for (auto & e : f) { + mx = std::max(mx, e.second); + } + for (auto & e : f) { + if (mx > 0 && e.second * 10 >= mx * 8) { + out.push_back(e.first); + } + } +#endif + return out; +} + +int default_threads() { + int n = 0; +#if defined(__APPLE__) + int v = 0; + size_t len = sizeof(v); + if (sysctlbyname("hw.perflevel0.physicalcpu", &v, &len, nullptr, 0) == 0) { + n = v; + } +#elif defined(__linux__) + n = (int) fast_cpus().size(); +#endif + if (n <= 0) { + n = (int) std::thread::hardware_concurrency() / 2; + } + return std::max(1, std::min(4, n)); +} + +pool::pool(int n_threads, bool pin) : n_(std::max(1, n_threads)) { + std::vector cpus; + if (pin) { + cpus = fast_cpus(); + } + for (int t = 1; t < n_; ++t) { + threads_.emplace_back([this, t, cpus] { +#if defined(__linux__) + if (!cpus.empty()) { + cpu_set_t set; + CPU_ZERO(&set); + for (int c : cpus) { + CPU_SET(c, &set); + } + sched_setaffinity((pid_t) syscall(SYS_gettid), sizeof(set), &set); + } +#else + (void) cpus; +#endif +#if defined(__APPLE__) + // keep workers on performance cores (no affinity API on Darwin) + pthread_set_qos_class_self_np(QOS_CLASS_USER_INTERACTIVE, 0); +#endif + worker(t); + }); + } +} + +pool::~pool() { + { + std::lock_guard lk(m_); + stop_.store(true); + gen_.fetch_add(1); + } + cv_.notify_all(); + for (auto & th : threads_) { + th.join(); + } +} + +void pool::worker(int tid) { + uint64_t seen = 0; + for (;;) { + uint64_t g = gen_.load(std::memory_order_acquire); + if (g == seen) { + // spin briefly, then yield (other threads may need this core when + // the CPU is oversubscribed), then sleep + for (int i = 0; i < 20000 && g == seen; ++i) { + if (i < 2000) { + PK_RELAX(); + } else { + std::this_thread::yield(); + } + g = gen_.load(std::memory_order_acquire); + } + if (g == seen) { + std::unique_lock lk(m_); + sleeping_.fetch_add(1); + cv_.wait(lk, [&] { return gen_.load() != seen; }); + sleeping_.fetch_sub(1); + g = gen_.load(); + } + } + seen = g; + if (stop_.load()) { + return; + } + (*job_)(tid); + remaining_.fetch_sub(1, std::memory_order_acq_rel); + } +} + +void pool::run(const std::function & f) { + if (n_ == 1) { + f(0); + return; + } + job_ = &f; + remaining_.store(n_ - 1, std::memory_order_relaxed); + gen_.fetch_add(1); + if (sleeping_.load() > 0) { + std::lock_guard lk(m_); + cv_.notify_all(); + } + f(0); + for (int i = 0; remaining_.load(std::memory_order_acquire) > 0; ++i) { + if (i < 2000) { + PK_RELAX(); + } else { + std::this_thread::yield(); // a worker was preempted; let it run + } + } +} + +} // namespace pk diff --git a/tools/parakeet/src/pk-pool.h b/tools/parakeet/src/pk-pool.h new file mode 100644 index 000000000..70c0a6788 --- /dev/null +++ b/tools/parakeet/src/pk-pool.h @@ -0,0 +1,70 @@ +// tools/parakeet/src/pk-pool.h - a small persistent worker pool. Workers spin +// briefly between jobs (the encoder issues a few hundred short parallel ops +// per second of audio), then sleep on a condition variable when idle. +#pragma once + +#include +#include +#include +#include +#include +#include +#include + +namespace pk { + +class pool { + public: + pool(int n_threads, bool pin); + ~pool(); + pool(const pool &) = delete; + pool & operator=(const pool &) = delete; + + int size() const { return n_; } + + // Run f(tid) on every thread; the caller runs tid 0. Returns when all + // threads finished. + void run(const std::function & f); + + // Hand out [0, n) in chunks of `grain` to all threads; f(begin, end). + template void parallel_for(int n, int grain, F && f) { + if (n <= 0) { + return; + } + if (n_ == 1 || n <= grain) { + f(0, n); + return; + } + std::atomic next{0}; + run([&](int) { + for (;;) { + const int b = next.fetch_add(grain, std::memory_order_relaxed); + if (b >= n) { + break; + } + f(b, b + grain < n ? b + grain : n); + } + }); + } + + private: + void worker(int tid); + + int n_; + std::vector threads_; + std::atomic gen_{0}; + std::atomic remaining_{0}; + std::atomic sleeping_{0}; + std::atomic stop_{false}; + const std::function * job_ = nullptr; + std::mutex m_; + std::condition_variable cv_; +}; + +// Threads to use when the caller passes n_threads <= 0: the fast cores +// (Android big/prime cores, Apple performance cores), at most 4. +int default_threads(); +// Linux/Android: CPU ids whose cpuinfo_max_freq is within 20% of the fastest. +std::vector fast_cpus(); + +} // namespace pk diff --git a/tools/parakeet/src/pk-stream.cpp b/tools/parakeet/src/pk-stream.cpp new file mode 100644 index 000000000..9bb85656f --- /dev/null +++ b/tools/parakeet/src/pk-stream.cpp @@ -0,0 +1,413 @@ +// tools/parakeet/src/pk-stream.cpp - streaming on a full-context model. +// +// Parakeet-TDT-v3 (and Phonon-2) attends over the whole input and has no +// cache-aware mode, so a stream re-decodes a window. Windows must not start +// in the middle of speech: the model (the dense fp32 reference included) +// drops or invents words when its input begins mid-word. So the stream works +// on pause-delimited segments: +// +// * every decode covers [segment start, newest sample); +// * tokens are committed when two consecutive decodes agree on them and +// they lie at least `right_context` before the newest sample, on word +// boundaries only; +// * when a pause (>= 150 ms of low energy) lies in committed audio, the next +// segment starts in it (preferably between the last committed token and +// the first uncommitted one), so later windows begin in silence and stay +// short; committed tokens a new window re-decodes are skipped by count, or +// by frame and id when the re-decode differs; +// * a segment longer than `max_window` without a pause is cut at the +// quietest 100 ms of its last 2 s of committed audio; +// * the decode interval grows with the window (at least min_hop, at most +// a sixth of the window), bounding the re-decode cost; +// * finish() decodes the open segment with full right context. +// +// Committed text never changes. Partials are committed + tentative tokens. +#include "parakeet.h" + +#include "pk-model.h" + +#include +#include +#include +#include +#include +#include +#include +#include +#include + +using namespace pk; + +struct pk_stream { + pk_model * m = nullptr; + pk_stream_params p{}; + std::vector audio; // samples [audio_base, total) + size_t audio_base = 0; + size_t total = 0; + size_t last_decode_total = 0; + size_t seg_start = 0; // absolute sample, multiple of one encoder frame + size_t commit_sample = 0; // all committed tokens lie before this + std::vector committed; // absolute frames + std::vector tentative; // absolute frames + bool done = false; + float speech_level = 0.0f; // loud (90th percentile) 10 ms energy of the last window with tokens + int n_decodes = 0; + int n_skipped = 0; // decodes skipped on quiet audio + size_t committed_in_segment(int f0) const { + size_t n = 0; + while (n < committed.size() && committed[committed.size() - 1 - n].frame >= f0) ++n; + return n; + } + double decode_ms = 0.0; + double first_ms = -1.0; +}; + +extern "C" pk_stream_params pk_stream_default_params(void) { + pk_stream_params p; + p.right_context_ms = 800; + p.min_hop_ms = 400; + p.max_window_ms = 12000; + return p; +} + +extern "C" pk_stream * pk_stream_open(pk_model * m, pk_stream_params p) { + if (!m) return nullptr; + auto * s = new pk_stream(); + s->m = m; + s->p = p; + return s; +} + +extern "C" int pk_stream_feed(pk_stream * s, const float * pcm, size_t n) { + if (!s || (!pcm && n)) return PK_E_INVALID; + if (s->done) return PK_E_INVALID; + s->audio.insert(s->audio.end(), pcm, pcm + n); + s->total += n; + return (int) n; +} + +// token identity for re-decode matching: the piece without its word-start +// space, lower-cased (a segment's first word is often re-decoded capitalised) +static bool same_piece(const pk_model * m, int a, int b) { + if (a == b) return true; + const std::string & x = m->pieces[a]; + const std::string & y = m->pieces[b]; + size_t i = x.find_first_not_of(' '), j = y.find_first_not_of(' '); + if (i == std::string::npos || j == std::string::npos) return false; + if (x.size() - i != y.size() - j) return false; + for (; i < x.size(); ++i, ++j) { + if (std::tolower((unsigned char) x[i]) != std::tolower((unsigned char) y[j])) return false; + } + return true; +} + +static size_t ms_to_samples(const pk_stream * s, int ms) { + return (size_t) ((long long) std::max(0, ms) * s->m->hp.sample_rate / 1000); +} + +static const float * at(const pk_stream * s, size_t abs_sample) { return s->audio.data() + (abs_sample - s->audio_base); } + +// 10 ms frame energies over [a, b) +static std::vector frame_energy(const pk_stream * s, size_t a, size_t b) { + const size_t hop = 160; + std::vector e; + for (size_t x = a; x + hop <= b; x += hop) { + const float * p = at(s, x); + double acc = 0.0; + for (size_t i = 0; i < hop; ++i) acc += (double) p[i] * p[i]; + e.push_back((float) (acc / hop)); + } + return e; +} + +// 90th percentile of the 10 ms frame energies over [a, b) +static float loud_level(const pk_stream * s, size_t a, size_t b) { + std::vector e = frame_energy(s, a, b); + if (e.empty()) return 0.0f; + std::nth_element(e.begin(), e.begin() + e.size() * 9 / 10, e.end()); + return e[e.size() * 9 / 10]; +} + +// Quiet: every 10 ms frame below -60 dBFS, or below -30 dB of the speech +// level of the last window that held tokens. +static bool is_quiet(const pk_stream * s, size_t a, size_t b) { + const float thr = std::max(1e-6f, 0.03f * 0.03f * s->speech_level); + for (float v : frame_energy(s, a, b)) { + if (v >= thr) return false; + } + return true; +} + +static void drop_before(pk_stream * s, size_t abs_sample) { + if (abs_sample > s->audio_base) { + s->audio.erase(s->audio.begin(), s->audio.begin() + (abs_sample - s->audio_base)); + s->audio_base = abs_sample; + } +} + +// true when sample x falls in no committed token's span [frame, frame + dur) +static bool between_tokens(const pk_stream * s, size_t x, size_t fs) { + for (auto it = s->committed.rbegin(); it != s->committed.rend(); ++it) { + const size_t b = (size_t) it->frame * fs; + const size_t e = (size_t) (it->frame + std::max(1, it->duration)) * fs; + if (e <= x) return true; // tokens are in time order + if (b <= x) return false; // x inside this token + } + return true; +} + +// Latest pause midpoint in [lo, hi): >= 15 consecutive 10 ms frames below +// 5% (amplitude) of the segment's loud level (90th percentile energy), and +// not inside a committed token. 0 when there is none. +static size_t find_pause(const pk_stream * s, size_t lo, size_t hi) { + if (hi <= lo + 3200) return 0; + const std::vector seg = frame_energy(s, s->seg_start, std::min(hi, s->total)); + if (seg.size() < 40) return 0; + std::vector sorted(seg); + std::nth_element(sorted.begin(), sorted.begin() + sorted.size() * 9 / 10, sorted.end()); + const float loud = sorted[sorted.size() * 9 / 10]; + const float thr = std::max(1e-7f, 0.05f * 0.05f * loud); + size_t best = 0; + size_t run = 0; + for (size_t i = 0; i < seg.size(); ++i) { + run = seg[i] < thr ? run + 1 : 0; + if (run >= 15) { + const size_t x = s->seg_start + (i + 1) * 160; // end of the run so far + const size_t mid = x - run * 80; + const size_t fs = (size_t) s->m->hp.hop_length * 8; + const size_t cut = (mid / fs) * fs; // where the segment would start + if (mid >= lo && mid < hi && between_tokens(s, cut, fs)) best = mid; + } + } + return best; +} + +// centre of the quietest 100 ms in [lo, hi) +static size_t quietest(const pk_stream * s, size_t lo, size_t hi) { + const std::vector e = frame_energy(s, lo, hi); + if (e.size() < 10) return hi; + size_t best = 0; + double be = 1e300; + for (size_t i = 0; i + 10 <= e.size(); ++i) { + double acc = 0; + for (size_t j = 0; j < 10; ++j) acc += e[i + j]; + if (acc < be) { + be = acc; + best = i; + } + } + return lo + best * 160 + 800; +} + +static int decode_window(pk_stream * s, bool final, pk_cancel_cb cancel, void * cu) { + pk_model * m = s->m; + const size_t fs = (size_t) m->hp.hop_length * 8; // samples per encoder frame (1280) + const size_t ws = s->seg_start; + if (s->total <= ws) return 0; + const auto t0 = std::chrono::steady_clock::now(); + std::vector mel, encp; + int nm = 0; + std::vector toks; + if (compute_mel(m, at(s, ws), s->total - ws, mel, nm)) { + const int T = encode(m, mel.data(), nm, encp, cancel, cu); + if (T < 0) return PK_E_CANCELLED; + if (tdt_decode(m, encp.data(), T, toks, cancel, cu) < 0) return PK_E_CANCELLED; + } + const double ms = std::chrono::duration(std::chrono::steady_clock::now() - t0).count(); + s->n_decodes++; + s->decode_ms += ms; + if (s->first_ms < 0) s->first_ms = ms; + s->last_decode_total = s->total; + + const int f0 = (int) (ws / fs); + for (auto & t : toks) t.frame += f0; + // PK_STREAM_DEBUG: window geometry and token counts only. Transcript + // text and token ids never go to a log. + const bool dbg = std::getenv("PK_STREAM_DEBUG") != nullptr; + if (dbg) { + std::fprintf(stderr, "[pk-stream] window %.2f-%.2fs: %zu tokens, %.1f ms\n", ws / 16000.0, + s->total / 16000.0, toks.size(), ms); + } + + // The window re-decodes this segment's committed tokens too. Skip them by + // count when the ids agree, otherwise by the commit point. + size_t seg_committed = 0; + while (seg_committed < s->committed.size() && + s->committed[s->committed.size() - 1 - seg_committed].frame >= f0) { + ++seg_committed; + } + std::vector hyp; + { + const size_t c0 = s->committed.size() - seg_committed; + bool same = toks.size() >= seg_committed; + for (size_t i = 0; same && i < seg_committed; ++i) same = same_piece(m, toks[i].id, s->committed[c0 + i].id); + size_t start = 0; + if (same) { + start = seg_committed; + } else { + // the segment's committed tokens re-decoded differently: keep what + // starts after the commit point, minus re-emissions of the + // committed tail (same ids within 3 frames) + const int boundary = (int) (s->commit_sample / fs); + while (start < toks.size() && toks[start].frame < boundary - 3) ++start; + size_t drop = 0; + const size_t avail = toks.size() - start; + const size_t kmax = std::min(6, std::min(s->committed.size(), avail)); + for (size_t kk = kmax; kk >= 1 && !drop; --kk) { + bool ok = true; + for (size_t i = 0; i < kk && ok; ++i) { + const pk_token & c = s->committed[s->committed.size() - kk + i]; + ok = same_piece(m, c.id, toks[start + i].id) && std::abs(c.frame - toks[start + i].frame) <= 3; + } + if (ok) drop = kk; + } + start += drop; + while (start < toks.size() && toks[start].frame < boundary) ++start; + } + hyp.assign(toks.begin() + start, toks.end()); + } + if (final) { + s->committed.insert(s->committed.end(), hyp.begin(), hyp.end()); + s->tentative.clear(); + s->commit_sample = s->total; + return 0; + } + const size_t rc = ms_to_samples(s, s->p.right_context_ms); + const int stable = (int) ((s->total > rc ? s->total - rc : 0) / fs); + size_t k = 0; + while (k < hyp.size() && k < s->tentative.size() && hyp[k].id == s->tentative[k].id && hyp[k].frame < stable) ++k; + const bool force = s->total - ws > ms_to_samples(s, s->p.max_window_ms); + if (force) { + while (k < hyp.size() && hyp[k].frame < stable) ++k; + } + // whole words only: the next token must start a word. The newest word + // also counts as finished when every token of the hypothesis ends before + // the stable point (right_context of audio with no further token): + // otherwise the last word before a pause would wait for the next word. + const bool tail_done = !hyp.empty() && hyp.back().frame + std::max(1, hyp.back().duration) <= stable; + while (k > 0 && !(k < hyp.size() ? m->word_start[hyp[k].id] != 0 : tail_done)) --k; + if (k > 0) { + s->committed.insert(s->committed.end(), hyp.begin(), hyp.begin() + k); + // commit point: the start frame of the first uncommitted token, or + // the end of the last committed one when nothing is left + const size_t cp = k < hyp.size() ? (size_t) hyp[k].frame * fs + : (size_t) (hyp.back().frame + std::max(1, hyp.back().duration)) * fs; + s->commit_sample = std::max(s->commit_sample, cp); + } + s->tentative.assign(hyp.begin() + k, hyp.end()); + if (!toks.empty()) s->speech_level = loud_level(s, ws, s->total); + + // Start a new segment in a pause, so the next window begins in silence. + // Preferred: a pause between the last committed token and the first + // uncommitted one (the new window then holds no committed token). Else + // the latest pause in the last 3 s of committed audio (its committed + // tokens are re-decoded and skipped above). Past max_window with no + // pause at all: the quietest 100 ms of the last 2 s of committed audio. + // Past max_window with no token at all (noise, no speech): keep the + // newest half window, so cost and memory stay bounded and the window is + // still long enough for the per-window feature normalisation. + if (hyp.empty() && s->committed_in_segment(f0) == 0 && force) { + const size_t keep = ms_to_samples(s, s->p.max_window_ms / 2); + const size_t cut = ((s->total - keep) / fs) * fs; + if (cut > ws) { + s->seg_start = cut; + s->commit_sample = std::max(s->commit_sample, cut); + drop_before(s, cut); + } + } else if (!s->committed.empty() && s->commit_sample > ws + fs) { + const pk_token & lc = s->committed.back(); + const size_t gap_lo = std::max(ws + fs, (size_t) (lc.frame + std::max(1, lc.duration)) * fs); + size_t cut = gap_lo < s->commit_sample ? find_pause(s, gap_lo, s->commit_sample) : 0; + if (!cut) { + const size_t lo3 = std::max(ws + fs, s->commit_sample > 48000 ? s->commit_sample - 48000 : 0); + cut = find_pause(s, lo3, s->commit_sample); + } + if (!cut && force) { + const size_t lo2 = std::max(ws + fs, s->commit_sample > 32000 ? s->commit_sample - 32000 : 0); + if (s->commit_sample > lo2 + 1600) cut = quietest(s, lo2, s->commit_sample); + } + cut = (cut / fs) * fs; + if (cut > ws && cut <= s->commit_sample) { + s->seg_start = cut; + drop_before(s, cut); + } + } + if (dbg) { + std::fprintf(stderr, "[pk-stream] commit+%zu -> %.2fs, %zu tentative, segment from %.2fs%s\n", k, + s->commit_sample / 16000.0, s->tentative.size(), s->seg_start / 16000.0, + force ? " (forced)" : ""); + } + return 0; +} + +static int write_text(pk_stream * s, char * out, size_t cap) { + std::vector all = s->committed; + all.insert(all.end(), s->tentative.begin(), s->tentative.end()); + const std::string t = detokenize(s->m, all.data(), all.size()); + if (!out || cap == 0) return (int) t.size(); + size_t n = std::min(t.size(), cap - 1); + // never cut a UTF-8 sequence: back off over continuation bytes + if (n < t.size()) { + while (n > 0 && ((unsigned char) t[n] & 0xC0) == 0x80) --n; + } + std::memcpy(out, t.data(), n); + out[n] = '\0'; + return (int) n; +} + +extern "C" int pk_stream_partial(pk_stream * s, char * out, size_t cap, pk_cancel_cb cancel, void * cu) { + if (!s) return PK_E_INVALID; + // decode interval: at least min_hop, growing to a sixth of the window + const size_t win = s->total - s->seg_start; + const size_t hop = std::max(ms_to_samples(s, s->p.min_hop_ms), win / 6); + if (!s->done && s->total - s->last_decode_total >= hop) { + if (s->tentative.empty() && s->last_decode_total > s->seg_start && + is_quiet(s, s->last_decode_total, s->total)) { + // Nothing pending and only quiet audio since the last decode: a + // decode could only re-read silence (and short noise windows are + // where the model invents words). Skip it, and keep at most the + // newest 1.5 s of a quiet stretch, so a long pause neither grows + // the next window nor the buffer. + s->n_skipped++; + s->last_decode_total = s->total; + const size_t fs = (size_t) s->m->hp.hop_length * 8; + size_t lo = s->seg_start + fs; + if (!s->committed.empty()) { + const pk_token & lc = s->committed.back(); + lo = std::max(lo, (size_t) (lc.frame + std::max(1, lc.duration)) * fs); + } + const size_t keep = ms_to_samples(s, 1500); + if (s->total > lo + keep + fs) { + const size_t cut = ((s->total - keep) / fs) * fs; + if (cut > lo && is_quiet(s, std::max(lo, cut - fs), s->total)) { + s->seg_start = cut; + s->commit_sample = std::max(s->commit_sample, cut); + drop_before(s, cut); + } + } + } else { + const int rc = decode_window(s, false, cancel, cu); + if (rc < 0) return rc; + } + } + return write_text(s, out, cap); +} + +extern "C" int pk_stream_finish(pk_stream * s, char * out, size_t cap, pk_cancel_cb cancel, void * cu) { + if (!s) return PK_E_INVALID; + if (!s->done) { + const int rc = decode_window(s, true, cancel, cu); + if (rc < 0) return rc; + s->done = true; + } + return write_text(s, out, cap); +} + +extern "C" void pk_stream_stats(const pk_stream * s, int * n, double * ms, double * first) { + if (!s) return; + if (n) *n = s->n_decodes; + if (ms) *ms = s->decode_ms; + if (first) *first = s->first_ms; +} + +extern "C" void pk_stream_close(pk_stream * s) { delete s; } diff --git a/tools/parakeet/tests/CMakeLists.txt b/tools/parakeet/tests/CMakeLists.txt new file mode 100644 index 000000000..90db2ca99 --- /dev/null +++ b/tools/parakeet/tests/CMakeLists.txt @@ -0,0 +1,20 @@ +# tools/parakeet/tests/CMakeLists.txt + +add_executable(test-parakeet-kernels test_parakeet_kernels.cpp) +target_link_libraries(test-parakeet-kernels PRIVATE parakeet_lib) +target_include_directories(test-parakeet-kernels PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/../src) +add_test(NAME test-parakeet-kernels COMMAND test-parakeet-kernels) +# weight-free and fast: runs in CI (ctest -L main) +set_property(TEST test-parakeet-kernels PROPERTY LABELS main) + +# Model parity gate against reference_transformers.py dumps +# (tools/dump_parakeet_reference.py). Needs the converted GGUF and the +# fixture directory, so it only runs when both are given: +# -DPARAKEET_TEST_MODEL=/path/phonon-2-fv5.gguf -DPARAKEET_TEST_FIXTURES=/path/fixtures +add_executable(test-parakeet-parity test_parakeet_parity.cpp) +target_link_libraries(test-parakeet-parity PRIVATE parakeet_lib) +target_include_directories(test-parakeet-parity PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/../src) +if (PARAKEET_TEST_MODEL AND PARAKEET_TEST_FIXTURES) + add_test(NAME test-parakeet-parity + COMMAND test-parakeet-parity ${PARAKEET_TEST_MODEL} ${PARAKEET_TEST_FIXTURES}) +endif() diff --git a/tools/parakeet/tests/test_parakeet_kernels.cpp b/tools/parakeet/tests/test_parakeet_kernels.cpp new file mode 100644 index 000000000..8476b463a --- /dev/null +++ b/tools/parakeet/tests/test_parakeet_kernels.cpp @@ -0,0 +1,191 @@ +// tools/parakeet/tests/test_parakeet_kernels.cpp - every kernel tier must give +// bit-identical int32 results to the scalar reference (decode, tile, matvec, +// and the whole gemm_q8 driver), and gemm_q8 must match a double-precision +// GEMM over the same int8 operands. +#include "pk-kernels.h" +#include "pk-ops.h" +#include "pk-pool.h" + +#include +#include +#include +#include +#include +#include + +using namespace pk; + +static int fails = 0; +#define CHECK(c, ...) do { if (!(c)) { std::fprintf(stderr, "FAIL %s:%d: ", __FILE__, __LINE__); std::fprintf(stderr, __VA_ARGS__); std::fprintf(stderr, "\n"); ++fails; } } while (0) + +static void make_fv5(std::mt19937 & rng, int O, int K, std::vector & codes, std::vector & lohi) { + std::uniform_int_distribution lv(0, 4); + std::uniform_real_distribution mag(0.001f, 0.05f), ratio(1.6f, 2.6f); + codes.assign((size_t) O / 2 * K, 0); + // build levels row-major, then pack in fv5-pair16-v1 order + std::vector L((size_t) O * K); + for (auto & v : L) v = (uint8_t) lv(rng); + for (int p = 0; p < O / 2; ++p) { + for (int g = 0; g < K / 16; ++g) { + for (int j = 0; j < 16; ++j) { + const int r = 2 * p + (j >= 8); + const int k = 16 * g + (j & 7); + const uint8_t a = L[(size_t) r * K + k] + (j >= 8 ? 8 : 0); + const uint8_t b = L[(size_t) r * K + k + 8] + (j >= 8 ? 8 : 0); + codes[(size_t) p * K + 16 * g + j] = (uint8_t) (a | (b << 4)); + } + } + } + lohi.resize((size_t) O * 2); + for (int r = 0; r < O; ++r) { + const float lo = mag(rng); + lohi[2 * r] = lo; + lohi[2 * r + 1] = lo * ratio(rng); + } +} + +int main() { + std::mt19937 rng(1234); + const kernel_table & ref = kernels_scalar(); + std::vector tiers; +#if PK_HAVE_ARM_EXT + if (cpu_has_dotprod()) tiers.push_back(&kernels_dotprod()); + if (cpu_has_i8mm()) tiers.push_back(&kernels_i8mm()); +#endif +#if PK_HAVE_X86_EXT + if (cpu_has_avx2()) tiers.push_back(&kernels_avx2()); +#endif + std::printf("tiers under test: scalar"); + for (auto * t : tiers) std::printf(", %s", t->name); + std::printf("\n"); + + for (int K : { 256, 1024, 4096 }) { + const int O = 64, N = 21; + std::vector codes; + std::vector lohi; + make_fv5(rng, O, K, codes, lohi); + wmat W; + wmat_from_fv5(W, codes.data(), lohi.data(), O, K); + CHECK(W.grid_err < 0.01, "grid error %g", W.grid_err); + std::vector X((size_t) N * K); + std::normal_distribution nd(0.0f, 1.0f); + for (auto & v : X) v = nd(rng); + pool P1(1, false), P3(3, false); + qact qa; + quantize_acts(P1, X.data(), N, K, K, qa); + // panel + tile identity + std::vector pr((size_t) 8 * K), pt((size_t) 8 * K); + int32_t a[64], b[64]; + ref.decode_fv5(codes.data(), K, W.lut.data(), K, pr.data()); + ref.tile(pr.data(), qa.buf.data(), qa.pair_stride(), K, 4, a); + for (auto * t : tiers) { + t->decode_fv5(codes.data(), K, W.lut.data(), K, pt.data()); + CHECK(std::memcmp(pr.data(), pt.data(), pr.size()) == 0, "%s decode differs (K=%d)", t->name, K); + t->tile(pr.data(), qa.buf.data(), qa.pair_stride(), K, 4, b); + CHECK(std::memcmp(a, b, sizeof(a)) == 0, "%s tile differs (K=%d)", t->name, K); + for (int np = 1; np <= 3; ++np) { + int32_t c[64]; + t->tile(pr.data(), qa.buf.data(), qa.pair_stride(), K, np, c); + for (int tk = 0; tk < 2 * np; ++tk) { + CHECK(std::memcmp(a + tk * 8, c + tk * 8, 32) == 0, "%s tile np=%d differs (K=%d)", t->name, np, K); + } + } + } + // tile vs direct dot products over the decoded levels + for (int r = 0; r < 8; ++r) { + for (int tk = 0; tk < 8; ++tk) { + int64_t s = 0; + for (int k = 0; k < K; ++k) { + const int c = k / 8, j = k % 8; + const int8_t w = pr[((size_t) c * 4 + r / 2) * 16 + (r & 1) * 8 + j]; + const int8_t x = qa.buf[(size_t) (tk / 2) * qa.pair_stride() + c * 16 + (tk & 1) * 8 + j]; + s += (int64_t) w * x; + } + CHECK(s == a[tk * 8 + r], "scalar tile wrong at r=%d t=%d", r, tk); + } + } + // full driver: tiers and thread counts agree bit for bit + std::vector y0((size_t) N * O), y1((size_t) N * O); + gemm_q8(P1, ref, W, qa, y0.data(), O, epi::store); + for (auto * t : tiers) { + gemm_q8(P3, *t, W, qa, y1.data(), O, epi::store); + CHECK(std::memcmp(y0.data(), y1.data(), y0.size() * 4) == 0, "%s gemm differs (K=%d)", t->name, K); + } + // and the driver result is the dequantised product + double maxrel = 0; + for (int t = 0; t < N; ++t) { + for (int o = 0; o < O; ++o) { + double s = 0, sw = 0; + for (int k = 0; k < K; ++k) { + const int c = k / 8, j = k % 8; + const int8_t x = qa.buf[(size_t) (t / 2) * qa.pair_stride() + c * 16 + (t & 1) * 8 + j]; + // decode the weight level from the codes + const int p = o / 2, g = k / 16, jj = (k % 16) % 8 + (o & 1) * 8; + const uint8_t byte = codes[(size_t) p * K + 16 * g + jj]; + const int lv = ((k % 16) < 8 ? (byte & 15) : (byte >> 4)) - (o & 1) * 8; + const double mags[5] = { -lohi[2 * o + 1], -lohi[2 * o], 0.0, lohi[2 * o], lohi[2 * o + 1] }; + s += mags[lv] * x * qa.scale[t]; + sw += std::fabs(mags[lv] * x * qa.scale[t]); + } + maxrel = std::max(maxrel, std::fabs(s - y0[(size_t) t * O + o]) / (sw + 1e-12)); + } + } + CHECK(maxrel < 2e-3, "gemm vs exact five-value product: rel err %g (K=%d)", maxrel, K); + std::printf("K=%d ok: grid_err %.2e, gemm rel err vs exact five-value %.2e\n", K, W.grid_err, maxrel); + // matvec + std::vector mw((size_t) 37 * K), mx(K); + std::uniform_int_distribution i8(-127, 127); + for (auto & v : mw) v = (int8_t) i8(rng); + for (auto & v : mx) v = (int8_t) i8(rng); + std::vector m0(37), m1(37); + ref.matvec(mw.data(), 37, K, mx.data(), m0.data()); + for (auto * t : tiers) { + t->matvec(mw.data(), 37, K, mx.data(), m1.data()); + CHECK(m0 == m1, "%s matvec differs", t->name); + } + } + // Random full-range operands (every int8 value, -128 included, which the + // model never produces): each tier's tile and matvec must still equal the + // scalar reference exactly, for every K and token-pair count. + { + std::uniform_int_distribution any8(-128, 127); + int cases = 0; + for (int it = 0; it < 200; ++it) { + const int K = 16 * (1 + (int) (rng() % 96)); + const int np = 1 + (int) (rng() % 4); + const size_t xs = (size_t) 2 * K + 16 * (rng() % 3); // padded pair stride + std::vector pan((size_t) 8 * K), xb(xs * 4); + for (auto & v : pan) v = (int8_t) any8(rng); + for (auto & v : xb) v = (int8_t) any8(rng); + if (it == 0) { // the extreme corner: every product is (-128)*(-128) + std::fill(pan.begin(), pan.end(), (int8_t) -128); + std::fill(xb.begin(), xb.end(), (int8_t) -128); + } + int32_t a[64], b[64]; + ref.tile(pan.data(), xb.data(), xs, K, np, a); + for (auto * t : tiers) { + std::memset(b, 0x5a, sizeof(b)); + t->tile(pan.data(), xb.data(), xs, K, np, b); + CHECK(std::memcmp(a, b, (size_t) 2 * np * 8 * 4) == 0, "%s random tile differs (K=%d np=%d)", t->name, K, np); + } + const int O = 1 + (int) (rng() % 23); + std::vector mw((size_t) O * K), mx(K); + for (auto & v : mw) v = (int8_t) any8(rng); + for (auto & v : mx) v = (int8_t) any8(rng); + if (it == 0) { + std::fill(mw.begin(), mw.end(), (int8_t) -128); + std::fill(mx.begin(), mx.end(), (int8_t) -128); + } + std::vector m0(O), m1(O); + ref.matvec(mw.data(), O, K, mx.data(), m0.data()); + for (auto * t : tiers) { + t->matvec(mw.data(), O, K, mx.data(), m1.data()); + CHECK(m0 == m1, "%s random matvec differs (K=%d O=%d)", t->name, K, O); + } + ++cases; + } + std::printf("random full-range int8: %d tile + matvec cases vs scalar\n", cases); + } + std::printf(fails ? "FAILED (%d)\n" : "all kernel checks passed\n", fails); + return fails ? 1 : 0; +} diff --git a/tools/parakeet/tests/test_parakeet_parity.cpp b/tools/parakeet/tests/test_parakeet_parity.cpp new file mode 100644 index 000000000..80d00e980 --- /dev/null +++ b/tools/parakeet/tests/test_parakeet_parity.cpp @@ -0,0 +1,239 @@ +// tools/parakeet/tests/test_parakeet_parity.cpp - compare the native pipeline +// with reference_transformers.py (dense fp32 ParakeetForTDT) stage by stage. +// +// test-parakeet-parity model.gguf fixtures_dir [--report out.tsv] +// +// fixtures_dir is written by tools/dump_parakeet_reference.py. Gates: +// mel max |diff| <= 2e-3 (normalised log-mel) +// encoder mean per-frame cosine >= PK_GATE_ENC_COS and relative L2 +// error <= PK_GATE_ENC_L2; the minimum frame cosine is reported +// (near-silent frames have tiny norms) +// +// Two modes: +// default the shipped int8-activation path. Per-token int8 +// activations through 24 layers leave a relative L2 +// of ~0.07-0.2 on the encoder output (chaotic in the +// input), so the gates are cos >= 0.98, L2 <= 0.25; +// the transcript gate is the functional one. +// PK_DEBUG_F32_ACTS=1 the same graph and weights with fp32 activations +// (slow diagnostic GEMM): checks the graph itself +// against the reference, cos >= 0.9999, L2 <= 0.02 +// (what remains is the fp16 key cache and the +// per-row int8 grid of the five-value levels). +// transcript at least PK_GATE_TEXT (default 0.9) of the utterances +// text-identical to the reference, and the word error rate +// against the reference transcripts (all utterances pooled, +// case-folded, punctuation kept) at most PK_GATE_REF_WER +// (default 0.02). The second gate bounds how far the +// utterances that differ may drift. +// An empty or unreadable fixture set fails: a parity gate that compared +// nothing must not pass. +// The native encoder quantises activations to int8 (see parakeet.cpp), so +// it is not bit-exact; the report prints the actual numbers. +#include "parakeet.h" +#include "pk-model.h" + +#include +#include +#include +#include +#include +#include +#include +#include +#include + +static bool read_f32(const std::string & path, std::vector & out) { + std::ifstream f(path, std::ios::binary | std::ios::ate); + if (!f) return false; + const size_t n = (size_t) f.tellg(); + f.seekg(0); + out.resize(n / 4); + f.read((char *) out.data(), (std::streamsize) n); + return (bool) f; +} + +struct cmp { + double max_abs = 0, min_cos = 1, mean_cos = 0, rel_l2 = 0; +}; + +static std::vector words_of(const std::string & s) { + std::vector w; + std::istringstream is(s); + std::string x; + while (is >> x) { + for (auto & ch : x) ch = (char) std::tolower((unsigned char) ch); + w.push_back(x); + } + return w; +} + +static size_t edit_distance(const std::vector & a, const std::vector & b) { + std::vector d(b.size() + 1), nd(b.size() + 1); + for (size_t j = 0; j <= b.size(); ++j) d[j] = j; + for (size_t i = 1; i <= a.size(); ++i) { + nd[0] = i; + for (size_t j = 1; j <= b.size(); ++j) { + nd[j] = std::min({ d[j] + 1, nd[j - 1] + 1, d[j - 1] + (a[i - 1] == b[j - 1] ? 0 : 1) }); + } + d.swap(nd); + } + return d[b.size()]; +} + +static cmp compare(const std::vector & a, const std::vector & b, int rows, int cols) { + cmp c; + double sc = 0, de = 0, nb = 0; + for (int r = 0; r < rows; ++r) { + double ab = 0, aa = 0, bb = 0; + for (int k = 0; k < cols; ++k) { + const double x = a[(size_t) r * cols + k], y = b[(size_t) r * cols + k]; + c.max_abs = std::max(c.max_abs, std::fabs(x - y)); + ab += x * y; + aa += x * x; + bb += y * y; + de += (x - y) * (x - y); + nb += y * y; + } + const double cs = ab / (std::sqrt(aa * bb) + 1e-30); + c.min_cos = std::min(c.min_cos, cs); + sc += cs; + } + c.mean_cos = rows ? sc / rows : 1; + c.rel_l2 = std::sqrt(de / (nb + 1e-30)); + return c; +} + +int main(int argc, char ** argv) { + if (argc < 3) { + std::fprintf(stderr, "usage: %s model.gguf fixtures_dir\n", argv[0]); + return 2; + } + const std::string dir = argv[2]; + const bool f32_acts = std::getenv("PK_DEBUG_F32_ACTS") != nullptr; + const double gate_cos = std::getenv("PK_GATE_ENC_COS") ? std::atof(std::getenv("PK_GATE_ENC_COS")) + : (f32_acts ? 0.9999 : 0.98); + const double gate_l2 = std::getenv("PK_GATE_ENC_L2") ? std::atof(std::getenv("PK_GATE_ENC_L2")) + : (f32_acts ? 0.02 : 0.25); + const double gate_text = std::getenv("PK_GATE_TEXT") ? std::atof(std::getenv("PK_GATE_TEXT")) : 0.9; + const double gate_wer = std::getenv("PK_GATE_REF_WER") ? std::atof(std::getenv("PK_GATE_REF_WER")) : 0.02; + char err[512] = { 0 }; + pk_params pp = pk_default_params(); + pk_model * m = pk_model_load(argv[1], pp, err, sizeof(err)); + if (!m) { + std::fprintf(stderr, "load: %s\n", err); + return 1; + } + std::printf("kernel %s, %d threads, fv5 int8 grid max rel err %.2e\n", pk_model_kernel_name(m), + pk_model_n_threads(m), m->fv5_grid_err); + std::ifstream man(dir + "/manifest.tsv"); + if (!man) { + std::fprintf(stderr, "no %s/manifest.tsv\n", dir.c_str()); + return 1; + } + std::string line; + int n = 0, same_text = 0, same_ids = 0, fails = 0, n_enc = 0; + size_t wer_err = 0, wer_words = 0; + double worst_mel = 0, worst_cos = 1, worst_mean_cos = 1, worst_l2 = 0; + while (std::getline(man, line)) { + std::vector f; + std::stringstream ss(line); + std::string x; + while (std::getline(ss, x, '\t')) f.push_back(x); + if (f.size() < 5) continue; + const std::string id = f[0]; + const int n_mel = std::atoi(f[2].c_str()), T = std::atoi(f[3].c_str()); + std::vector ref_ids; + { + std::stringstream is(f[4]); + std::string t; + while (std::getline(is, t, ',')) if (!t.empty()) ref_ids.push_back(std::atoi(t.c_str())); + } + const std::string ref_text = f.size() > 5 ? f[5] : ""; + std::vector pcm, rmel, renc; + if (!read_f32(dir + "/" + id + ".pcm.f32", pcm)) { + std::fprintf(stderr, "%s: missing pcm\n", id.c_str()); + ++fails; + continue; + } + std::vector mel; + int nm = 0; + pk::compute_mel(m, pcm.data(), pcm.size(), mel, nm); + std::printf("%-18s n_mel %d/%d", id.c_str(), nm, n_mel); + if (nm != n_mel) ++fails; + if (read_f32(dir + "/" + id + ".mel.f32", rmel) && (int) rmel.size() == n_mel * 128 && nm == n_mel) { + const cmp c = compare(mel, rmel, nm, 128); + worst_mel = std::max(worst_mel, c.max_abs); + std::printf(" mel max|d| %.2e", c.max_abs); + } + std::vector encp, enc, sub, l0; + const int Tn = pk::encode(m, mel.data(), nm, encp, nullptr, nullptr, &enc, &sub, &l0); + std::printf(" T %d/%d", Tn, T); + if (Tn != T) ++fails; + std::vector r; + if (Tn == T && read_f32(dir + "/" + id + ".sub.f32", r) && (int) r.size() == T * 1024) { + const cmp c = compare(sub, r, T, 1024); + std::printf(" sub cos min %.5f", c.min_cos); + } + if (Tn == T && read_f32(dir + "/" + id + ".l0.f32", r) && (int) r.size() == T * 1024) { + const cmp c = compare(l0, r, T, 1024); + std::printf(" l0 cos min %.5f", c.min_cos); + } + if (Tn == T && read_f32(dir + "/" + id + ".enc.f32", renc) && (int) renc.size() == T * 1024) { + const cmp c = compare(enc, renc, T, 1024); + worst_cos = std::min(worst_cos, c.min_cos); + worst_mean_cos = std::min(worst_mean_cos, c.mean_cos); + worst_l2 = std::max(worst_l2, c.rel_l2); + ++n_enc; + std::printf(" enc cos min %.5f mean %.5f relL2 %.4f", c.min_cos, c.mean_cos, c.rel_l2); + } + std::vector toks; + pk::tdt_decode(m, encp.data(), Tn, toks, nullptr, nullptr); + std::vector ids; + for (auto & t : toks) ids.push_back(t.id); + const std::string text = pk::detokenize(m, toks.data(), toks.size()); + const bool st = text == ref_text, si = ids == ref_ids; + { + const auto rw = words_of(ref_text), gw = words_of(text); + wer_err += edit_distance(rw, gw); + wer_words += rw.size(); + } + same_text += st; + same_ids += si; + std::printf(" text %s\n", st ? "same" : "DIFF"); + if (!st) std::printf(" ref: %s\n got: %s\n", ref_text.c_str(), text.c_str()); + ++n; + } + std::printf("\n%d utterances: text identical %d, token ids identical %d; worst mel max|d| %.2e; encoder over %d " + "dumps: worst mean frame cosine %.5f, worst min frame cosine %.5f, worst relative L2 %.4f\n", + n, same_text, same_ids, worst_mel, n_enc, worst_mean_cos, worst_cos, worst_l2); + const double ref_wer = wer_words ? (double) wer_err / (double) wer_words : 0.0; + std::printf("WER against the reference transcripts: %.4f (%zu errors / %zu words)\n", ref_wer, wer_err, wer_words); + if (n == 0 || wer_words == 0) { + std::printf("FAIL: no fixtures with reference transcripts were compared\n"); + ++fails; + } + if (ref_wer > gate_wer) { + std::printf("FAIL: WER against the reference %.4f > %.4f\n", ref_wer, gate_wer); + ++fails; + } + if (worst_mel > 2e-3) { + std::printf("FAIL: mel differs by %.2e > 2e-3\n", worst_mel); + ++fails; + } + if (n_enc && worst_mean_cos < gate_cos) { + std::printf("FAIL: encoder mean frame cosine %.5f < %.5f\n", worst_mean_cos, gate_cos); + ++fails; + } + if (n_enc && worst_l2 > gate_l2) { + std::printf("FAIL: encoder relative L2 %.4f > %.4f\n", worst_l2, gate_l2); + ++fails; + } + if (n && same_text < gate_text * n) { + std::printf("FAIL: %d/%d transcripts identical < %.0f%%\n", same_text, n, gate_text * 100); + ++fails; + } + pk_model_free(m); + std::printf(fails ? "PARITY FAILED\n" : "PARITY OK\n"); + return fails ? 1 : 0; +} diff --git a/tools/parakeet/tools/bench_parakeet.py b/tools/parakeet/tools/bench_parakeet.py new file mode 100644 index 000000000..140942ac2 --- /dev/null +++ b/tools/parakeet/tools/bench_parakeet.py @@ -0,0 +1,142 @@ +#!/usr/bin/env python3 +# SPDX-License-Identifier: MIT +""" +bench_parakeet.py - run parakeet-bench and score it: RTFx (wall and CPU time), +first-window partial latency, peak RSS, WER against the references and the +share of transcripts identical to reference_transformers.py. + + bench_parakeet.py --bench build/bin/parakeet-bench --model phonon-2-fv5.gguf \ + --rows rows.json [--rows cmd_rows.json] [--reference ref/manifest.json] \ + [--threads 1 4] [--stream] [--out summary.json] [--gate-wer 5.0] + +rows.json: [{id, path, ref, split}] (the same rows dump_parakeet_reference.py +reads). WER uses whisper_normalizer's EnglishTextNormalizer when installed +(the Open ASR Leaderboard normaliser), otherwise lower-case + strip +punctuation. CPU-time RTFx is the robust number on a shared host; wall-clock +RTFx is only meaningful on an idle machine. Host numbers are not device +evidence. +""" +from __future__ import annotations + +import argparse +import json +import re +import statistics as st +import subprocess +import sys +import tempfile +from pathlib import Path +from typing import Any + + +def normaliser(): + try: + from whisper_normalizer.english import EnglishTextNormalizer # type: ignore + n = EnglishTextNormalizer() + return n, "whisper_normalizer.EnglishTextNormalizer" + except ImportError: + def n(s: str) -> str: + s = s.lower().replace("-", " ") + s = re.sub(r"[^a-z0-9' ]+", " ", s) + return " ".join(s.split()) + return n, "lowercase+strip-punctuation" + + +def wer(refs: list[str], hyps: list[str]) -> float: + errs = words = 0 + for r, h in zip(refs, hyps): + rw, hw = r.split(), h.split() + d = list(range(len(hw) + 1)) + for i in range(1, len(rw) + 1): + prev, d[0] = d[0], i + for j in range(1, len(hw) + 1): + cur = d[j] + d[j] = min(d[j] + 1, d[j - 1] + 1, prev + (rw[i - 1] != hw[j - 1])) + prev = cur + errs += d[len(hw)] + words += len(rw) + return 100.0 * errs / max(1, words) + + +def run(bench: str, model: str, rows: list[dict], threads: int, stream: bool, extra: list[str]) -> list[dict]: + with tempfile.NamedTemporaryFile("w", suffix=".tsv", delete=False) as fh: + for r in rows: + fh.write(f"{r['id']}\t{r['path']}\n") + lst = fh.name + cmd = [bench, "-m", model, "-t", str(threads), "--list", lst] + (["--stream"] if stream else []) + extra + out = subprocess.run(cmd, check=True, stdout=subprocess.PIPE, text=True).stdout + return [json.loads(l) for l in out.splitlines() if l.strip()] + + +def main(): + ap = argparse.ArgumentParser() + ap.add_argument("--bench", required=True) + ap.add_argument("--model", required=True) + ap.add_argument("--rows", action="append", required=True) + ap.add_argument("--reference", action="append", default=[], help="manifest.json from dump_parakeet_reference.py") + ap.add_argument("--threads", type=int, nargs="+", default=[1]) + ap.add_argument("--stream", action="store_true") + ap.add_argument("--extra", default="", help="extra parakeet-bench arguments") + ap.add_argument("--out") + ap.add_argument("--gate-wer", type=float, default=None, help="fail when the overall WER exceeds this") + ap.add_argument("--gate-ref-identical", type=float, default=None, + help="fail when fewer than this fraction of transcripts equal the reference") + a = ap.parse_args() + norm, norm_name = normaliser() + rows = [] + for p in a.rows: + rows += json.loads(Path(p).read_text()) + ref_text = {} + for p in a.reference: + for r in json.loads(Path(p).read_text())["rows"]: + ref_text[r["id"]] = r["text"] + summary: dict[str, Any] = {"normaliser": norm_name, "n": len(rows), + "host_note": "host CPU numbers; not device evidence", "runs": {}} + ok = True + for th in a.threads: + res = run(a.bench, a.model, rows, th, a.stream, a.extra.split()) + tail = res[-1] + per = {r["id"]: r for r in res if not r.get("summary")} + byid = {r["id"]: r for r in rows} + splits = sorted({r.get("split", "all") for r in rows}) + s: dict[str, Any] = {k: tail[k] for k in ("kernel", "threads", "audio_s", "wall_s", "cpu_s", "rtfx_wall", "rtfx_cpu", + "peak_rss_mb", "load_ms", "loadavg")} + refs = [norm(byid[i]["ref"]) for i in per] + hyps = [norm(per[i]["text"]) for i in per] + s["wer"] = round(wer(refs, hyps), 2) + for sp in splits: + ids = [i for i in per if byid[i].get("split", "all") == sp] + s[f"wer_{sp}"] = round(wer([norm(byid[i]["ref"]) for i in ids], [norm(per[i]["text"]) for i in ids]), 2) + s["prefix_cpu_ms_median"] = round(st.median(r["prefix_cpu_ms"] for r in per.values()), 1) + s["prefix_wall_ms_median"] = round(st.median(r["prefix_wall_ms"] for r in per.values()), 1) + s["utt_cpu_ms_median"] = round(st.median(r["cpu_ms"] for r in per.values()), 1) + if ref_text: + ids = [i for i in per if i in ref_text] + same = sum(per[i]["text"] == ref_text[i] for i in ids) + s["identical_to_reference"] = f"{same}/{len(ids)}" + s["reference_wer"] = round(wer([norm(byid[i]["ref"]) for i in ids], [norm(ref_text[i]) for i in ids]), 2) + if a.gate_ref_identical is not None and same < a.gate_ref_identical * len(ids): + ok = False + if a.stream: + s["stream_wer"] = round(wer(refs, [norm(per[i]["stream_text"]) for i in per]), 2) + same = sum(per[i]["stream_text"] == per[i]["text"] for i in per) + s["stream_equals_batch"] = f"{same}/{len(per)}" + s["stream_first_partial_ms_median"] = round( + st.median(r["stream_first_partial_ms"] for r in per.values() if r["stream_first_partial_ms"] >= 0), 1) + s["stream_cpu_rtfx"] = round(sum(r["audio_s"] for r in per.values()) * 1e3 / + sum(r["stream_cpu_ms"] for r in per.values()), 2) + s["stream_finish_ms_median"] = round(st.median(r["stream_finish_ms"] for r in per.values()), 1) + if a.gate_wer is not None and s["wer"] > a.gate_wer: + ok = False + summary["runs"][f"t{th}"] = s + summary.setdefault("rows", {})[f"t{th}"] = res + print(json.dumps(s), flush=True) + if a.out: + Path(a.out).write_text(json.dumps(summary, indent=1)) + if not ok: + print("GATE FAILED", file=sys.stderr) + sys.exit(1) + + +if __name__ == "__main__": + main() diff --git a/tools/parakeet/tools/dump_parakeet_reference.py b/tools/parakeet/tools/dump_parakeet_reference.py new file mode 100644 index 000000000..2cf520079 --- /dev/null +++ b/tools/parakeet/tools/dump_parakeet_reference.py @@ -0,0 +1,128 @@ +#!/usr/bin/env python3 +# SPDX-License-Identifier: MIT +""" +dump_parakeet_reference.py - write the reference outputs the tools/parakeet +parity tests compare against. + +The reference is Fermion's own `reference_transformers.py` (dense fp32 +`transformers.ParakeetForTDT`, every container record expanded exactly) from +the pinned FermionResearch/Phonon-2 snapshot, with the processor of the base +repo nvidia/parakeet-tdt-0.6b-v3. Both scripts are checked against +tools/parakeet/pins.json before they are imported. + +For every utterance it writes, into --out-dir: + .mel.f32 normalised log-mel features [n_frames, 128] (valid frames only) + .sub.f32 subsampling output [T, 1024] + .l0.f32 encoder layer 0 output [T, 1024] + .enc.f32 encoder output [T, 1024] + .pcm.f32 the 16 kHz mono input +and manifest.json (shapes, token ids without blanks, text) plus the same rows +as manifest.tsv for the C++ test: + TAB TAB TAB TAB TAB + +usage: + dump_parakeet_reference.py --snapshot --base-dir \ + --container model.fermion --utts rows.json [--n 8] [--text-only] --out-dir fixtures/ +""" +from __future__ import annotations + +import argparse +import hashlib +import json +import sys +from pathlib import Path + +import numpy as np + +_HERE = Path(__file__).resolve().parent +PINS = json.loads((_HERE.parent / "pins.json").read_text()) + + +def check(path: Path, pin: dict): + data = path.read_bytes() + if len(data) != pin["size"] or hashlib.sha256(data).hexdigest() != pin["sha256"]: + sys.exit(f"{path}: does not match pins.json") + + +def load_pcm(path: str) -> np.ndarray: + import soundfile as sf + w, sr = sf.read(path, dtype="float32", always_2d=True) + w = w.mean(axis=1) if w.shape[1] > 1 else w[:, 0] + if sr != 16000: + import librosa + w = librosa.resample(w, orig_sr=sr, target_sr=16000).astype(np.float32) + return np.ascontiguousarray(w, dtype=np.float32) + + +def write_tsv(path: Path, rows: list[dict]): + with open(path, "w") as fh: + for r in rows: + text = r["text"].replace("\t", " ").replace("\n", " ") + fh.write(f"{r['id']}\t{r['n_samples']}\t{r['n_mel']}\t{r['T']}\t" + f"{','.join(str(i) for i in r['token_ids'])}\t{text}\n") + + +def main(): + ap = argparse.ArgumentParser() + ap.add_argument("--snapshot", type=Path, required=True) + ap.add_argument("--base-dir", required=True) + ap.add_argument("--container", type=Path, required=True) + ap.add_argument("--utts", type=Path, required=True, help="JSON list of {id, path, ref}") + ap.add_argument("--n", type=int, default=0, help="first N utterances (0 = all)") + ap.add_argument("--text-only", action="store_true", help="write transcripts only, no tensors") + ap.add_argument("--out-dir", type=Path, required=True) + a = ap.parse_args() + for f in ("reference_transformers.py", "fermion_container.py"): + check(a.snapshot / f, PINS["files"][f]) + check(a.container, PINS["archive_members"]["model.fermion"]) + sys.path.insert(0, str(a.snapshot)) + import torch + import reference_transformers as rt # type: ignore # noqa: E402 (from the pinned snapshot) + + torch.set_num_threads(4) + model, processor, info = rt.load_model(str(a.container), a.base_dir) + utts = json.loads(a.utts.read_text()) + if a.n: + utts = utts[: a.n] + a.out_dir.mkdir(parents=True, exist_ok=True) + caps = {} + hooks = [model.encoder.subsampling.register_forward_hook(lambda m, i, o: caps.__setitem__("sub", o)), + model.encoder.layers[0].register_forward_hook(lambda m, i, o: caps.__setitem__("l0", o))] + rows = [] + for u in utts: + pcm = load_pcm(u["path"]) + inp = processor([pcm], sampling_rate=16000, return_tensors="pt", padding=True) + feats = inp["input_features"] + mask = inp["attention_mask"] + n_valid = int(mask.sum()) + with torch.no_grad(): + enc = model.encoder(input_features=feats, attention_mask=mask) + out = model.generate(input_features=feats, attention_mask=mask) + seq = getattr(out, "sequences", out)[0].tolist() + text = processor.batch_decode(getattr(out, "sequences", out), skip_special_tokens=True)[0].strip() + blank = model.config.blank_token_id + ids = [int(x) for x in seq if int(x) != blank][0:] + # the generate output starts with decoder_start_token_id (= blank) and keeps blank steps + T = int(model._get_subsampling_output_length(torch.tensor([n_valid])).item()) + row = {"id": u["id"], "ref": u.get("ref"), "text": text, "token_ids": ids, + "n_samples": int(pcm.size), "n_mel": n_valid, "T": T, "split": u.get("split")} + if not a.text_only: + pcm.tofile(a.out_dir / f"{u['id']}.pcm.f32") + feats[0, :n_valid].float().numpy().astype(np.float32).tofile(a.out_dir / f"{u['id']}.mel.f32") + caps["sub"][0, :T].float().numpy().astype(np.float32).tofile(a.out_dir / f"{u['id']}.sub.f32") + l0 = caps["l0"][0] if isinstance(caps["l0"], torch.Tensor) else caps["l0"][0][0] + l0[:T].float().numpy().astype(np.float32).tofile(a.out_dir / f"{u['id']}.l0.f32") + enc.last_hidden_state[0, :T].float().numpy().astype(np.float32).tofile(a.out_dir / f"{u['id']}.enc.f32") + rows.append(row) + print(f"{u['id']:>20} T={T:4d} {text[:90]}", flush=True) + for h in hooks: + h.remove() + write_tsv(a.out_dir / "manifest.tsv", rows) + (a.out_dir / "manifest.json").write_text(json.dumps( + {"reference": "reference_transformers.py (dense fp32 ParakeetForTDT)", "load": { + k: v for k, v in info.items() if k in ("container_records", "params")}, + "torch": torch.__version__, "rows": rows}, indent=1)) + + +if __name__ == "__main__": + main() diff --git a/tools/parakeet/tools/parakeet-bench.cpp b/tools/parakeet/tools/parakeet-bench.cpp new file mode 100644 index 000000000..e748bdb10 --- /dev/null +++ b/tools/parakeet/tools/parakeet-bench.cpp @@ -0,0 +1,192 @@ +// tools/parakeet/tools/parakeet-bench.cpp - speed harness. +// +// parakeet-bench -m model.gguf -t N --list utts.tsv [--prefix-ms 1000] +// [--stream] [--chunk-ms 80] [--warmup 1] [--repeat R] > results.jsonl +// +// utts.tsv: one "\t" per line. For every utterance it prints a +// JSON line with the transcript, wall time, process CPU time (getrusage, so +// worker threads are included), the time to decode the first prefix-ms of +// audio (the first-window partial of a stream) and, with --stream, the +// streaming result fed in chunk-ms pieces. The last line is a summary with +// the peak RSS. With --repeat R every timing is the minimum of R runs, which +// filters out preemption on a shared host. Times are measured on the host this runs on; an Apple +// Silicon Mac is an ARM CPU proxy, not phone evidence. +#include "parakeet.h" +#include "pk-wav.h" + +#include +#include +#include +#include +#include +#include +#include +#include + +#include +#include +#if defined(__APPLE__) +#include +#include +#endif + +static double cpu_ms() { + struct rusage ru; + getrusage(RUSAGE_SELF, &ru); + return (ru.ru_utime.tv_sec + ru.ru_stime.tv_sec) * 1e3 + (ru.ru_utime.tv_usec + ru.ru_stime.tv_usec) / 1e3; +} + +static double peak_rss_mb() { + struct rusage ru; + getrusage(RUSAGE_SELF, &ru); +#if defined(__APPLE__) + return ru.ru_maxrss / (1024.0 * 1024.0); +#else + return ru.ru_maxrss / 1024.0; +#endif +} + +static std::string jesc(const std::string & s) { + std::string o; + for (char c : s) { + if (c == '"' || c == '\\') { o += '\\'; o += c; } + else if ((unsigned char) c < 0x20) { char b[8]; std::snprintf(b, sizeof(b), "\\u%04x", c); o += b; } + else o += c; + } + return o; +} + +int main(int argc, char ** argv) { +#if defined(__APPLE__) + pthread_set_qos_class_self_np(QOS_CLASS_USER_INTERACTIVE, 0); +#endif + std::string model, list; + pk_params pp = pk_default_params(); + int prefix_ms = 1000, chunk_ms = 80, warmup = 1, repeat = 1; + bool stream = false; + pk_stream_params sp = pk_stream_default_params(); + for (int i = 1; i < argc; ++i) { + const std::string a = argv[i]; + auto next = [&]() -> const char * { return i + 1 < argc ? argv[++i] : ""; }; + if (a == "-m") model = next(); + else if (a == "-t") pp.n_threads = std::atoi(next()); + else if (a == "--kernel") { const std::string k = next(); pp.kernel = k == "scalar" ? 1 : k == "dotprod" ? 2 : k == "i8mm" ? 3 : k == "avx2" ? 4 : 0; } + else if (a == "--no-pin") pp.pin_threads = 0; + else if (a == "--list") list = next(); + else if (a == "--prefix-ms") prefix_ms = std::atoi(next()); + else if (a == "--stream") stream = true; + else if (a == "--chunk-ms") chunk_ms = std::atoi(next()); + else if (a == "--warmup") warmup = std::atoi(next()); + else if (a == "--repeat") repeat = std::max(1, std::atoi(next())); + else if (a == "--right-ms") sp.right_context_ms = std::atoi(next()); + else if (a == "--hop-ms") sp.min_hop_ms = std::atoi(next()); + else if (a == "--max-window-ms") sp.max_window_ms = std::atoi(next()); + } + if (model.empty() || list.empty()) { + std::fprintf(stderr, "usage: %s -m model.gguf -t N --list utts.tsv [--stream]\n", argv[0]); + return 2; + } + char err[512] = { 0 }; + const auto tl = std::chrono::steady_clock::now(); + const double rss0 = peak_rss_mb(); + pk_model * m = pk_model_load(model.c_str(), pp, err, sizeof(err)); + if (!m) { std::fprintf(stderr, "load failed: %s\n", err); return 1; } + const double load_ms = std::chrono::duration(std::chrono::steady_clock::now() - tl).count(); + std::vector> utts; + { + std::ifstream in(list); + std::string line; + while (std::getline(in, line)) { + const size_t tab = line.find('\t'); + if (tab == std::string::npos) continue; + utts.emplace_back(line.substr(0, tab), line.substr(tab + 1)); + } + } + std::vector> audio(utts.size()); + for (size_t i = 0; i < utts.size(); ++i) { + int sr = 0; + std::string e; + if (!pkwav::read_wav(utts[i].second, audio[i], sr, e)) { std::fprintf(stderr, "%s\n", e.c_str()); return 1; } + audio[i] = pkwav::resample(audio[i], sr, 16000); + } + for (int w = 0; w < warmup && !audio.empty(); ++w) { + pk_result r; + pk_transcribe(m, audio[0].data(), audio[0].size(), nullptr, nullptr, &r); + pk_result_free(&r); + } + double tot_audio = 0, tot_wall = 0, tot_cpu = 0; + for (size_t i = 0; i < utts.size(); ++i) { + const auto & pcm = audio[i]; + const double audio_s = pcm.size() / 16000.0; + pk_result r{}; + double wall = 1e30, cpu = 1e30, pwall = 1e30, pcpu = 1e30; + double c0; + std::chrono::steady_clock::time_point t0; + const size_t np = std::min(pcm.size(), (size_t) prefix_ms * 16); + for (int rep = 0; rep < repeat; ++rep) { + if (rep) pk_result_free(&r); + c0 = cpu_ms(); + t0 = std::chrono::steady_clock::now(); + pk_transcribe(m, pcm.data(), pcm.size(), nullptr, nullptr, &r); + wall = std::min(wall, std::chrono::duration(std::chrono::steady_clock::now() - t0).count()); + cpu = std::min(cpu, cpu_ms() - c0); + // first-window partial: decode the first prefix_ms of audio + pk_result rp; + c0 = cpu_ms(); + t0 = std::chrono::steady_clock::now(); + pk_transcribe(m, pcm.data(), np, nullptr, nullptr, &rp); + pwall = std::min(pwall, std::chrono::duration(std::chrono::steady_clock::now() - t0).count()); + pcpu = std::min(pcpu, cpu_ms() - c0); + pk_result_free(&rp); + } + tot_audio += audio_s; tot_wall += wall; tot_cpu += cpu; + std::ostringstream js; + js << "{\"id\":\"" << jesc(utts[i].first) << "\",\"audio_s\":" << audio_s << ",\"wall_ms\":" << wall + << ",\"cpu_ms\":" << cpu << ",\"mel_ms\":" << r.t_mel_ms << ",\"enc_ms\":" << r.t_encoder_ms + << ",\"dec_ms\":" << r.t_decoder_ms << ",\"prefix_wall_ms\":" << pwall << ",\"prefix_cpu_ms\":" << pcpu + << ",\"text\":\"" << jesc(r.text ? r.text : "") << "\""; + pk_result_free(&r); + if (stream) { + pk_stream * s = pk_stream_open(m, sp); + const size_t step = (size_t) chunk_ms * 16; + std::vector buf(1 << 16); + c0 = cpu_ms(); + double first_wall = -1; + for (size_t p = 0; p < pcm.size(); p += step) { + pk_stream_feed(s, pcm.data() + p, std::min(step, pcm.size() - p)); + t0 = std::chrono::steady_clock::now(); + pk_stream_partial(s, buf.data(), buf.size(), nullptr, nullptr); + int nd = 0; + pk_stream_stats(s, &nd, nullptr, nullptr); + if (first_wall < 0 && nd > 0) { + first_wall = std::chrono::duration(std::chrono::steady_clock::now() - t0).count(); + } + } + t0 = std::chrono::steady_clock::now(); + pk_stream_finish(s, buf.data(), buf.size(), nullptr, nullptr); + const double fin = std::chrono::duration(std::chrono::steady_clock::now() - t0).count(); + const double scpu = cpu_ms() - c0; + int nd = 0; + double dms = 0, first = 0; + pk_stream_stats(s, &nd, &dms, &first); + js << ",\"stream_text\":\"" << jesc(buf.data()) << "\",\"stream_decodes\":" << nd + << ",\"stream_decode_ms\":" << dms << ",\"stream_cpu_ms\":" << scpu + << ",\"stream_first_partial_ms\":" << first_wall << ",\"stream_finish_ms\":" << fin; + pk_stream_close(s); + } + js << "}"; + std::printf("%s\n", js.str().c_str()); + std::fflush(stdout); + } + double la[3] = { 0, 0, 0 }; +#if defined(__APPLE__) || defined(__linux__) + getloadavg(la, 3); +#endif + std::printf("{\"summary\":true,\"kernel\":\"%s\",\"threads\":%d,\"n\":%zu,\"audio_s\":%.2f,\"wall_s\":%.3f," + "\"cpu_s\":%.3f,\"rtfx_wall\":%.2f,\"rtfx_cpu\":%.2f,\"load_ms\":%.0f,\"peak_rss_mb\":%.0f," + "\"rss_before_load_mb\":%.0f,\"loadavg\":%.1f}\n", + pk_model_kernel_name(m), pk_model_n_threads(m), utts.size(), tot_audio, tot_wall / 1e3, tot_cpu / 1e3, + tot_audio * 1e3 / tot_wall, tot_audio * 1e3 / tot_cpu, load_ms, peak_rss_mb(), rss0, la[0]); + pk_model_free(m); + return 0; +} diff --git a/tools/parakeet/tools/parakeet-cli.cpp b/tools/parakeet/tools/parakeet-cli.cpp new file mode 100644 index 000000000..70ee3c1ce --- /dev/null +++ b/tools/parakeet/tools/parakeet-cli.cpp @@ -0,0 +1,108 @@ +// tools/parakeet/tools/parakeet-cli.cpp - transcribe WAV files with a +// parakeet-tdt GGUF. +// +// parakeet-cli -m phonon-2-fv5.gguf [-t N] [--kernel scalar|dotprod|i8mm|avx2] +// [--stream] [--chunk-ms 80] [--words] file.wav... +#include "parakeet.h" +#include "pk-wav.h" + +#include +#include +#include +#include +#include +#include +#include + +int main(int argc, char ** argv) { + std::string model; + pk_params pp = pk_default_params(); + bool stream = false, words = false; + int chunk_ms = 80; + std::vector files; + for (int i = 1; i < argc; ++i) { + const std::string a = argv[i]; + if ((a == "-m" || a == "--model") && i + 1 < argc) model = argv[++i]; + else if ((a == "-t" || a == "--threads") && i + 1 < argc) pp.n_threads = std::atoi(argv[++i]); + else if (a == "--kernel" && i + 1 < argc) { + const std::string k = argv[++i]; + pp.kernel = k == "scalar" ? 1 : k == "dotprod" ? 2 : k == "i8mm" ? 3 : k == "avx2" ? 4 : 0; + } else if (a == "--stream") stream = true; + else if (a == "--words") words = true; + else if (a == "--chunk-ms" && i + 1 < argc) chunk_ms = std::atoi(argv[++i]); + else if (a == "-h" || a == "--help") { + std::printf("usage: %s -m model.gguf [-t N] [--kernel K] [--stream] [--words] file.wav...\n", argv[0]); + return 0; + } else files.push_back(a); + } + if (model.empty() || files.empty()) { + std::fprintf(stderr, "usage: %s -m model.gguf [-t N] file.wav...\n", argv[0]); + return 2; + } + char err[512] = { 0 }; + auto t0 = std::chrono::steady_clock::now(); + pk_model * m = pk_model_load(model.c_str(), pp, err, sizeof(err)); + if (!m) { + std::fprintf(stderr, "load failed: %s\n", err); + return 1; + } + std::fprintf(stderr, "loaded %s in %.0f ms (kernel %s, %d threads)\n", model.c_str(), + std::chrono::duration(std::chrono::steady_clock::now() - t0).count(), + pk_model_kernel_name(m), pk_model_n_threads(m)); + int rc = 0; + for (const auto & f : files) { + std::vector pcm; + int sr = 0; + std::string e; + if (!pkwav::read_wav(f, pcm, sr, e)) { + std::fprintf(stderr, "%s\n", e.c_str()); + rc = 1; + continue; + } + pcm = pkwav::resample(pcm, sr, 16000); + const double audio_s = pcm.size() / 16000.0; + if (stream) { + pk_stream * s = pk_stream_open(m, pk_stream_default_params()); + const size_t step = (size_t) chunk_ms * 16; + std::string last; + char buf[8192]; + for (size_t p = 0; p < pcm.size(); p += step) { + pk_stream_feed(s, pcm.data() + p, std::min(step, pcm.size() - p)); + if (pk_stream_partial(s, buf, sizeof(buf), nullptr, nullptr) >= 0 && last != buf) { + last = buf; + std::fprintf(stderr, " [%6.2fs] %s\n", (p + step) / 16000.0, buf); + } + } + pk_stream_finish(s, buf, sizeof(buf), nullptr, nullptr); + int nd = 0; + double dms = 0, first = 0; + pk_stream_stats(s, &nd, &dms, &first); + std::printf("%s\t%s\n", f.c_str(), buf); + std::fprintf(stderr, " stream: %d decodes, %.0f ms total, first partial %.0f ms, audio %.2fs\n", nd, dms, + first, audio_s); + pk_stream_close(s); + continue; + } + pk_result r; + t0 = std::chrono::steady_clock::now(); + const int st = pk_transcribe(m, pcm.data(), pcm.size(), nullptr, nullptr, &r); + const double ms = std::chrono::duration(std::chrono::steady_clock::now() - t0).count(); + if (st != PK_OK) { + std::fprintf(stderr, "%s: transcribe failed (%d)\n", f.c_str(), st); + rc = 1; + continue; + } + std::printf("%s\t%s\n", f.c_str(), r.text); + std::fprintf(stderr, " %.2fs audio, %.1f ms (mel %.1f, encoder %.1f, decoder %.1f), RTFx %.1f\n", audio_s, ms, + r.t_mel_ms, r.t_encoder_ms, r.t_decoder_ms, audio_s * 1000.0 / ms); + if (words) { + std::vector a(4096), b(4096); + size_t n = a.size(); + pk_result_words(m, &r, a.data(), b.data(), &n); + for (size_t i = 0; i < n; ++i) std::fprintf(stderr, " [%5d-%5d ms]\n", a[i], b[i]); + } + pk_result_free(&r); + } + pk_model_free(m); + return rc; +} diff --git a/tools/parakeet/tools/pk-wav.h b/tools/parakeet/tools/pk-wav.h new file mode 100644 index 000000000..3c4ceb15e --- /dev/null +++ b/tools/parakeet/tools/pk-wav.h @@ -0,0 +1,100 @@ +// tools/parakeet/tools/pk-wav.h - minimal RIFF/WAVE reader for the CLI tools +// (PCM16 / PCM32 / float32, any channel count, mixed to mono) plus a +// windowed-sinc resampler to 16 kHz. +#pragma once + +#include +#include +#include +#include +#include +#include +#include + +namespace pkwav { + +constexpr double kPi = 3.14159265358979323846; + +inline bool read_wav(const std::string & path, std::vector & out, int & sr, std::string & err) { + FILE * f = std::fopen(path.c_str(), "rb"); + if (!f) { + err = "cannot open " + path; + return false; + } + std::vector d; + uint8_t buf[65536]; + size_t n; + while ((n = std::fread(buf, 1, sizeof(buf), f)) > 0) d.insert(d.end(), buf, buf + n); + std::fclose(f); + auto u16 = [&](size_t o) { return (uint32_t) d[o] | ((uint32_t) d[o + 1] << 8); }; + auto u32 = [&](size_t o) { return u16(o) | (u16(o + 2) << 16); }; + if (d.size() < 12 || std::memcmp(d.data(), "RIFF", 4) || std::memcmp(d.data() + 8, "WAVE", 4)) { + err = path + ": not a RIFF/WAVE file"; + return false; + } + int fmt = 0, ch = 0, bits = 0; + size_t pos = 12, data_off = 0, data_len = 0; + while (pos + 8 <= d.size()) { + const uint32_t len = u32(pos + 4); + if (!std::memcmp(d.data() + pos, "fmt ", 4) && pos + 24 <= d.size()) { + fmt = (int) u16(pos + 8); + ch = (int) u16(pos + 10); + sr = (int) u32(pos + 12); + bits = (int) u16(pos + 22); + if (fmt == 0xFFFE && len >= 26) fmt = (int) u16(pos + 32); + } else if (!std::memcmp(d.data() + pos, "data", 4)) { + data_off = pos + 8; + data_len = std::min(len, d.size() - data_off); + break; + } + pos += 8 + len + (len & 1); + } + if (!data_off || ch <= 0) { + err = path + ": no fmt/data chunk"; + return false; + } + const int bps = bits / 8; + const size_t frames = data_len / (size_t) (bps * ch); + out.assign(frames, 0.0f); + for (size_t i = 0; i < frames; ++i) { + double s = 0.0; + for (int c = 0; c < ch; ++c) { + const uint8_t * p = d.data() + data_off + (i * ch + c) * bps; + if (fmt == 1 && bits == 16) s += (int16_t) (p[0] | (p[1] << 8)) / 32768.0; + else if (fmt == 1 && bits == 32) s += (int32_t) (p[0] | (p[1] << 8) | (p[2] << 16) | ((uint32_t) p[3] << 24)) / 2147483648.0; + else if (fmt == 1 && bits == 24) s += ((int32_t) ((p[0] << 8) | (p[1] << 16) | ((uint32_t) p[2] << 24)) >> 8) / 8388608.0; + else if (fmt == 3 && bits == 32) { float v; std::memcpy(&v, p, 4); s += v; } + else { err = path + ": unsupported sample format"; return false; } + } + out[i] = (float) (s / ch); + } + return true; +} + +// Kaiser-windowed sinc, 32 zero crossings; adequate for test tools. +inline std::vector resample(const std::vector & x, int sr_in, int sr_out) { + if (sr_in == sr_out) return x; + const double ratio = (double) sr_out / sr_in; + const double cutoff = std::min(1.0, ratio) * 0.97; + const int half = 32; + const size_t n_out = (size_t) std::floor(x.size() * ratio); + std::vector y(n_out); + for (size_t i = 0; i < n_out; ++i) { + const double t = i / ratio; + const long c = (long) std::floor(t); + double acc = 0.0, wsum = 0.0; + for (long k = c - half / cutoff; k <= c + half / cutoff; ++k) { + if (k < 0 || k >= (long) x.size()) continue; + const double u = (t - k) * cutoff; + const double sinc = u == 0.0 ? 1.0 : std::sin(kPi * u) / (kPi * u); + const double r = u / half; + const double w = std::fabs(r) >= 1.0 ? 0.0 : 0.5 + 0.5 * std::cos(kPi * r); + acc += x[k] * sinc * w * cutoff; + wsum += sinc * w * cutoff; + } + y[i] = (float) (wsum != 0.0 ? acc : 0.0); + } + return y; +} + +} // namespace pkwav diff --git a/ty.toml b/ty.toml index dbc0520e4..a63fd139d 100644 --- a/ty.toml +++ b/ty.toml @@ -27,6 +27,10 @@ include = [ # Same for the Paradee reference dumper (kokoro, misaki, espeakng_loader # and the paradee training scripts are pinned in its own environment). "./tools/paradee/tools/**", + # Same for the Parakeet reference dumper and bench scorer (torch, + # transformers, soundfile, librosa and whisper_normalizer live in the + # pinned reference env, see tools/parakeet/pins.json). + "./tools/parakeet/tools/**", # conversion/* defines mixins (e.g. _Qwen35MtpMixin) whose super() # chains resolve to attributes provided by the composed Model # subclasses in convert_hf_to_gguf.py. ty cannot see that