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
117 lines
3.8 KiB
C++
117 lines
3.8 KiB
C++
#pragma once
|
|
|
|
#include "engine/community_models/nemotron_transcribe/frontend.h"
|
|
#include "engine/framework/assets/tensor_source.h"
|
|
#include "engine/framework/core/execution_context.h"
|
|
|
|
#include <memory>
|
|
#include <functional>
|
|
#include <vector>
|
|
|
|
namespace engine::models::nemotron_transcribe {
|
|
|
|
struct RopeExpertConfig {
|
|
int feature_size = 128;
|
|
int hidden_size = 1280;
|
|
int heads = 10;
|
|
int layers = 32;
|
|
int intermediate_size = 5120;
|
|
int subsampling_factor = 8;
|
|
int rotary_dimensions = 64;
|
|
bool qk_norm = true;
|
|
bool qkv_bias = false;
|
|
bool pre_block_norm = true;
|
|
float layer_norm_epsilon = 1e-5f;
|
|
float rope_base = 10000.0f;
|
|
};
|
|
|
|
struct SpeakerCacheConfig {
|
|
// PEE online-inference overrides in NVIDIA's native reference. These differ
|
|
// from the standalone diarizer's training geometry (264/0/264).
|
|
int cache_frames = 188;
|
|
int fifo_frames = 40;
|
|
int update_period = 300;
|
|
int silence_frames_per_speaker = 1;
|
|
float prediction_score_threshold = 0.25f;
|
|
float latest_score_boost = 0.05f;
|
|
float strong_boost_rate = 0.75f;
|
|
float weak_boost_rate = 1.5f;
|
|
float min_positive_score_rate = 0.5f;
|
|
};
|
|
|
|
struct EncoderConfig {
|
|
RopeExpertConfig asr;
|
|
RopeExpertConfig diar = [] {
|
|
RopeExpertConfig c;
|
|
c.hidden_size = 512;
|
|
c.heads = 8;
|
|
c.layers = 31;
|
|
c.intermediate_size = 2048;
|
|
c.qk_norm = false;
|
|
return c;
|
|
}();
|
|
int projection_dim = 2688;
|
|
int speaker_count = 8;
|
|
int speaker_hidden_size = 192;
|
|
int upsample_factor = 8;
|
|
float speaker_activity_threshold = 0.5f;
|
|
// Output frames, each 80 ms with the released model. Export metadata owns
|
|
// these values; do not substitute the separate HF decoder chunk setting.
|
|
int online_inference_length = 500;
|
|
int chunk_left_context = 50;
|
|
int chunk_right_context = 50;
|
|
SpeakerCacheConfig speaker_cache;
|
|
bool use_flash_attention = true;
|
|
// Callers may lower this bound; the released checkpoint supports at most one hour.
|
|
double maximum_audio_seconds = 3600.0;
|
|
};
|
|
|
|
struct EncoderWindow {
|
|
int64_t mel_begin = 0;
|
|
int64_t mel_end = 0;
|
|
int64_t left_drop = 0;
|
|
int64_t core_frames = 0;
|
|
int64_t right_mel_frames = 0;
|
|
};
|
|
|
|
struct EncoderOutput {
|
|
int64_t frames = 0;
|
|
int64_t native_speaker_frames = 0;
|
|
int embedding_dim = 0;
|
|
int asr_dim = 0;
|
|
int speaker_count = 0;
|
|
double duration_seconds = 0.0;
|
|
// Row-major [frames, embedding_dim], [frames, asr_dim], [frames, speakers].
|
|
std::vector<float> audio_embeddings;
|
|
std::vector<float> asr_states;
|
|
std::vector<float> speaker_probs80;
|
|
// Native 10 ms grid. Padded subframes are excluded from this public result.
|
|
std::vector<float> speaker_probs10;
|
|
};
|
|
|
|
void validate_encoder_config(const EncoderConfig & config);
|
|
std::vector<EncoderWindow> plan_encoder_windows(int64_t mel_frames, const EncoderConfig & config);
|
|
|
|
class ParallelExpertEncoder {
|
|
public:
|
|
// Source is read only while constructing/uploading weights. Execution must
|
|
// outlive the encoder; its backend and graph resources stay framework-owned.
|
|
ParallelExpertEncoder(
|
|
const core::ExecutionContext & execution,
|
|
const assets::TensorSource & source,
|
|
EncoderConfig config = {},
|
|
assets::TensorStorageType storage = assets::TensorStorageType::Native);
|
|
~ParallelExpertEncoder();
|
|
ParallelExpertEncoder(const ParallelExpertEncoder &) = delete;
|
|
ParallelExpertEncoder & operator=(const ParallelExpertEncoder &) = delete;
|
|
|
|
// Cancellation is observed between graph launches and before every window.
|
|
EncoderOutput encode(const MelFeatures & features, const std::function<bool()> & should_stop = {});
|
|
const EncoderConfig & config() const noexcept;
|
|
|
|
private:
|
|
struct Impl;
|
|
std::unique_ptr<Impl> impl_;
|
|
};
|
|
|
|
} // namespace engine::models::nemotron_transcribe
|