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
99 lines
3.4 KiB
C++
99 lines
3.4 KiB
C++
#pragma once
|
|
|
|
#include "engine/framework/core/execution_context.h"
|
|
#include "engine/framework/runtime/session.h"
|
|
#include "engine/community_models/parakeet_tdt/assets.h"
|
|
#include "engine/community_models/parakeet_tdt/encoder.h"
|
|
#include "engine/community_models/parakeet_tdt/weights.h"
|
|
#include "engine/community_models/parakeet_tdt/hotwords.h"
|
|
|
|
#include <cstdint>
|
|
#include <array>
|
|
#include <functional>
|
|
#include <memory>
|
|
#include <string>
|
|
#include <vector>
|
|
|
|
namespace engine::community_models::parakeet_tdt {
|
|
|
|
struct ParakeetDecodeOptions {
|
|
int64_t max_tokens = 0;
|
|
bool keep_language_tags = false;
|
|
int64_t audio_end_sample = -1;
|
|
std::shared_ptr<const PhononHotwordAutomaton> hotwords;
|
|
};
|
|
|
|
struct ParakeetDecodedText {
|
|
std::string text;
|
|
std::vector<int32_t> token_ids;
|
|
std::vector<int32_t> token_frame_indices;
|
|
std::vector<int32_t> durations;
|
|
std::vector<runtime::WordTimestamp> word_timestamps;
|
|
};
|
|
|
|
class ParakeetTDTDecoderRuntime {
|
|
public:
|
|
ParakeetTDTDecoderRuntime(
|
|
std::shared_ptr<const ParakeetTDTAssets> assets,
|
|
std::shared_ptr<const ParakeetWeights> weights,
|
|
engine::core::ExecutionContext & execution_context,
|
|
size_t graph_arena_bytes);
|
|
~ParakeetTDTDecoderRuntime();
|
|
|
|
void prepare();
|
|
|
|
ParakeetDecodedText decode(
|
|
const ParakeetEncodedAudio & encoded,
|
|
const ParakeetDecodeOptions & options);
|
|
|
|
void reset_state();
|
|
ParakeetDecodedText decode_incremental(
|
|
const ParakeetEncodedAudio & encoded,
|
|
const ParakeetDecodeOptions & options,
|
|
int64_t frame_offset = 0);
|
|
ParakeetDecodedText format_tokens(
|
|
std::vector<int32_t> token_ids,
|
|
std::vector<int32_t> token_frame_indices,
|
|
std::vector<int32_t> durations,
|
|
const ParakeetDecodeOptions & options,
|
|
int64_t audio_end_frame) const;
|
|
|
|
int32_t run_step(int32_t input_token, const float * encoder_frame, bool decoder_cache_valid, int32_t * out_dur_id = nullptr);
|
|
int32_t run_joint_step(const float * encoder_frame, int32_t * out_dur_id = nullptr);
|
|
|
|
private:
|
|
struct StepGraph;
|
|
struct JointGraph;
|
|
|
|
std::shared_ptr<const ParakeetTDTAssets> assets_;
|
|
std::shared_ptr<const ParakeetWeights> weights_;
|
|
engine::core::ExecutionContext * execution_context_ = nullptr;
|
|
size_t graph_arena_bytes_ = 0;
|
|
std::array<int64_t, 2> state_widths_{};
|
|
std::unique_ptr<StepGraph> step_graph_;
|
|
std::unique_ptr<JointGraph> joint_graph_;
|
|
std::vector<float> logits_scratch_;
|
|
std::vector<float> hidden_scratch_;
|
|
std::vector<float> cell_scratch_;
|
|
std::vector<float> decoder_cache_scratch_;
|
|
std::vector<float> hidden_read_scratch_;
|
|
std::vector<float> cell_read_scratch_;
|
|
int32_t pending_input_token_ = 0;
|
|
bool predictor_cache_valid_ = false;
|
|
bool state_initialized_ = false;
|
|
std::shared_ptr<const PhononHotwordAutomaton> hotwords_;
|
|
int32_t hotword_state_ = 0;
|
|
int32_t select_token();
|
|
|
|
void ensure_step_graph();
|
|
void ensure_joint_graph();
|
|
std::string decode_text(const std::vector<int32_t> & token_ids, bool keep_language_tags) const;
|
|
std::vector<runtime::WordTimestamp> build_word_timestamps(
|
|
const std::vector<int32_t> & token_ids,
|
|
const std::vector<int32_t> & token_frame_indices,
|
|
const std::vector<int32_t> & durations,
|
|
int64_t audio_end_frame,
|
|
int64_t audio_end_sample) const;
|
|
};
|
|
|
|
} // namespace engine::community_models::parakeet_tdt
|