Fabio
c36510c21e
Some checks failed
Linux Build / Linux cpu (push) Has been cancelled
Linux Build / Linux cpu-clang (push) Has been cancelled
Linux Build / Linux vulkan (push) Has been cancelled
mac-build / macOS CPU (push) Has been cancelled
Nix Build / Nix (cpu) (push) Has been cancelled
Nix Build / Nix (vulkan) (push) Has been cancelled
windows-build / Windows CPU (push) Has been cancelled
89 lines
3.2 KiB
C++
89 lines
3.2 KiB
C++
#include "engine/community_models/kroko_asr/assets.h"
|
|
#include "engine/community_models/kroko_asr/encoder.h"
|
|
#include "engine/community_models/kroko_asr/zipformer.h"
|
|
#include "engine/framework/core/execution_context.h"
|
|
|
|
#include <cmath>
|
|
#include <exception>
|
|
#include <filesystem>
|
|
#include <fstream>
|
|
#include <iostream>
|
|
#include <vector>
|
|
|
|
int main(int argc, char ** argv) try {
|
|
if (argc != 3 && argc != 4) {
|
|
std::cerr
|
|
<< "usage: kroko_asr_streaming_probe <model-directory> "
|
|
"<output-f32.bin> [embedded-f32.bin]\n";
|
|
return 2;
|
|
}
|
|
const auto assets =
|
|
engine::models::kroko_asr::load_kroko_asr_assets(
|
|
std::filesystem::path(argv[1]));
|
|
engine::core::ExecutionContext execution(
|
|
engine::core::BackendConfig{
|
|
engine::core::BackendType::Cpu, 0, 4});
|
|
engine::models::kroko_asr::KrokoEncoderRuntime subsampling(
|
|
assets, execution);
|
|
engine::models::kroko_asr::KrokoZipformerRuntime zipformer(
|
|
assets, execution);
|
|
|
|
const int64_t chunk_size = assets->config.chunk_size;
|
|
const int64_t chunk_shift = assets->config.chunk_shift;
|
|
const int64_t total_frames = chunk_shift + chunk_size;
|
|
constexpr int64_t feature_dim = 80;
|
|
std::vector<float> features(
|
|
static_cast<size_t>(total_frames * feature_dim));
|
|
for (size_t index = 0; index < features.size(); ++index) {
|
|
features[index] =
|
|
std::sin(static_cast<float>(index) * 0.0013F) * 0.7F +
|
|
std::cos(static_cast<float>(index) * 0.0007F) * 0.2F;
|
|
}
|
|
|
|
std::ofstream output(
|
|
std::filesystem::path(argv[2]),
|
|
std::ios::binary | std::ios::trunc);
|
|
if (!output) {
|
|
throw std::runtime_error(
|
|
"failed to open Kroko streaming probe output");
|
|
}
|
|
std::ofstream embedded_output;
|
|
if (argc == 4) {
|
|
embedded_output.open(
|
|
std::filesystem::path(argv[3]),
|
|
std::ios::binary | std::ios::trunc);
|
|
if (!embedded_output) {
|
|
throw std::runtime_error(
|
|
"failed to open Kroko embedding probe output");
|
|
}
|
|
}
|
|
for (int64_t chunk = 0; chunk < 2; ++chunk) {
|
|
const int64_t offset = chunk * chunk_shift;
|
|
std::vector<float> input(
|
|
features.begin() + offset * feature_dim,
|
|
features.begin() + (offset + chunk_size) * feature_dim);
|
|
const auto embedded =
|
|
subsampling.encode_subsampled_chunk(input);
|
|
if (embedded_output) {
|
|
embedded_output.write(
|
|
reinterpret_cast<const char *>(
|
|
embedded.values.data()),
|
|
static_cast<std::streamsize>(
|
|
embedded.values.size() * sizeof(float)));
|
|
}
|
|
const auto encoded =
|
|
zipformer.encode_chunk(embedded.values);
|
|
output.write(
|
|
reinterpret_cast<const char *>(encoded.values.data()),
|
|
static_cast<std::streamsize>(
|
|
encoded.values.size() * sizeof(float)));
|
|
}
|
|
std::cout << "{\"chunks\":2,\"frames_per_chunk\":"
|
|
<< chunk_shift / 4 << ','
|
|
<< "\"dims\":512}\n";
|
|
return 0;
|
|
} catch (const std::exception & error) {
|
|
std::cerr << "kroko_asr_streaming_probe failed: "
|
|
<< error.what() << '\n';
|
|
return 1;
|
|
}
|