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
108 lines
2.6 KiB
C++
108 lines
2.6 KiB
C++
#pragma once
|
|
|
|
#include "engine/framework/assets/resource_bundle.h"
|
|
#include "engine/framework/audio/nemo_mel_frontend.h"
|
|
#include "engine/framework/tokenizers/sentencepiece.h"
|
|
|
|
#include <cstdint>
|
|
#include <filesystem>
|
|
#include <memory>
|
|
#include <optional>
|
|
#include <string>
|
|
#include <vector>
|
|
|
|
namespace engine::models::citrinet_asr {
|
|
|
|
struct JasperBlockConfig {
|
|
int64_t filters = 0;
|
|
int64_t repeat = 0;
|
|
int64_t kernel = 0;
|
|
int64_t stride = 1;
|
|
int64_t dilation = 1;
|
|
float dropout = 0.0f;
|
|
bool residual = false;
|
|
std::string residual_mode;
|
|
bool separable = false;
|
|
bool se = false;
|
|
int64_t se_reduction_ratio = 8;
|
|
};
|
|
|
|
struct CitrinetConfig {
|
|
int64_t sample_rate = 0;
|
|
int64_t n_mels = 0;
|
|
int64_t n_fft = 0;
|
|
int64_t hop_length = 0;
|
|
int64_t win_length = 0;
|
|
int64_t pad_to = 0;
|
|
int64_t vocab_size = 0;
|
|
int64_t num_classes = 0;
|
|
int64_t blank_id = 0;
|
|
int64_t output_stride = 1;
|
|
std::string normalize;
|
|
std::string window;
|
|
std::vector<JasperBlockConfig> jasper;
|
|
};
|
|
|
|
struct Conv1dWeights {
|
|
std::string weight_name;
|
|
std::vector<int64_t> weight_source_shape;
|
|
int64_t in_channels = 0;
|
|
int64_t out_channels = 0;
|
|
int64_t kernel = 0;
|
|
int64_t stride = 1;
|
|
int64_t dilation = 1;
|
|
int64_t padding = 0;
|
|
int64_t groups = 1;
|
|
bool use_bias = false;
|
|
std::optional<std::string> bias_name;
|
|
};
|
|
|
|
struct BatchNorm1dWeights {
|
|
int64_t channels = 0;
|
|
std::vector<float> weight;
|
|
std::vector<float> bias;
|
|
std::vector<float> running_mean;
|
|
std::vector<float> running_var;
|
|
};
|
|
|
|
struct ConvBn {
|
|
Conv1dWeights conv;
|
|
BatchNorm1dWeights bn;
|
|
};
|
|
|
|
struct SeparableConvBn {
|
|
Conv1dWeights depthwise;
|
|
Conv1dWeights pointwise;
|
|
BatchNorm1dWeights bn;
|
|
};
|
|
|
|
struct SqueezeExciteWeights {
|
|
Conv1dWeights fc1;
|
|
Conv1dWeights fc2;
|
|
};
|
|
|
|
struct JasperBlockWeights {
|
|
bool separable = false;
|
|
std::vector<SeparableConvBn> separable_repeats;
|
|
std::vector<ConvBn> conv_repeats;
|
|
bool has_residual = false;
|
|
Conv1dWeights residual_conv;
|
|
BatchNorm1dWeights residual_bn;
|
|
bool has_se = false;
|
|
SqueezeExciteWeights se;
|
|
};
|
|
|
|
struct CitrinetWeights {
|
|
CitrinetConfig config;
|
|
std::shared_ptr<const assets::TensorSource> source;
|
|
std::vector<float> window;
|
|
std::vector<float> fb;
|
|
std::shared_ptr<const engine::audio::NemoMelFrontend> frontend;
|
|
std::vector<JasperBlockWeights> blocks;
|
|
Conv1dWeights decoder;
|
|
std::vector<tokenizers::SentencePiecePiece> tokenizer_pieces;
|
|
};
|
|
|
|
std::shared_ptr<const CitrinetWeights> load_citrinet_weights_cached(const std::filesystem::path & model_path);
|
|
|
|
} // namespace engine::models::citrinet_asr
|