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
115 lines
3.3 KiB
C++
115 lines
3.3 KiB
C++
#pragma once
|
|
|
|
#include "engine/framework/assets/resource_bundle.h"
|
|
#include "engine/framework/assets/tensor_source.h"
|
|
#include "engine/framework/audio/nemo_mel_frontend.h"
|
|
#include "engine/framework/core/backend_weight_store.h"
|
|
#include "engine/framework/modules/attention/types.h"
|
|
#include "engine/framework/modules/conv_modules.h"
|
|
#include "engine/framework/modules/linear_module.h"
|
|
#include "engine/framework/modules/norm_modules.h"
|
|
#include "engine/framework/modules/weight_binding.h"
|
|
|
|
#include <cstdint>
|
|
#include <filesystem>
|
|
#include <memory>
|
|
#include <unordered_map>
|
|
#include <vector>
|
|
|
|
namespace engine::models::nemotron_3_diar {
|
|
|
|
struct EncoderConfig {
|
|
int64_t feature_size = 128;
|
|
int64_t hidden_size = 512;
|
|
int64_t intermediate_size = 2048;
|
|
int64_t heads = 8;
|
|
int64_t layers = 31;
|
|
int64_t subsampling_factor = 8;
|
|
float rope_theta = 10000.0F;
|
|
float layer_norm_eps = 1.0e-5F;
|
|
};
|
|
|
|
struct HeadConfig {
|
|
int64_t hidden_size = 192;
|
|
int64_t upsample_factor = 8;
|
|
int64_t upsample_kernel_size = 3;
|
|
};
|
|
|
|
struct StreamingConfig {
|
|
int64_t spkcache_len = 264;
|
|
int64_t fifo_len = 0;
|
|
int64_t chunk_len = 264;
|
|
int64_t chunk_left_context = 0;
|
|
int64_t chunk_right_context = 0;
|
|
int64_t spkcache_update_period = 188;
|
|
int64_t spkcache_sil_frames_per_spk = 1;
|
|
float pred_score_threshold = 0.25F;
|
|
float scores_boost_latest = 0.05F;
|
|
float sil_threshold = 0.2F;
|
|
float strong_boost_rate = 0.75F;
|
|
float weak_boost_rate = 1.5F;
|
|
float min_pos_scores_rate = 0.5F;
|
|
int64_t max_index = 99999;
|
|
};
|
|
|
|
struct FeatureConfig {
|
|
int64_t sample_rate = 16000;
|
|
int64_t n_fft = 512;
|
|
int64_t win_length = 400;
|
|
int64_t hop_length = 160;
|
|
int64_t num_mel_bins = 128;
|
|
float preemphasis = 0.97F;
|
|
float dither = 1.0e-5F;
|
|
};
|
|
|
|
struct ModelConfig {
|
|
int64_t num_speakers = 8;
|
|
bool high_resolution = true;
|
|
int64_t output_subsampling_factor = 1;
|
|
EncoderConfig encoder;
|
|
HeadConfig head;
|
|
StreamingConfig streaming;
|
|
};
|
|
|
|
struct EncoderLayerWeights {
|
|
modules::NormWeights norm1;
|
|
modules::AttentionWeights attention;
|
|
modules::NormWeights norm2;
|
|
modules::LinearWeights ffn_in;
|
|
modules::LinearWeights ffn_out;
|
|
};
|
|
|
|
struct ModelWeights {
|
|
std::shared_ptr<core::BackendWeightStore> store;
|
|
modules::LinearWeights pre_encode;
|
|
modules::NormWeights embed_norm;
|
|
std::vector<EncoderLayerWeights> layers;
|
|
modules::NormWeights final_norm;
|
|
modules::LinearWeights encoder_projection;
|
|
modules::Conv1dWeights subpixel_upsample;
|
|
modules::LinearWeights head_hidden;
|
|
modules::LinearWeights speaker_head;
|
|
};
|
|
|
|
struct Assets {
|
|
assets::ResourceBundle resources;
|
|
ModelConfig model_config;
|
|
FeatureConfig feature_config;
|
|
std::vector<float> mel_filterbank;
|
|
std::vector<float> silence_embedding;
|
|
std::shared_ptr<const audio::NemoMelFrontend> frontend;
|
|
std::shared_ptr<const assets::TensorSource> model_weights;
|
|
};
|
|
|
|
std::shared_ptr<const Assets> load_assets(
|
|
const std::filesystem::path & model_root,
|
|
const std::filesystem::path & spec_path);
|
|
|
|
std::shared_ptr<ModelWeights> load_weights(
|
|
const Assets & assets,
|
|
ggml_backend_t backend,
|
|
core::BackendType backend_type,
|
|
assets::TensorStorageType storage_type,
|
|
size_t context_bytes);
|
|
|
|
} // namespace engine::models::nemotron_3_diar
|