#pragma once
#include "transcribe.h"
#include <cstdint>
#include <string>
#include <utility>
#include <vector>
struct gguf_context;
struct ggml_context;
struct ggml_tensor;
namespace transcribe::parakeet {
enum class HeadKind { TDT, RNNT, CTC };
struct ParakeetHParams {
HeadKind head_kind = HeadKind::TDT;
int32_t enc_n_layers = 0;
int32_t enc_d_model = 0;
int32_t enc_n_heads = 0;
int32_t enc_d_ff = 0; int32_t enc_conv_kernel = 0; int32_t enc_subsampling_factor = 0; int32_t enc_subsampling_channels = 0; int32_t enc_pos_emb_max_len = 0;
bool enc_use_bias = false;
bool enc_xscaling = false;
int32_t enc_att_context_left = -1;
int32_t enc_att_context_right = -1;
std::vector<std::pair<int32_t, int32_t>> enc_att_context_size_choices;
std::vector<int32_t> enc_att_chunk_left_choices;
std::vector<int32_t> enc_att_chunk_chunk_choices;
std::vector<int32_t> enc_att_chunk_right_choices;
enum class AttContextStyle { Regular, ChunkedLimited, ChunkedLimitedWithRc };
AttContextStyle enc_att_context_style = AttContextStyle::Regular;
int32_t enc_conv_context_left = -1;
int32_t enc_conv_context_right = -1;
enum class ConvNormType { BatchNorm, LayerNorm };
ConvNormType enc_conv_norm_type = ConvNormType::BatchNorm;
int32_t enc_stream_pre_encode_cache_size = 0;
int32_t enc_stream_drop_extra_pre_encoded = 0;
int32_t enc_stream_sampling_frames_first = 0;
int32_t pred_hidden = 0;
int32_t pred_n_layers = 0;
int32_t pred_vocab = 0;
int32_t joint_hidden = 0;
int32_t joint_num_extra_outputs = 0; std::string joint_activation;
std::vector<int32_t> tdt_durations;
int32_t tdt_max_symbols = 10;
std::string fe_type; int32_t fe_num_mels = 0;
int32_t fe_sample_rate = 0;
int32_t fe_n_fft = 0;
int32_t fe_win_length = 0; int32_t fe_hop_length = 0; std::string fe_window; std::string fe_normalize; float fe_dither = 0.0f;
float fe_pre_emphasis = 0.0f; float fe_f_min = 0.0f;
float fe_f_max = 0.0f;
int32_t head_ctc_n_classes = 0;
bool has_prompt = false;
int32_t prompt_num_prompts = 0;
int32_t prompt_hidden = 0;
std::string prompt_field;
std::string prompt_activation;
std::vector<std::string> prompt_dictionary_locales;
std::vector<int32_t> prompt_dictionary_indices;
int32_t prompt_auto_id = -1;
int32_t enc_head_dim() const { return enc_n_heads > 0 ? enc_d_model / enc_n_heads : 0; }
int32_t joint_n_classes() const { return (pred_vocab - 1) + joint_num_extra_outputs + 1; }
};
transcribe_status read_parakeet_hparams(const gguf_context * gguf,
ParakeetHParams & hp);
struct ParakeetPreEncode {
ggml_tensor * conv0_w = nullptr; ggml_tensor * conv0_b = nullptr; ggml_tensor * conv2_w = nullptr; ggml_tensor * conv2_b = nullptr; ggml_tensor * conv3_w = nullptr; ggml_tensor * conv3_b = nullptr; ggml_tensor * conv5_w = nullptr; ggml_tensor * conv5_b = nullptr; ggml_tensor * conv6_w = nullptr; ggml_tensor * conv6_b = nullptr; ggml_tensor * out_w = nullptr;
ggml_tensor * out_b = nullptr;
};
struct ParakeetBlock {
ggml_tensor * norm_ff1_w = nullptr; ggml_tensor * norm_ff1_b = nullptr; ggml_tensor * ff1_lin1_w = nullptr; ggml_tensor * ff1_lin1_b = nullptr; ggml_tensor * ff1_lin2_w = nullptr; ggml_tensor * ff1_lin2_b = nullptr;
ggml_tensor * norm_attn_w = nullptr; ggml_tensor * norm_attn_b = nullptr; ggml_tensor * attn_q_w = nullptr; ggml_tensor * attn_q_b = nullptr; ggml_tensor * attn_k_w = nullptr; ggml_tensor * attn_k_b = nullptr; ggml_tensor * attn_v_w = nullptr; ggml_tensor * attn_v_b = nullptr; ggml_tensor * attn_out_w = nullptr; ggml_tensor * attn_out_b = nullptr; ggml_tensor * attn_pos_w = nullptr; ggml_tensor * attn_pos_u = nullptr; ggml_tensor * attn_pos_v = nullptr;
ggml_tensor * norm_conv_w = nullptr; ggml_tensor * norm_conv_b = nullptr; ggml_tensor * conv_pw1_w = nullptr; ggml_tensor * conv_pw1_b = nullptr; ggml_tensor * conv_dw_w = nullptr; ggml_tensor * conv_dw_b = nullptr; ggml_tensor * conv_pw2_w = nullptr; ggml_tensor * conv_pw2_b = nullptr; ggml_tensor * conv_bn_w = nullptr; ggml_tensor * conv_bn_b = nullptr; ggml_tensor * conv_bn_rm = nullptr; ggml_tensor * conv_bn_rv = nullptr;
ggml_tensor * conv_bn_fused_scale = nullptr; ggml_tensor * conv_bn_fused_bias = nullptr;
ggml_tensor * norm_ff2_w = nullptr; ggml_tensor * norm_ff2_b = nullptr; ggml_tensor * ff2_lin1_w = nullptr; ggml_tensor * ff2_lin1_b = nullptr; ggml_tensor * ff2_lin2_w = nullptr; ggml_tensor * ff2_lin2_b = nullptr;
ggml_tensor * norm_out_w = nullptr; ggml_tensor * norm_out_b = nullptr; };
struct ParakeetPredictor {
ggml_tensor * embed_w = nullptr;
struct LstmLayer {
ggml_tensor * Wx = nullptr; ggml_tensor * Wh = nullptr; ggml_tensor * b = nullptr; };
std::vector<LstmLayer> lstm; };
struct ParakeetJoint {
ggml_tensor * enc_w = nullptr; ggml_tensor * enc_b = nullptr; ggml_tensor * pred_w = nullptr; ggml_tensor * pred_b = nullptr; ggml_tensor * out_w = nullptr; ggml_tensor * out_b = nullptr; };
struct ParakeetCtcHead {
ggml_tensor * weight = nullptr; ggml_tensor * bias = nullptr; };
struct ParakeetPromptMlp {
ggml_tensor * mlp0_w = nullptr; ggml_tensor * mlp0_b = nullptr; ggml_tensor * mlp2_w = nullptr; ggml_tensor * mlp2_b = nullptr; };
struct ParakeetWeights {
ParakeetPreEncode pre_encode;
std::vector<ParakeetBlock> blocks; ParakeetPredictor predictor; ParakeetJoint joint; ParakeetCtcHead ctc_head; ParakeetPromptMlp prompt; };
transcribe_status build_parakeet_weights(ggml_context * ctx_meta,
const ParakeetHParams & hp,
ParakeetWeights & weights);
}