240 lines
13 KiB
C++
240 lines
13 KiB
C++
#pragma once
|
|
|
|
#include "ggml.h"
|
|
|
|
#include <map>
|
|
|
|
enum parakeet_tensor {
|
|
// Encoder pre_encode
|
|
PARAKEET_TENSOR_ENC_PRE_OUT_WEIGHT,
|
|
PARAKEET_TENSOR_ENC_PRE_OUT_BIAS,
|
|
PARAKEET_TENSOR_ENC_PRE_CONV_0_WEIGHT,
|
|
PARAKEET_TENSOR_ENC_PRE_CONV_0_BIAS,
|
|
PARAKEET_TENSOR_ENC_PRE_CONV_2_WEIGHT,
|
|
PARAKEET_TENSOR_ENC_PRE_CONV_2_BIAS,
|
|
PARAKEET_TENSOR_ENC_PRE_CONV_3_WEIGHT,
|
|
PARAKEET_TENSOR_ENC_PRE_CONV_3_BIAS,
|
|
PARAKEET_TENSOR_ENC_PRE_CONV_5_WEIGHT,
|
|
PARAKEET_TENSOR_ENC_PRE_CONV_5_BIAS,
|
|
PARAKEET_TENSOR_ENC_PRE_CONV_6_WEIGHT,
|
|
PARAKEET_TENSOR_ENC_PRE_CONV_6_BIAS,
|
|
|
|
// Encoder layers (per-layer)
|
|
PARAKEET_TENSOR_ENC_NORM_FF1_WEIGHT,
|
|
PARAKEET_TENSOR_ENC_NORM_FF1_BIAS,
|
|
PARAKEET_TENSOR_ENC_FF1_LINEAR1_WEIGHT,
|
|
PARAKEET_TENSOR_ENC_FF1_LINEAR2_WEIGHT,
|
|
PARAKEET_TENSOR_ENC_NORM_CONV_WEIGHT,
|
|
PARAKEET_TENSOR_ENC_NORM_CONV_BIAS,
|
|
PARAKEET_TENSOR_ENC_CONV_PW1_WEIGHT,
|
|
PARAKEET_TENSOR_ENC_CONV_DW_WEIGHT,
|
|
PARAKEET_TENSOR_ENC_CONV_BN_WEIGHT,
|
|
PARAKEET_TENSOR_ENC_CONV_BN_BIAS,
|
|
PARAKEET_TENSOR_ENC_CONV_BN_MEAN,
|
|
PARAKEET_TENSOR_ENC_CONV_BN_VAR,
|
|
PARAKEET_TENSOR_ENC_CONV_BN_NUM_BATCHES,
|
|
PARAKEET_TENSOR_ENC_CONV_PW2_WEIGHT,
|
|
PARAKEET_TENSOR_ENC_NORM_ATTN_WEIGHT,
|
|
PARAKEET_TENSOR_ENC_NORM_ATTN_BIAS,
|
|
PARAKEET_TENSOR_ENC_ATTN_POS_BIAS_U,
|
|
PARAKEET_TENSOR_ENC_ATTN_POS_BIAS_V,
|
|
PARAKEET_TENSOR_ENC_ATTN_Q_WEIGHT,
|
|
PARAKEET_TENSOR_ENC_ATTN_K_WEIGHT,
|
|
PARAKEET_TENSOR_ENC_ATTN_V_WEIGHT,
|
|
PARAKEET_TENSOR_ENC_ATTN_OUT_WEIGHT,
|
|
PARAKEET_TENSOR_ENC_ATTN_POS_WEIGHT,
|
|
PARAKEET_TENSOR_ENC_NORM_FF2_WEIGHT,
|
|
PARAKEET_TENSOR_ENC_NORM_FF2_BIAS,
|
|
PARAKEET_TENSOR_ENC_FF2_LINEAR1_WEIGHT,
|
|
PARAKEET_TENSOR_ENC_FF2_LINEAR2_WEIGHT,
|
|
PARAKEET_TENSOR_ENC_NORM_OUT_WEIGHT,
|
|
PARAKEET_TENSOR_ENC_NORM_OUT_BIAS,
|
|
|
|
// Prediction network
|
|
PARAKEET_TENSOR_PRED_EMBED_WEIGHT,
|
|
PARAKEET_TENSOR_PRED_LSTM_WEIGHT_IH,
|
|
PARAKEET_TENSOR_PRED_LSTM_WEIGHT_HH,
|
|
PARAKEET_TENSOR_PRED_LSTM_BIAS_H,
|
|
|
|
// Joint network
|
|
PARAKEET_TENSOR_JOINT_PRED_WEIGHT,
|
|
PARAKEET_TENSOR_JOINT_PRED_BIAS,
|
|
PARAKEET_TENSOR_JOINT_ENC_WEIGHT,
|
|
PARAKEET_TENSOR_JOINT_ENC_BIAS,
|
|
PARAKEET_TENSOR_JOINT_NET_WEIGHT,
|
|
PARAKEET_TENSOR_JOINT_NET_BIAS,
|
|
};
|
|
|
|
enum parakeet_hparam {
|
|
PARAKEET_HPARAM_N_VOCAB,
|
|
PARAKEET_HPARAM_N_AUDIO_CTX,
|
|
PARAKEET_HPARAM_N_AUDIO_STATE,
|
|
PARAKEET_HPARAM_N_AUDIO_HEAD,
|
|
PARAKEET_HPARAM_N_AUDIO_LAYER,
|
|
PARAKEET_HPARAM_N_MELS,
|
|
PARAKEET_HPARAM_N_FFT,
|
|
PARAKEET_HPARAM_SUBSAMPLING_FACTOR,
|
|
PARAKEET_HPARAM_N_SUBSAMPLING_CHANNELS,
|
|
PARAKEET_HPARAM_N_CONV_KERNEL,
|
|
PARAKEET_HPARAM_N_PRED_DIM,
|
|
PARAKEET_HPARAM_N_PRED_LAYERS,
|
|
PARAKEET_HPARAM_N_TDT_DURATIONS,
|
|
PARAKEET_HPARAM_N_MAX_TOKENS,
|
|
};
|
|
|
|
static const std::map<parakeet_tensor, const char *> PARAKEET_TENSOR_NAMES = {
|
|
// Encoder pre_encode
|
|
{PARAKEET_TENSOR_ENC_PRE_OUT_WEIGHT, "encoder.pre_encode.out.weight"},
|
|
{PARAKEET_TENSOR_ENC_PRE_OUT_BIAS, "encoder.pre_encode.out.bias"},
|
|
{PARAKEET_TENSOR_ENC_PRE_CONV_0_WEIGHT, "encoder.pre_encode.conv.0.weight"},
|
|
{PARAKEET_TENSOR_ENC_PRE_CONV_0_BIAS, "encoder.pre_encode.conv.0.bias"},
|
|
{PARAKEET_TENSOR_ENC_PRE_CONV_2_WEIGHT, "encoder.pre_encode.conv.2.weight"},
|
|
{PARAKEET_TENSOR_ENC_PRE_CONV_2_BIAS, "encoder.pre_encode.conv.2.bias"},
|
|
{PARAKEET_TENSOR_ENC_PRE_CONV_3_WEIGHT, "encoder.pre_encode.conv.3.weight"},
|
|
{PARAKEET_TENSOR_ENC_PRE_CONV_3_BIAS, "encoder.pre_encode.conv.3.bias"},
|
|
{PARAKEET_TENSOR_ENC_PRE_CONV_5_WEIGHT, "encoder.pre_encode.conv.5.weight"},
|
|
{PARAKEET_TENSOR_ENC_PRE_CONV_5_BIAS, "encoder.pre_encode.conv.5.bias"},
|
|
{PARAKEET_TENSOR_ENC_PRE_CONV_6_WEIGHT, "encoder.pre_encode.conv.6.weight"},
|
|
{PARAKEET_TENSOR_ENC_PRE_CONV_6_BIAS, "encoder.pre_encode.conv.6.bias"},
|
|
|
|
// Encoder layers (use %d for layer number)
|
|
{PARAKEET_TENSOR_ENC_NORM_FF1_WEIGHT, "encoder.layers.%d.norm_feed_forward1.weight"},
|
|
{PARAKEET_TENSOR_ENC_NORM_FF1_BIAS, "encoder.layers.%d.norm_feed_forward1.bias"},
|
|
{PARAKEET_TENSOR_ENC_FF1_LINEAR1_WEIGHT, "encoder.layers.%d.feed_forward1.linear1.weight"},
|
|
{PARAKEET_TENSOR_ENC_FF1_LINEAR2_WEIGHT, "encoder.layers.%d.feed_forward1.linear2.weight"},
|
|
{PARAKEET_TENSOR_ENC_NORM_CONV_WEIGHT, "encoder.layers.%d.norm_conv.weight"},
|
|
{PARAKEET_TENSOR_ENC_NORM_CONV_BIAS, "encoder.layers.%d.norm_conv.bias"},
|
|
{PARAKEET_TENSOR_ENC_CONV_PW1_WEIGHT, "encoder.layers.%d.conv.pointwise_conv1.weight"},
|
|
{PARAKEET_TENSOR_ENC_CONV_DW_WEIGHT, "encoder.layers.%d.conv.depthwise_conv.weight"},
|
|
{PARAKEET_TENSOR_ENC_CONV_BN_WEIGHT, "encoder.layers.%d.conv.batch_norm.weight"},
|
|
{PARAKEET_TENSOR_ENC_CONV_BN_BIAS, "encoder.layers.%d.conv.batch_norm.bias"},
|
|
{PARAKEET_TENSOR_ENC_CONV_BN_MEAN, "encoder.layers.%d.conv.batch_norm.running_mean"},
|
|
{PARAKEET_TENSOR_ENC_CONV_BN_VAR, "encoder.layers.%d.conv.batch_norm.running_var"},
|
|
{PARAKEET_TENSOR_ENC_CONV_BN_NUM_BATCHES, "encoder.layers.%d.conv.batch_norm.num_batches_tracked"},
|
|
{PARAKEET_TENSOR_ENC_CONV_PW2_WEIGHT, "encoder.layers.%d.conv.pointwise_conv2.weight"},
|
|
{PARAKEET_TENSOR_ENC_NORM_ATTN_WEIGHT, "encoder.layers.%d.norm_self_att.weight"},
|
|
{PARAKEET_TENSOR_ENC_NORM_ATTN_BIAS, "encoder.layers.%d.norm_self_att.bias"},
|
|
{PARAKEET_TENSOR_ENC_ATTN_POS_BIAS_U, "encoder.layers.%d.self_attn.pos_bias_u"},
|
|
{PARAKEET_TENSOR_ENC_ATTN_POS_BIAS_V, "encoder.layers.%d.self_attn.pos_bias_v"},
|
|
{PARAKEET_TENSOR_ENC_ATTN_Q_WEIGHT, "encoder.layers.%d.self_attn.linear_q.weight"},
|
|
{PARAKEET_TENSOR_ENC_ATTN_K_WEIGHT, "encoder.layers.%d.self_attn.linear_k.weight"},
|
|
{PARAKEET_TENSOR_ENC_ATTN_V_WEIGHT, "encoder.layers.%d.self_attn.linear_v.weight"},
|
|
{PARAKEET_TENSOR_ENC_ATTN_OUT_WEIGHT, "encoder.layers.%d.self_attn.linear_out.weight"},
|
|
{PARAKEET_TENSOR_ENC_ATTN_POS_WEIGHT, "encoder.layers.%d.self_attn.linear_pos.weight"},
|
|
{PARAKEET_TENSOR_ENC_NORM_FF2_WEIGHT, "encoder.layers.%d.norm_feed_forward2.weight"},
|
|
{PARAKEET_TENSOR_ENC_NORM_FF2_BIAS, "encoder.layers.%d.norm_feed_forward2.bias"},
|
|
{PARAKEET_TENSOR_ENC_FF2_LINEAR1_WEIGHT, "encoder.layers.%d.feed_forward2.linear1.weight"},
|
|
{PARAKEET_TENSOR_ENC_FF2_LINEAR2_WEIGHT, "encoder.layers.%d.feed_forward2.linear2.weight"},
|
|
{PARAKEET_TENSOR_ENC_NORM_OUT_WEIGHT, "encoder.layers.%d.norm_out.weight"},
|
|
{PARAKEET_TENSOR_ENC_NORM_OUT_BIAS, "encoder.layers.%d.norm_out.bias"},
|
|
|
|
// Prediction network
|
|
{PARAKEET_TENSOR_PRED_EMBED_WEIGHT, "decoder.prediction.embed.weight"},
|
|
{PARAKEET_TENSOR_PRED_LSTM_WEIGHT_IH, "decoder.prediction.dec_rnn.lstm.weight_ih_l%d"},
|
|
{PARAKEET_TENSOR_PRED_LSTM_WEIGHT_HH, "decoder.prediction.dec_rnn.lstm.weight_hh_l%d"},
|
|
{PARAKEET_TENSOR_PRED_LSTM_BIAS_H, "decoder.prediction.dec_rnn.lstm.bias_h_l%d"},
|
|
|
|
// Joint network
|
|
{PARAKEET_TENSOR_JOINT_PRED_WEIGHT, "joint.pred.weight"},
|
|
{PARAKEET_TENSOR_JOINT_PRED_BIAS, "joint.pred.bias"},
|
|
{PARAKEET_TENSOR_JOINT_ENC_WEIGHT, "joint.enc.weight"},
|
|
{PARAKEET_TENSOR_JOINT_ENC_BIAS, "joint.enc.bias"},
|
|
{PARAKEET_TENSOR_JOINT_NET_WEIGHT, "joint.joint_net.2.weight"},
|
|
{PARAKEET_TENSOR_JOINT_NET_BIAS, "joint.joint_net.2.bias"},
|
|
};
|
|
|
|
static const std::map<parakeet_tensor, ggml_op> PARAKEET_TENSOR_INFO = {
|
|
// Encoder pre_encode
|
|
{PARAKEET_TENSOR_ENC_PRE_OUT_WEIGHT, GGML_OP_MUL_MAT},
|
|
{PARAKEET_TENSOR_ENC_PRE_OUT_BIAS, GGML_OP_ADD},
|
|
{PARAKEET_TENSOR_ENC_PRE_CONV_0_WEIGHT, GGML_OP_IM2COL},
|
|
{PARAKEET_TENSOR_ENC_PRE_CONV_0_BIAS, GGML_OP_ADD},
|
|
{PARAKEET_TENSOR_ENC_PRE_CONV_2_WEIGHT, GGML_OP_IM2COL},
|
|
{PARAKEET_TENSOR_ENC_PRE_CONV_2_BIAS, GGML_OP_ADD},
|
|
{PARAKEET_TENSOR_ENC_PRE_CONV_3_WEIGHT, GGML_OP_IM2COL},
|
|
{PARAKEET_TENSOR_ENC_PRE_CONV_3_BIAS, GGML_OP_ADD},
|
|
{PARAKEET_TENSOR_ENC_PRE_CONV_5_WEIGHT, GGML_OP_IM2COL},
|
|
{PARAKEET_TENSOR_ENC_PRE_CONV_5_BIAS, GGML_OP_ADD},
|
|
{PARAKEET_TENSOR_ENC_PRE_CONV_6_WEIGHT, GGML_OP_IM2COL},
|
|
{PARAKEET_TENSOR_ENC_PRE_CONV_6_BIAS, GGML_OP_ADD},
|
|
|
|
// Encoder layers
|
|
{PARAKEET_TENSOR_ENC_NORM_FF1_WEIGHT, GGML_OP_MUL},
|
|
{PARAKEET_TENSOR_ENC_NORM_FF1_BIAS, GGML_OP_ADD},
|
|
{PARAKEET_TENSOR_ENC_FF1_LINEAR1_WEIGHT, GGML_OP_MUL_MAT},
|
|
{PARAKEET_TENSOR_ENC_FF1_LINEAR2_WEIGHT, GGML_OP_MUL_MAT},
|
|
{PARAKEET_TENSOR_ENC_NORM_CONV_WEIGHT, GGML_OP_MUL},
|
|
{PARAKEET_TENSOR_ENC_NORM_CONV_BIAS, GGML_OP_ADD},
|
|
{PARAKEET_TENSOR_ENC_CONV_PW1_WEIGHT, GGML_OP_IM2COL},
|
|
{PARAKEET_TENSOR_ENC_CONV_DW_WEIGHT, GGML_OP_IM2COL},
|
|
{PARAKEET_TENSOR_ENC_CONV_BN_WEIGHT, GGML_OP_MUL},
|
|
{PARAKEET_TENSOR_ENC_CONV_BN_BIAS, GGML_OP_ADD},
|
|
{PARAKEET_TENSOR_ENC_CONV_BN_MEAN, GGML_OP_SUB},
|
|
{PARAKEET_TENSOR_ENC_CONV_BN_VAR, GGML_OP_DIV},
|
|
{PARAKEET_TENSOR_ENC_CONV_BN_NUM_BATCHES, GGML_OP_NONE},
|
|
{PARAKEET_TENSOR_ENC_CONV_PW2_WEIGHT, GGML_OP_IM2COL},
|
|
{PARAKEET_TENSOR_ENC_NORM_ATTN_WEIGHT, GGML_OP_MUL},
|
|
{PARAKEET_TENSOR_ENC_NORM_ATTN_BIAS, GGML_OP_ADD},
|
|
{PARAKEET_TENSOR_ENC_ATTN_POS_BIAS_U, GGML_OP_ADD},
|
|
{PARAKEET_TENSOR_ENC_ATTN_POS_BIAS_V, GGML_OP_ADD},
|
|
{PARAKEET_TENSOR_ENC_ATTN_Q_WEIGHT, GGML_OP_MUL_MAT},
|
|
{PARAKEET_TENSOR_ENC_ATTN_K_WEIGHT, GGML_OP_MUL_MAT},
|
|
{PARAKEET_TENSOR_ENC_ATTN_V_WEIGHT, GGML_OP_MUL_MAT},
|
|
{PARAKEET_TENSOR_ENC_ATTN_OUT_WEIGHT, GGML_OP_MUL_MAT},
|
|
{PARAKEET_TENSOR_ENC_ATTN_POS_WEIGHT, GGML_OP_MUL_MAT},
|
|
{PARAKEET_TENSOR_ENC_NORM_FF2_WEIGHT, GGML_OP_MUL},
|
|
{PARAKEET_TENSOR_ENC_NORM_FF2_BIAS, GGML_OP_ADD},
|
|
{PARAKEET_TENSOR_ENC_FF2_LINEAR1_WEIGHT, GGML_OP_MUL_MAT},
|
|
{PARAKEET_TENSOR_ENC_FF2_LINEAR2_WEIGHT, GGML_OP_MUL_MAT},
|
|
{PARAKEET_TENSOR_ENC_NORM_OUT_WEIGHT, GGML_OP_MUL},
|
|
{PARAKEET_TENSOR_ENC_NORM_OUT_BIAS, GGML_OP_ADD},
|
|
|
|
// Prediction network
|
|
{PARAKEET_TENSOR_PRED_EMBED_WEIGHT, GGML_OP_GET_ROWS},
|
|
{PARAKEET_TENSOR_PRED_LSTM_WEIGHT_IH, GGML_OP_MUL_MAT},
|
|
{PARAKEET_TENSOR_PRED_LSTM_WEIGHT_HH, GGML_OP_MUL_MAT},
|
|
{PARAKEET_TENSOR_PRED_LSTM_BIAS_H, GGML_OP_ADD},
|
|
|
|
// Joint network
|
|
{PARAKEET_TENSOR_JOINT_PRED_WEIGHT, GGML_OP_MUL_MAT},
|
|
{PARAKEET_TENSOR_JOINT_PRED_BIAS, GGML_OP_ADD},
|
|
{PARAKEET_TENSOR_JOINT_ENC_WEIGHT, GGML_OP_MUL_MAT},
|
|
{PARAKEET_TENSOR_JOINT_ENC_BIAS, GGML_OP_ADD},
|
|
{PARAKEET_TENSOR_JOINT_NET_WEIGHT, GGML_OP_MUL_MAT},
|
|
{PARAKEET_TENSOR_JOINT_NET_BIAS, GGML_OP_ADD},
|
|
};
|
|
|
|
static const std::map<parakeet_hparam, const char *> PARAKEET_HPARAM_NAMES = {
|
|
{PARAKEET_HPARAM_N_VOCAB, "n_vocab"},
|
|
{PARAKEET_HPARAM_N_AUDIO_CTX, "n_audio_ctx"},
|
|
{PARAKEET_HPARAM_N_AUDIO_STATE, "n_audio_state"},
|
|
{PARAKEET_HPARAM_N_AUDIO_HEAD, "n_audio_head"},
|
|
{PARAKEET_HPARAM_N_AUDIO_LAYER, "n_audio_layer"},
|
|
{PARAKEET_HPARAM_N_MELS, "n_mels"},
|
|
{PARAKEET_HPARAM_N_FFT, "n_fft"},
|
|
{PARAKEET_HPARAM_SUBSAMPLING_FACTOR, "subsampling_factor"},
|
|
{PARAKEET_HPARAM_N_SUBSAMPLING_CHANNELS, "n_subsampling_channels"},
|
|
{PARAKEET_HPARAM_N_CONV_KERNEL, "n_conv_kernel"},
|
|
{PARAKEET_HPARAM_N_PRED_DIM, "n_pred_dim"},
|
|
{PARAKEET_HPARAM_N_PRED_LAYERS, "n_pred_layers"},
|
|
{PARAKEET_HPARAM_N_TDT_DURATIONS, "n_tdt_durations"},
|
|
{PARAKEET_HPARAM_N_MAX_TOKENS, "n_max_tokens"},
|
|
};
|
|
|
|
static const std::map<parakeet_hparam, int32_t> PARAKEET_HPARAM_MODEL_VALUES = {
|
|
{PARAKEET_HPARAM_N_VOCAB, 8192},
|
|
{PARAKEET_HPARAM_N_AUDIO_CTX, 5000},
|
|
{PARAKEET_HPARAM_N_AUDIO_STATE, 1024},
|
|
{PARAKEET_HPARAM_N_AUDIO_HEAD, 8},
|
|
{PARAKEET_HPARAM_N_AUDIO_LAYER, 24},
|
|
{PARAKEET_HPARAM_N_MELS, 128},
|
|
{PARAKEET_HPARAM_N_FFT, 512},
|
|
{PARAKEET_HPARAM_SUBSAMPLING_FACTOR, 8},
|
|
{PARAKEET_HPARAM_N_SUBSAMPLING_CHANNELS, 256},
|
|
{PARAKEET_HPARAM_N_CONV_KERNEL, 9},
|
|
{PARAKEET_HPARAM_N_PRED_DIM, 640},
|
|
{PARAKEET_HPARAM_N_PRED_LAYERS, 2},
|
|
{PARAKEET_HPARAM_N_TDT_DURATIONS, 5},
|
|
{PARAKEET_HPARAM_N_MAX_TOKENS, 10},
|
|
};
|