#include "parakeet.h"
#include "parakeet-arch.h"

#include "ggml.h"
#include "ggml-cpp.h"
#include "ggml-alloc.h"
#include "ggml-backend.h"

#include <atomic>
#include <algorithm>
#include <cassert>
#include <cfloat>
#define _USE_MATH_DEFINES
#include <cmath>
#include <climits>
#include <cstdarg>
#include <cstdio>
#include <cstring>
#include <fstream>
#include <functional>
#include <cctype>
#include <map>
#include <random>
#include <set>
#include <string>
#include <thread>
#include <vector>

#ifdef _MSC_VER
#include <codecvt>
#endif

#if defined(PARAKEET_BIG_ENDIAN)
template<typename T>
static T byteswap(T value) {
    T value_swapped;
    char * source = reinterpret_cast<char *>(&value);
    char * target = reinterpret_cast<char *>(&value_swapped);
    int size = sizeof(T);
    for (int i = 0; i < size; i++) {
        target[size - 1 - i] = source[i];
    }
    return value_swapped;
}

template<typename T>
static void byteswap_tensor_data(ggml_tensor * tensor) {
    T * datum = reinterpret_cast<T *>(tensor->data);
    for (int i = 0; i < ggml_nelements(tensor); i++) {
        datum[i] = byteswap(datum[i]);
    }
}

static void byteswap_tensor(ggml_tensor * tensor) {
    switch (tensor->type) {
        case GGML_TYPE_I16: {
            byteswap_tensor_data<int16_t>(tensor);
            break;
        }
        case GGML_TYPE_F16: {
            byteswap_tensor_data<ggml_fp16_t>(tensor);
            break;
        }
        case GGML_TYPE_I32: {
            byteswap_tensor_data<int32_t>(tensor);
            break;
        }
        case GGML_TYPE_F32: {
            byteswap_tensor_data<float>(tensor);
            break;
        }
        default: { // GML_TYPE_I8
            break;
        }
    }
}

#define BYTESWAP_VALUE(d) d = byteswap(d)
#define BYTESWAP_FILTERS(f)           \
    do {                              \
        for (auto & datum : f.data) { \
            datum = byteswap(datum);  \
        }                             \
    } while (0)
#define BYTESWAP_TENSOR(t)  \
    do {                    \
        byteswap_tensor(t); \
    } while (0)
#else
#define BYTESWAP_VALUE(d) do {} while (0)
#define BYTESWAP_FILTERS(f) do {} while (0)
#define BYTESWAP_TENSOR(t) do {} while (0)
#endif

#ifdef __GNUC__
#ifdef __MINGW32__
#define PARAKEET_ATTRIBUTE_FORMAT(...) __attribute__((format(gnu_printf, __VA_ARGS__)))
#else
#define PARAKEET_ATTRIBUTE_FORMAT(...) __attribute__((format(printf, __VA_ARGS__)))
#endif
#else
#define PARAKEET_ATTRIBUTE_FORMAT(...)
#endif

//
// logging
//

PARAKEET_ATTRIBUTE_FORMAT(2, 3)
static void parakeet_log_internal        (ggml_log_level level, const char * format, ...);
static void parakeet_log_callback_default(ggml_log_level level, const char * text, void * user_data);

#define PARAKEET_LOG_ERROR(...) parakeet_log_internal(GGML_LOG_LEVEL_ERROR, __VA_ARGS__)
#define PARAKEET_LOG_WARN(...)  parakeet_log_internal(GGML_LOG_LEVEL_WARN , __VA_ARGS__)
#define PARAKEET_LOG_INFO(...)  parakeet_log_internal(GGML_LOG_LEVEL_INFO , __VA_ARGS__)

// define this to enable verbose trace logging - useful for debugging purposes
//#define PARAKEET_DEBUG

#if defined(PARAKEET_DEBUG)
#define PARAKEET_LOG_DEBUG(...) parakeet_log_internal(GGML_LOG_LEVEL_DEBUG, __VA_ARGS__)
#else
#define PARAKEET_LOG_DEBUG(...)
#endif

#define PARAKEET_ASSERT(x) \
    do { \
        if (!(x)) { \
            PARAKEET_LOG_ERROR("PARAKEET_ASSERT: %s:%d: %s\n", __FILE__, __LINE__, #x); \
            abort(); \
        } \
    } while (0)

#define PARAKEET_MAX_NODES 8192

// Threshold for when local attention should be used.
// 8192 frames x 80ms = 655 s (about 10.9 mins)
static constexpr int PARAKEET_LOCAL_ATTN_THRESHOLD = 8192;
// Window of context in each director of the current token.
// 128 frames * 80ms = 10.24 s
static constexpr int PARAKEET_LOCAL_ATTN_WINDOW    = 128;

static std::string format(const char * fmt, ...) {
    va_list ap;
    va_list ap2;
    va_start(ap, fmt);
    va_copy(ap2, ap);
    int size = vsnprintf(NULL, 0, fmt, ap);
    GGML_ASSERT(size >= 0 && size < INT_MAX); // NOLINT
    std::vector<char> buf(size + 1);
    int size2 = vsnprintf(buf.data(), size + 1, fmt, ap2);
    GGML_ASSERT(size2 == size);
    va_end(ap2);
    va_end(ap);
    return std::string(buf.data(), size);
}

//
// ggml helpers
//

static bool ggml_graph_compute_helper(
          struct ggml_cgraph * graph,
                         int   n_threads,
         ggml_abort_callback   abort_callback,
                        void * abort_callback_data) {
    ggml_backend_ptr backend { ggml_backend_init_by_type(GGML_BACKEND_DEVICE_TYPE_CPU, nullptr) };

    auto * reg = ggml_backend_dev_backend_reg(ggml_backend_get_device(backend.get()));

    auto * set_abort_callback_fn = (ggml_backend_set_abort_callback_t) ggml_backend_reg_get_proc_address(reg, "ggml_backend_set_abort_callback");
    if (set_abort_callback_fn) {
        set_abort_callback_fn(backend.get(), abort_callback, abort_callback_data);
    }

    auto ggml_backend_set_n_threads_fn = (ggml_backend_set_n_threads_t) ggml_backend_reg_get_proc_address(reg, "ggml_backend_set_n_threads");
    if (ggml_backend_set_n_threads_fn) {
        ggml_backend_set_n_threads_fn(backend.get(), n_threads);
    }

    return ggml_backend_graph_compute(backend.get(), graph) == GGML_STATUS_SUCCESS;
}

static bool ggml_graph_compute_helper(
      ggml_backend_sched_t   sched,
        struct ggml_cgraph * graph,
                       int   n_threads,
                      bool   sched_reset = true) {
    for (int i = 0; i < ggml_backend_sched_get_n_backends(sched); ++i) {
        ggml_backend_t backend = ggml_backend_sched_get_backend(sched, i);
        ggml_backend_dev_t dev = ggml_backend_get_device(backend);
        ggml_backend_reg_t reg = dev ? ggml_backend_dev_backend_reg(dev) : nullptr;

        auto * fn_set_n_threads = (ggml_backend_set_n_threads_t) ggml_backend_reg_get_proc_address(reg, "ggml_backend_set_n_threads");
        if (fn_set_n_threads) {
            fn_set_n_threads(backend, n_threads);
        }
    }

    const bool t = (ggml_backend_sched_graph_compute(sched, graph) == GGML_STATUS_SUCCESS);

    if (!t || sched_reset) {
        ggml_backend_sched_reset(sched);
    }

    return t;
}

// TODO: move these functions to ggml-base with support for ggml-backend?


struct parakeet_mel {
    int n_len     = 0;
    int n_len_org = 0;
    int n_mel     = 0;

    std::vector<float> data;
};

struct parakeet_filters {
    int32_t n_mel = 0;
    int32_t n_fb  = 0;  // number of frequency bins

    std::vector<float> data;
};

struct parakeet_vocab {
    using id    = int32_t;
    using token = std::string;

    int n_vocab = 8192;
    size_t max_token_length = 0;

    std::map<token, id> token_to_id;
    std::map<id, token> id_to_token;

    id token_unk;
    id token_bos;
    id token_blank;
    id token_eos;
};

struct parakeet_segment {
    int64_t t0;
    int64_t t1;

    std::string text;

    std::vector<parakeet_token_data> tokens;
};

struct parakeet_batch {
    int32_t n_tokens;

    parakeet_token  *  token;
    int32_t         *  i_time;   // index of the audio frame
    parakeet_pos    *  pos;
    int32_t         *  n_seq_id; // always 1, here for consistency with llama.cpp
    parakeet_seq_id ** seq_id;   // null terminated
    int8_t          *  logits;
};

// ggml_backend_sched wrapper for parakeet usage
struct parakeet_sched {
    ggml_backend_sched_t sched = nullptr;

    std::vector<uint8_t> meta;
};

// TODO: Find out is there a multiple version types. It is not yet clear to me
// at this point.
enum parakeet_arch {
    PARAKEET_ARCH_UNKNOWN = 0,
    PARAKEET_ARCH_TDT     = 1,  // NVIDIA Parakeet TDT (RNN-T)
};

struct parakeet_hparams {
    int32_t n_vocab                = 8192;
    int32_t n_audio_ctx            = 0;  // 0 = unlimited, will be set based on input
    int32_t n_audio_state          = 1024;
    int32_t n_audio_head           = 8;
    int32_t n_audio_layer          = 24;
    int32_t n_mels                 = 128;
    int32_t ftype                  = 1;
    int32_t n_fft                  = 512;  // FFT size for mel spectrogram
    float   eps                    = 1e-5f;
    int32_t subsampling_factor     = 8;
    int32_t n_subsampling_channels = 256;
    int32_t n_conv_kernel          = 9;
    int32_t n_pred_dim             = 640;
    int32_t n_pred_layers          = 2;
    int32_t n_tdt_durations        = 5;
    int32_t n_max_tokens           = 10;

    parakeet_arch arch     = PARAKEET_ARCH_TDT;
};

struct parakeet_layer_encoder {
    struct ggml_tensor * norm_ff1_w = nullptr;
    struct ggml_tensor * norm_ff1_b = nullptr;

    struct ggml_tensor * ff1_linear1_w = nullptr;
    struct ggml_tensor * ff1_linear2_w = nullptr;

    struct ggml_tensor * norm_conv_w = nullptr;
    struct ggml_tensor * norm_conv_b = nullptr;

    struct ggml_tensor * conv_pw1_w          = nullptr;  // pointwise_conv1
    struct ggml_tensor * conv_dw_w           = nullptr;  // depthwise_conv
    struct ggml_tensor * conv_bn_w           = nullptr;  // batch_norm weight
    struct ggml_tensor * conv_bn_b           = nullptr;  // batch_norm bias
    struct ggml_tensor * conv_bn_mean        = nullptr;  // batch_norm running_mean
    struct ggml_tensor * conv_bn_var         = nullptr;  // batch_norm running_var
    struct ggml_tensor * conv_bn_num_batches = nullptr;  // batch_norm num_batches_tracked
    struct ggml_tensor * conv_pw2_w          = nullptr;  // pointwise_conv2

    struct ggml_tensor * norm_attn_w = nullptr;
    struct ggml_tensor * norm_attn_b = nullptr;

    struct ggml_tensor * attn_pos_bias_u = nullptr;
    struct ggml_tensor * attn_pos_bias_v = nullptr;
    struct ggml_tensor * attn_q_w        = nullptr;
    struct ggml_tensor * attn_k_w        = nullptr;
    struct ggml_tensor * attn_v_w        = nullptr;
    struct ggml_tensor * attn_out_w      = nullptr;
    struct ggml_tensor * attn_pos_w      = nullptr;

    struct ggml_tensor * norm_ff2_w      = nullptr;
    struct ggml_tensor * norm_ff2_b      = nullptr;

    struct ggml_tensor * ff2_linear1_w = nullptr;
    struct ggml_tensor * ff2_linear2_w = nullptr;

    struct ggml_tensor * norm_out_w = nullptr;
    struct ggml_tensor * norm_out_b = nullptr;
};

struct parakeet_lsmt_layer {
    struct ggml_tensor * ih_w = nullptr;  // input-to-hidden weight
    struct ggml_tensor * hh_w = nullptr;  // hidden-to-hidden weight
    struct ggml_tensor * b_h = nullptr;   // bias (ih folded into hh at conversion time)
};

struct parakeet_prediction_network {
    struct ggml_tensor * embed_w = nullptr;

    std::vector<parakeet_lsmt_layer> lstm_layer;
};

struct parakeet_joint_network {
    struct ggml_tensor * pred_w = nullptr;
    struct ggml_tensor * pred_b = nullptr;
    struct ggml_tensor * enc_w  = nullptr;
    struct ggml_tensor * enc_b  = nullptr;
    struct ggml_tensor * net_w  = nullptr;
    struct ggml_tensor * net_b  = nullptr;
};

struct parakeet_model {
    parakeet_filters filters;
    parakeet_hparams hparams;

    struct ggml_tensor * enc_pre_out_w    = nullptr;
    struct ggml_tensor * enc_pre_out_b    = nullptr;
    struct ggml_tensor * enc_pre_conv_0_w = nullptr;
    struct ggml_tensor * enc_pre_conv_0_b = nullptr;
    struct ggml_tensor * enc_pre_conv_2_w = nullptr;
    struct ggml_tensor * enc_pre_conv_2_b = nullptr;
    struct ggml_tensor * enc_pre_conv_3_w = nullptr;
    struct ggml_tensor * enc_pre_conv_3_b = nullptr;
    struct ggml_tensor * enc_pre_conv_5_w = nullptr;
    struct ggml_tensor * enc_pre_conv_5_b = nullptr;
    struct ggml_tensor * enc_pre_conv_6_w = nullptr;
    struct ggml_tensor * enc_pre_conv_6_b = nullptr;

    std::vector<parakeet_layer_encoder> layers;

    parakeet_prediction_network prediction;

    parakeet_joint_network joint;

    std::vector<uint32_t> tdt_durations;

    std::vector<ggml_context *> ctxs;

    std::vector<ggml_backend_buffer_t> buffers;

    int n_loaded = 0;
    std::map<std::string, struct ggml_tensor *> tensors;
};

struct parakeet_lstm_state_layer {
    struct ggml_tensor * h_state = nullptr;
    struct ggml_tensor * c_state = nullptr;
};

struct parakeet_lstm_state {
    std::vector<parakeet_lstm_state_layer> layer;

    std::vector<uint8_t> ctx_buf;

    ggml_backend_buffer_t buffer = nullptr;
};

struct parakeet_state {
    int64_t t_sample_us = 0;
    int64_t t_encode_us = 0;
    int64_t t_decode_us = 0;
    int64_t t_predict_us = 0;
    int64_t t_predict_build_us   = 0; // time spent building the prediction graph
    int64_t t_predict_alloc_us   = 0; // time spent in ggml_backend_sched_alloc_graph
    int64_t t_predict_compute_us = 0; // time spent in ggml_graph_compute_helper
    int64_t t_mel_us = 0;

    int32_t n_sample = 0; // number of tokens sampled
    int32_t n_encode = 0; // number of encoder calls
    int32_t n_decode = 0; // number of decoder calls with n_tokens == 1  (text-generation)
    int32_t n_predict = 0; // number of prediction network calls
    int32_t n_fail_p = 0; // number of logprob threshold failures
    int32_t n_fail_h = 0; // number of entropy threshold failures

    parakeet_mel mel;

    parakeet_batch batch;

    int n_frames = 0;

    std::vector<ggml_backend_t> backends;

    parakeet_sched sched_encode;
    parakeet_sched sched_decode;

    // outputs from encoder stages
    struct ggml_tensor * enc_out     = nullptr;
    struct ggml_tensor * pred_out    = nullptr;

    std::vector<uint8_t> enc_out_buf;
    ggml_backend_buffer_t enc_out_buffer = nullptr;

    std::vector<uint8_t> pred_out_buf;
    ggml_backend_buffer_t pred_out_buffer = nullptr;

    struct ggml_tensor * attn_mask = nullptr;

    std::vector<float> inp_mel;
    std::vector<float> inp_mask;

    std::vector<float> logits;

    std::vector<parakeet_segment> result_all;

    std::vector<parakeet_token>      decoded_tokens;
    std::vector<parakeet_token_data> decoded_token_data;

    std::string path_model;

    int32_t n_audio_ctx = 0;
    int32_t sched_encode_n_audio_ctx = 0;

    parakeet_lstm_state lstm_state;
};

// FFT cache for mel spectrogram computation
struct parakeet_mel_cache {
    int n_fft = 0;

    // In FFT, we frequently use sine and cosine operations with the same values.
    // We can use precalculated values to speed up the process.
    std::vector<float> sin_vals;
    std::vector<float> cos_vals;

    // Hann window (Use cosf to eliminate difference)
    // ref: https://pytorch.org/docs/stable/generated/torch.hann_window.html
    // ref: https://github.com/openai/whisper/blob/main/whisper/audio.py#L147
    std::vector<float> hann_window;

    // Window function from model (Parakeet uses actual window from training)
    std::vector<float> window;

    void init(int fft_size) {
        n_fft = fft_size;
        sin_vals.resize(n_fft);
        cos_vals.resize(n_fft);
        hann_window.resize(n_fft);

        fill_sin_cos_table();
        fill_hann_window(n_fft, true, hann_window.data());
    }

    void fill_sin_cos_table() {
        for (int i = 0; i < n_fft; i++) {
            double theta = (2 * M_PI * i) / n_fft;
            sin_vals[i] = sinf(theta);
            cos_vals[i] = cosf(theta);
        }
    }

    void fill_hann_window(int length, bool periodic, float * output) {
        int offset = -1;
        if (periodic) {
            offset = 0;
        }
        for (int i = 0; i < length; i++) {
            output[i] = 0.5 * (1.0 - cosf((2.0 * M_PI * i) / (length + offset)));
        }
    }
};

struct parakeet_context {
    int64_t t_load_us  = 0;
    int64_t t_start_us = 0;

    ggml_type wtype = ggml_type::GGML_TYPE_F16;
    ggml_type itype = ggml_type::GGML_TYPE_F16;

    parakeet_context_params params;

    parakeet_model model;
    parakeet_vocab vocab;

    parakeet_state * state = nullptr;

    parakeet_mel_cache mel_cache;

    std::string path_model;
};

struct parakeet_global {
    // We save the log callback globally
    ggml_log_callback log_callback = parakeet_log_callback_default;
    void * log_callback_user_data = nullptr;
};

static parakeet_global g_state;

static const std::string PARAKEET_SPM_SPACE = "\xE2\x96\x81";

static inline int utf8_codepoint_len(unsigned char c) {
    if ((c & 0x80) == 0x00) return 1;
    if ((c & 0xE0) == 0xC0) return 2;
    if ((c & 0xF0) == 0xE0) return 3;
    if ((c & 0xF8) == 0xF0) return 4;
    return 1;
}

static bool is_sentencepiece_control(const std::string & piece) {
    return piece == "<unk>" || piece == "<s>" || piece == "</s>" || piece == "[BLANK]";
}

static std::string sentencepiece_normalize(const std::string & text) {
    std::string normalized;
    normalized.reserve(text.size() + PARAKEET_SPM_SPACE.size());
    normalized += PARAKEET_SPM_SPACE; // SentencePiece dummy prefix

    for (unsigned char c : text) {
        if (std::isspace(c)) {
            normalized += PARAKEET_SPM_SPACE;
        } else {
            normalized += static_cast<char>(c);
        }
    }

    return normalized;
}

static std::string sentencepiece_piece_to_text(const std::string & piece, bool is_first_piece) {
    if (is_sentencepiece_control(piece)) {
        return "";
    }

    std::string text;
    text.reserve(piece.size());

    size_t pos = 0;
    while (pos < piece.size()) {
        if (piece.compare(pos, PARAKEET_SPM_SPACE.size(), PARAKEET_SPM_SPACE) == 0) {
            if (!is_first_piece || !text.empty()) {
                text += ' ';
            }
            pos += PARAKEET_SPM_SPACE.size();
            continue;
        }

        text += piece[pos];
        ++pos;
    }

    return text;
}


static struct parakeet_batch parakeet_batch_init(int32_t n_tokens) {
    parakeet_batch batch = { 0, nullptr, nullptr, nullptr, nullptr, nullptr, nullptr, };

    batch.token    = (parakeet_token *  ) malloc(sizeof(parakeet_token)    * (n_tokens));
    batch.i_time   = (int32_t *)          malloc(sizeof(int32_t)           * (n_tokens));
    batch.pos      = (parakeet_pos *)     malloc(sizeof(parakeet_pos)      * (n_tokens));
    batch.n_seq_id = (int32_t *)          malloc(sizeof(int32_t)           * (n_tokens));
    batch.seq_id   = (parakeet_seq_id **) malloc(sizeof(parakeet_seq_id *) * (n_tokens + 1));
    for (int i = 0; i < n_tokens; ++i) {
        batch.seq_id[i] = (parakeet_seq_id *) malloc(sizeof(parakeet_seq_id));
    }
    batch.seq_id[n_tokens] = nullptr;
    batch.logits   = (int8_t *)          malloc(sizeof(int8_t)           * n_tokens);

    return batch;
}

static void parakeet_batch_free(struct parakeet_batch batch) {
    if (batch.token)    free(batch.token);
    if (batch.i_time)   free(batch.i_time);
    if (batch.pos)      free(batch.pos);
    if (batch.n_seq_id) free(batch.n_seq_id);
    if (batch.seq_id) {
        for (int i = 0; batch.seq_id[i]; ++i) {
            free(batch.seq_id[i]);
        }
        free(batch.seq_id);
    }
    if (batch.logits)   free(batch.logits);
}

static void parakeet_batch_prep_legacy(parakeet_batch & batch, const parakeet_token * tokens, int n_tokens, int n_past, int seq_id) {
    batch.n_tokens = n_tokens;
    for (int i = 0; i < n_tokens; ++i) {
        if (tokens) {
            batch.token[i] = tokens[i];
        }
        batch.pos     [i]    = n_past + i;
        batch.n_seq_id[i]    = 1;
        batch.seq_id  [i][0] = seq_id;
        batch.logits  [i]    = 0;
    }
    batch.logits[n_tokens - 1] = 1;
}


static size_t parakeet_sched_size(struct parakeet_sched & allocr) {
    size_t size = allocr.meta.size();
    for (int i = 0; i < ggml_backend_sched_get_n_backends(allocr.sched); ++i) {
        ggml_backend_t backend = ggml_backend_sched_get_backend(allocr.sched, i);
        size += ggml_backend_sched_get_buffer_size(allocr.sched, backend);
    }
    return size;
}

static bool parakeet_sched_graph_init(struct parakeet_sched & allocr, std::vector<ggml_backend_t> backends, std::function<struct ggml_cgraph *()> && get_graph) {
    auto & sched = allocr.sched;
    auto & meta  = allocr.meta;

    sched = ggml_backend_sched_new(backends.data(), nullptr, backends.size(), PARAKEET_MAX_NODES, false, true);

    if (!sched) {
        PARAKEET_LOG_ERROR("%s: failed to create scheduler\n", __func__);
        return false;
    }

    meta.resize(ggml_tensor_overhead()*PARAKEET_MAX_NODES + ggml_graph_overhead());

    if (!ggml_backend_sched_alloc_graph(sched, get_graph())) {
        PARAKEET_LOG_ERROR("%s: failed to allocate the compute buffer\n", __func__);
        ggml_backend_sched_free(sched);
        sched = nullptr;
        return false;
    }

    ggml_backend_sched_reset(sched);

    return true;
}

static void parakeet_sched_free(struct parakeet_sched & sched) {
    if (sched.sched) {
        ggml_backend_sched_free(sched.sched);
        sched.sched = nullptr;
    }

    sched.meta.clear();
}


template<typename T>
static void read_safe(parakeet_model_loader * loader, T & dest) {
    loader->read(loader->context, &dest, sizeof(T));
    BYTESWAP_VALUE(dest);
}

static bool parakeet_lstm_state_init(
               struct parakeet_state & pstate,
                      ggml_backend_t   backend,
                                 int   n_layer,
                                 int   n_pred_dim) {
    parakeet_lstm_state & lstm_state = pstate.lstm_state;

    lstm_state.ctx_buf.resize(ggml_tensor_overhead() * n_layer * 2);
    lstm_state.layer.resize(n_layer);

    struct ggml_init_params params = {
        /*.mem_size   =*/ lstm_state.ctx_buf.size(),
        /*.mem_buffer =*/ lstm_state.ctx_buf.data(),
        /*.no_alloc   =*/ true,
    };

    struct ggml_context * ctx = ggml_init(params);

    if (!ctx) {
        PARAKEET_LOG_ERROR("%s: failed to allocate memory for the lstm states context\n", __func__);
        return false;
    }


    for (int il = 0; il < n_layer; ++il) {
        lstm_state.layer[il].h_state = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, n_pred_dim);
        lstm_state.layer[il].c_state = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, n_pred_dim);
    }

    lstm_state.buffer = ggml_backend_alloc_ctx_tensors(ctx, backend);
    if (!lstm_state.buffer) {
        PARAKEET_LOG_ERROR("%s: failed to allocate memory for the lstm states\n", __func__);
        return false;
    }

    ggml_backend_buffer_clear(lstm_state.buffer, 0);

    ggml_free(ctx);

    return true;
}

static bool parakeet_pred_state_init(
               struct parakeet_state & pstate,
                      ggml_backend_t   backend,
                                 int   n_pred_dim) {
    pstate.pred_out_buf.resize(ggml_tensor_overhead());

    struct ggml_init_params params = {
        /*.mem_size   =*/ pstate.pred_out_buf.size(),
        /*.mem_buffer =*/ pstate.pred_out_buf.data(),
        /*.no_alloc   =*/ true,
    };

    struct ggml_context * ctx = ggml_init(params);
    if (!ctx) {
        PARAKEET_LOG_ERROR("%s: failed to allocate memory for pred tensor context\n", __func__);
        return false;
    }

    pstate.pred_out = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, n_pred_dim);
    pstate.pred_out_buffer = ggml_backend_alloc_ctx_tensors(ctx, backend);
    if (!pstate.pred_out_buffer) {
        PARAKEET_LOG_ERROR("%s: failed to allocate memory for pred tensor\n", __func__);
        ggml_free(ctx);
        return false;
    }

    ggml_free(ctx);

    return true;
}

static bool parakeet_enc_state_init(
               struct parakeet_state & pstate,
                      ggml_backend_t   backend,
                                 int   n_audio_state,
                                 int   n_frames_max) {
    pstate.enc_out_buf.resize(ggml_tensor_overhead());

    struct ggml_init_params params = {
        /*.mem_size   =*/ pstate.enc_out_buf.size(),
        /*.mem_buffer =*/ pstate.enc_out_buf.data(),
        /*.no_alloc   =*/ true,
    };

    struct ggml_context * ctx = ggml_init(params);
    if (!ctx) {
        PARAKEET_LOG_ERROR("%s: failed to allocate memory for enc_out tensor context\n", __func__);
        return false;
    }

    pstate.enc_out = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, n_audio_state, n_frames_max);
    pstate.enc_out_buffer = ggml_backend_alloc_ctx_tensors(ctx, backend);
    if (!pstate.enc_out_buffer) {
        PARAKEET_LOG_ERROR("%s: failed to allocate memory for enc_out tensor\n", __func__);
        ggml_free(ctx);
        return false;
    }

    ggml_free(ctx);

    return true;
}

static ggml_backend_t parakeet_backend_init_gpu(const parakeet_context_params & params) {
    ggml_log_set(g_state.log_callback, g_state.log_callback_user_data);

    ggml_backend_dev_t dev = nullptr;

    int cnt = 0;
    if (params.use_gpu) {
        for (size_t i = 0; i < ggml_backend_dev_count(); ++i) {
            ggml_backend_dev_t dev_cur = ggml_backend_dev_get(i);
            enum ggml_backend_dev_type dev_type = ggml_backend_dev_type(dev_cur);
            const char * dev_name = ggml_backend_dev_name(dev_cur);
            PARAKEET_LOG_INFO("%s: device %zu: %s (type: %d)\n", __func__, i, dev_name, dev_type);
            if (dev_type == GGML_BACKEND_DEVICE_TYPE_GPU || dev_type == GGML_BACKEND_DEVICE_TYPE_IGPU) {
                PARAKEET_LOG_INFO("%s: found GPU device %zu: %s (type: %d, cnt: %d)\n", __func__, i, dev_name, dev_type, cnt);
                if (cnt == params.gpu_device) {
                    dev = dev_cur;
                }

                if (++cnt > params.gpu_device) {
                    break;
                }
            }
        }
    }

    if (dev == nullptr) {
        PARAKEET_LOG_INFO("%s: no GPU found\n", __func__);
        return nullptr;
    }

    PARAKEET_LOG_INFO("%s: using %s backend\n", __func__, ggml_backend_dev_name(dev));
    ggml_backend_t result = ggml_backend_dev_init(dev, nullptr);
    if (!result) {
        PARAKEET_LOG_ERROR("%s: failed to initialize %s backend\n", __func__, ggml_backend_dev_name(dev));
    }

    return result;
}

static std::vector<ggml_backend_t> parakeet_backend_init(const parakeet_context_params & params) {
    std::vector<ggml_backend_t> result;

    ggml_backend_t backend_gpu = parakeet_backend_init_gpu(params);

    if (backend_gpu) {
        result.push_back(backend_gpu);
    }

    // ACCEL backends
    for (size_t i = 0; i < ggml_backend_dev_count(); ++i) {
        ggml_backend_dev_t dev = ggml_backend_dev_get(i);
        if (ggml_backend_dev_type(dev) == GGML_BACKEND_DEVICE_TYPE_ACCEL) {
            PARAKEET_LOG_INFO("%s: using %s backend\n", __func__, ggml_backend_dev_name(dev));
            ggml_backend_t backend = ggml_backend_dev_init(dev, nullptr);
            if (!backend) {
                PARAKEET_LOG_ERROR("%s: failed to initialize %s backend\n", __func__, ggml_backend_dev_name(dev));
                continue;
            }
            result.push_back(backend);
        }
    }

    ggml_backend_t backend_cpu = ggml_backend_init_by_type(GGML_BACKEND_DEVICE_TYPE_CPU, nullptr);
    if (backend_cpu == nullptr) {
        throw std::runtime_error("failed to initialize CPU backend");
    }
    result.push_back(backend_cpu);

    return result;
}

using buft_list_t = std::vector<std::pair<ggml_backend_dev_t, ggml_backend_buffer_type_t>>;

static buft_list_t make_buft_list(parakeet_context_params & params) {
    // Prio order: GPU -> CPU Extra -> CPU
    buft_list_t buft_list;

    // GPU
    if (params.use_gpu) {
        int cnt = 0;
        for (size_t i = 0; i < ggml_backend_dev_count(); ++i) {
            ggml_backend_dev_t dev = ggml_backend_dev_get(i);
            if (ggml_backend_dev_type(dev) == GGML_BACKEND_DEVICE_TYPE_GPU || ggml_backend_dev_type(dev) == GGML_BACKEND_DEVICE_TYPE_IGPU) {
                if (cnt == params.gpu_device) {
                    auto * buft = ggml_backend_dev_buffer_type(dev);
                    if (buft) {
                        buft_list.emplace_back(dev, buft);
                    }
                }

                if (++cnt > params.gpu_device) {
                    break;
                }
            }
        }
    }

    // CPU Extra
    auto * cpu_dev = ggml_backend_dev_by_type(GGML_BACKEND_DEVICE_TYPE_CPU);
    auto * cpu_reg = ggml_backend_dev_backend_reg(cpu_dev);
    auto get_extra_bufts_fn = (ggml_backend_dev_get_extra_bufts_t)
        ggml_backend_reg_get_proc_address(cpu_reg, "ggml_backend_dev_get_extra_bufts");
    if (get_extra_bufts_fn) {
        ggml_backend_buffer_type_t * extra_bufts = get_extra_bufts_fn(cpu_dev);
        while (extra_bufts && *extra_bufts) {
            buft_list.emplace_back(cpu_dev, *extra_bufts);
            ++extra_bufts;
        }
    }

    // CPU
    buft_list.emplace_back(cpu_dev, ggml_backend_cpu_buffer_type());

    return buft_list;
}

static bool weight_buft_supported(const parakeet_hparams & hparams, ggml_tensor * w, ggml_op op, ggml_backend_buffer_type_t buft, ggml_backend_dev_t dev) {
    bool op_supported = true;

    if (ggml_backend_dev_type(dev) == GGML_BACKEND_DEVICE_TYPE_GPU ||
        ggml_backend_dev_type(dev) == GGML_BACKEND_DEVICE_TYPE_IGPU ||
        (ggml_backend_dev_type(dev) == GGML_BACKEND_DEVICE_TYPE_CPU && buft == ggml_backend_cpu_buffer_type())) {
        // GPU and default CPU backend support all operators
        op_supported = true;
    } else {
        switch (op) {
            // The current extra_buffer_type implementations only support GGML_OP_MUL_MAT and GGML_OP_GET_ROWS
            case GGML_OP_GET_ROWS:
            case GGML_OP_MUL_MAT: {
                ggml_init_params params = {
                    /*.mem_size   =*/ 2 * ggml_tensor_overhead(),
                    /*.mem_buffer =*/ nullptr,
                    /*.no_alloc   =*/ true,
                };

                ggml_context_ptr ctx_ptr { ggml_init(params) };
                if (!ctx_ptr) {
                    throw std::runtime_error("failed to create ggml context");
                }
                ggml_context * ctx = ctx_ptr.get();

                ggml_tensor * op_tensor = nullptr;

                if (op == GGML_OP_MUL_MAT) {
                    int64_t n_ctx = hparams.n_audio_ctx;
                    ggml_tensor * b = ggml_new_tensor_4d(ctx, GGML_TYPE_F32, w->ne[0], n_ctx, w->ne[2], w->ne[3]);
                    op_tensor = ggml_mul_mat(ctx, w, b);
                } else if (op == GGML_OP_GET_ROWS) {
                    int64_t num_indices = 8;
                    ggml_tensor * indices = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, num_indices);
                    op_tensor = ggml_get_rows(ctx, w, indices);
                }

                // create a temporary dummy buffer for the weight so that supports_op can check the buffer type
                GGML_ASSERT(w->buffer == nullptr);
                w->buffer = ggml_backend_buft_alloc_buffer(buft, 0);
                op_supported = ggml_backend_dev_supports_op(dev, op_tensor);
                ggml_backend_buffer_free(w->buffer);
                w->buffer = nullptr;
                break;
            }
            default: {
                op_supported = false;
                break;
            }
        };
    }

    return op_supported;
}

static ggml_backend_buffer_type_t select_weight_buft(const parakeet_hparams & hparams, ggml_tensor * w, ggml_op op, buft_list_t buft_list) {
    GGML_ASSERT(!buft_list.empty());
    for (const auto & p : buft_list) {
        ggml_backend_dev_t dev = p.first;
        ggml_backend_buffer_type_t buft = p.second;
        if (weight_buft_supported(hparams, w, op, buft, dev)) {
            return buft;
        }
    }

    return nullptr;
}


// load the model from a ggml file
//

// see the convert-parakeet-to-ggml.py script for details
//
static bool parakeet_model_load(struct parakeet_model_loader * loader, parakeet_context & wctx) {
    PARAKEET_LOG_INFO("%s: loading model\n", __func__);

    const int64_t t_start_us = ggml_time_us();

    wctx.t_start_us = t_start_us;

    auto & model = wctx.model;
    auto & vocab = wctx.vocab;

    // verify magic
    {
        uint32_t magic;
        read_safe(loader, magic);
        if (magic != GGML_FILE_MAGIC) {
            PARAKEET_LOG_ERROR("%s: invalid model data (bad magic)\n", __func__);
            return false;
        }
    }

    //load hparams
    parakeet_hparams hparams;
    {
        read_safe(loader, hparams.n_vocab);
        read_safe(loader, hparams.n_audio_ctx);
        read_safe(loader, hparams.n_audio_state);
        read_safe(loader, hparams.n_audio_head);
        read_safe(loader, hparams.n_audio_layer);
        read_safe(loader, hparams.n_mels);
        read_safe(loader, hparams.ftype);
        read_safe(loader, hparams.n_fft);
        read_safe(loader, hparams.subsampling_factor);
        read_safe(loader, hparams.n_subsampling_channels);
        read_safe(loader, hparams.n_conv_kernel);
        read_safe(loader, hparams.n_pred_dim);
        read_safe(loader, hparams.n_pred_layers);
        read_safe(loader, hparams.n_tdt_durations);
        read_safe(loader, hparams.n_max_tokens);

        hparams.arch = PARAKEET_ARCH_TDT;
        wctx.model.hparams = hparams;

        const int32_t qntvr = hparams.ftype / GGML_QNT_VERSION_FACTOR;

        hparams.ftype %= GGML_QNT_VERSION_FACTOR;

        // for the big tensors, we have the option to store the data in 16-bit floats or quantized
        // in order to save memory and also to speed up the computation
        wctx.wtype = ggml_ftype_to_ggml_type((ggml_ftype) hparams.ftype);
        if (wctx.wtype == GGML_TYPE_COUNT) {
            PARAKEET_LOG_ERROR("%s: invalid model (bad ftype value %d)\n", __func__, hparams.ftype);
            return false;
        }

        const char* arch_name = hparams.arch == PARAKEET_ARCH_TDT ? "Parakeet TDT" : "unknown";
        PARAKEET_LOG_INFO("%s: arch                   = %s\n", __func__, arch_name);
        PARAKEET_LOG_INFO("%s: n_vocab                = %d\n", __func__, hparams.n_vocab);
        PARAKEET_LOG_INFO("%s: n_audio_ctx            = %d\n", __func__, hparams.n_audio_ctx);
        PARAKEET_LOG_INFO("%s: n_audio_state          = %d\n", __func__, hparams.n_audio_state);
        PARAKEET_LOG_INFO("%s: n_audio_head           = %d\n", __func__, hparams.n_audio_head);
        PARAKEET_LOG_INFO("%s: n_audio_layer          = %d\n", __func__, hparams.n_audio_layer);
        PARAKEET_LOG_INFO("%s: n_mels                 = %d\n", __func__, hparams.n_mels);
        PARAKEET_LOG_INFO("%s: n_fft                  = %d\n", __func__, hparams.n_fft);
        PARAKEET_LOG_INFO("%s: eps                    = %f\n", __func__, hparams.eps);
        PARAKEET_LOG_INFO("%s: ftype                  = %d\n", __func__, hparams.ftype);
        PARAKEET_LOG_INFO("%s: qntvr                  = %d\n", __func__, qntvr);
        PARAKEET_LOG_INFO("%s: subsampling_factor     = %d\n", __func__, hparams.subsampling_factor);
        PARAKEET_LOG_INFO("%s: n_subsampling_channels = %d\n", __func__, hparams.n_subsampling_channels);
        PARAKEET_LOG_INFO("%s: n_conv_kernel          = %d\n", __func__, hparams.n_conv_kernel);
        PARAKEET_LOG_INFO("%s: n_pred_dim             = %d\n", __func__, hparams.n_pred_dim);
        PARAKEET_LOG_INFO("%s: n_pred_layers          = %d\n", __func__, hparams.n_pred_layers);
        PARAKEET_LOG_INFO("%s: n_tdt_durations        = %d\n", __func__, hparams.n_tdt_durations);
        PARAKEET_LOG_INFO("%s: n_max_tokens           = %d\n", __func__, hparams.n_max_tokens);
    }

    // load mel filters
    {
        auto & filters = wctx.model.filters;

        read_safe(loader, filters.n_mel);
        read_safe(loader, filters.n_fb);

        filters.data.resize(filters.n_mel * filters.n_fb);
        loader->read(loader->context, filters.data.data(), filters.data.size() * sizeof(float));
        BYTESWAP_FILTERS(filters);
    }

    // load window function
    {
        int32_t n_window = 0;
        read_safe(loader, n_window);

        wctx.mel_cache.window.resize(n_window);
        loader->read(loader->context, wctx.mel_cache.window.data(), n_window * sizeof(float));

#ifdef GGML_BIG_ENDIAN
        for (auto & datum : wctx.mel_cache.window) {
            datum = byteswap(datum);
        }
#endif

        PARAKEET_LOG_INFO("%s: loaded window function with %d samples\n", __func__, n_window);
    }

    // load TDT (Token and Duration Transducer) values
    {
        auto & tdt_durations = wctx.model.tdt_durations;
        tdt_durations.resize(hparams.n_tdt_durations);
        loader->read(loader->context, tdt_durations.data(), hparams.n_tdt_durations * sizeof(uint32_t));

        PARAKEET_LOG_INFO("%s: loaded tdt_durations: [", __func__);
        for (const auto value : tdt_durations) {
            PARAKEET_LOG_INFO("%u ", value);
        }
        PARAKEET_LOG_INFO("]\n");
    }

    // load vocab
    {
        int32_t n_vocab = 0;
        read_safe(loader, n_vocab);

        std::string word;
        std::vector<char> tmp;

        tmp.reserve(128);

        for (int i = 0; i < n_vocab; i++) {
            uint32_t len;
            read_safe(loader, len);

            if (len > 0) {
                tmp.resize(len);
                loader->read(loader->context, &tmp[0], tmp.size()); // read to buffer
                word.assign(&tmp[0], tmp.size());
            } else {
                PARAKEET_LOG_WARN("%s: warning: empty-string token in vocab, i = %d\n", __func__, i);
                word = "";
            }

            vocab.token_to_id[word] = i;
            vocab.id_to_token[i] = word;
            vocab.max_token_length = std::max(vocab.max_token_length, word.size());
        }
        // Blank token for transducer is at index n_vocab (8192), outside the vocabulary
        int blank_id = n_vocab;
        vocab.token_blank = blank_id;
        vocab.id_to_token[blank_id] = "[BLANK]";
        vocab.token_to_id["[BLANK]"] = blank_id;

        // Set special token IDs by looking them up in the loaded vocabulary
        // These are from the SentencePiece vocab file loaded above
        if (vocab.token_to_id.find("<unk>") != vocab.token_to_id.end()) {
            vocab.token_unk = vocab.token_to_id.at("<unk>");
        } else {
            vocab.token_unk = 0;  // Fallback
        }

        if (vocab.token_to_id.find("<s>") != vocab.token_to_id.end()) {
            vocab.token_bos = vocab.token_to_id.at("<s>");
        } else if (vocab.token_to_id.find("<|startoftranscript|>") != vocab.token_to_id.end()) {
            vocab.token_bos = vocab.token_to_id.at("<|startoftranscript|>");
        } else {
            vocab.token_bos = 0;  // Fallback
        }

        if (vocab.token_to_id.find("</s>") != vocab.token_to_id.end()) {
            vocab.token_eos = vocab.token_to_id.at("</s>");
        } else if (vocab.token_to_id.find("<|endoftext|>") != vocab.token_to_id.end()) {
            vocab.token_eos = vocab.token_to_id.at("<|endoftext|>");
        } else {
            vocab.token_eos = 0;  // Fallback
        }

        vocab.n_vocab = model.hparams.n_vocab;

        PARAKEET_LOG_INFO("%s: loaded vocab with %d tokens (blank_id=%d, unk=%d, bos=%d, eos=%d)\n",
            __func__, n_vocab, blank_id, vocab.token_unk, vocab.token_bos, vocab.token_eos);
    }

    const ggml_type wtype = wctx.wtype;


    const int n_audio_layer = hparams.n_audio_layer;

    // Calculate tensor count: pre_encode (12) + encoder layers (29 per layer) + prediction (9) + joint (6)
    size_t n_tensors = 12 + (29 * n_audio_layer) + 9 + 6;

    std::map<ggml_backend_buffer_type_t, ggml_context *> ctx_map;
    auto get_ctx = [&](ggml_backend_buffer_type_t buft) -> ggml_context * {
        auto it = ctx_map.find(buft);
        if (it == ctx_map.end()) {
            ggml_init_params params = {
                /*.mem_size   =*/ n_tensors * ggml_tensor_overhead(),
                /*.mem_buffer =*/ nullptr,
                /*.no_alloc   =*/ true,
            };

            ggml_context * ctx = ggml_init(params);
            if (!ctx) {
                throw std::runtime_error("failed to create ggml context");
            }

            ctx_map[buft] = ctx;
            wctx.model.ctxs.emplace_back(ctx);

            return ctx;
        }

        return it->second;
    };

    // Create a list of available bufts, in priority order
    buft_list_t buft_list = make_buft_list(wctx.params);

    auto create_tensor = [&](parakeet_tensor type, ggml_tensor * meta, int layer = -1) -> ggml_tensor * {
        ggml_op op = PARAKEET_TENSOR_INFO.at(type);
        ggml_backend_buffer_type_t buft = select_weight_buft(hparams, meta, op, buft_list);
        if (!buft) {
            throw std::runtime_error(format("failed to find a compatible buffer type for parakeet tensor %s",
                        PARAKEET_TENSOR_NAMES.at(type)));
        }

        ggml_context * ctx = get_ctx(buft);
        ggml_tensor * tensor = ggml_dup_tensor(ctx, meta);

        std::string tensor_name;
        if (layer >= 0) {
            tensor_name = format(PARAKEET_TENSOR_NAMES.at(type), layer);
        } else {
            tensor_name = PARAKEET_TENSOR_NAMES.at(type);
        }

        wctx.model.tensors[tensor_name] = tensor;

        return tensor;
    };

    // prepare tensors for the weights

    ggml_init_params params = {
        /*.mem_size   =*/ n_tensors * ggml_tensor_overhead(),
        /*.mem_buffer =*/ nullptr,
        /*.no_alloc   =*/ true,
    };

    ggml_context * ctx = ggml_init(params);

    const int n_audio_state = hparams.n_audio_state;

    model.layers.resize(n_audio_layer);

    // Encoder pre_encode
    const int n_subsampling_channels = hparams.n_subsampling_channels;
    const int n_pre_enc_features     = (hparams.n_mels / hparams.subsampling_factor) * n_subsampling_channels;
    model.enc_pre_out_w = create_tensor(PARAKEET_TENSOR_ENC_PRE_OUT_WEIGHT, ggml_new_tensor_2d(ctx, wtype, n_pre_enc_features, n_audio_state));
    ggml_set_name(model.enc_pre_out_w, "enc_pre_out_w");
    model.enc_pre_out_b = create_tensor(PARAKEET_TENSOR_ENC_PRE_OUT_BIAS, ggml_new_tensor_1d(ctx, GGML_TYPE_F32, n_audio_state));
    ggml_set_name(model.enc_pre_out_b, "enc_pre_out_b");

    model.enc_pre_conv_0_w = create_tensor(PARAKEET_TENSOR_ENC_PRE_CONV_0_WEIGHT, ggml_new_tensor_4d(ctx, GGML_TYPE_F32, 3, 3, 1, n_subsampling_channels));
    ggml_set_name(model.enc_pre_conv_0_w, "enc_pre_conv_0_w");
    model.enc_pre_conv_0_b = create_tensor(PARAKEET_TENSOR_ENC_PRE_CONV_0_BIAS, ggml_new_tensor_4d(ctx, GGML_TYPE_F32, 1, 1, n_subsampling_channels, 1));
    ggml_set_name(model.enc_pre_conv_0_b, "enc_pre_conv_0_b");

    model.enc_pre_conv_2_w = create_tensor(PARAKEET_TENSOR_ENC_PRE_CONV_2_WEIGHT, ggml_new_tensor_4d(ctx, GGML_TYPE_F32, 3, 3, 1, n_subsampling_channels));
    ggml_set_name(model.enc_pre_conv_2_w, "enc_pre_conv_2_w");
    model.enc_pre_conv_2_b = create_tensor(PARAKEET_TENSOR_ENC_PRE_CONV_2_BIAS, ggml_new_tensor_4d(ctx, GGML_TYPE_F32, 1, 1, n_subsampling_channels, 1));
    ggml_set_name(model.enc_pre_conv_2_b, "enc_pre_conv_2_b");

    model.enc_pre_conv_3_w = create_tensor(PARAKEET_TENSOR_ENC_PRE_CONV_3_WEIGHT, ggml_new_tensor_4d(ctx, GGML_TYPE_F32, 1, 1, n_subsampling_channels, n_subsampling_channels));
    ggml_set_name(model.enc_pre_conv_3_w, "enc_pre_conv_3_w");
    model.enc_pre_conv_3_b = create_tensor(PARAKEET_TENSOR_ENC_PRE_CONV_3_BIAS, ggml_new_tensor_4d(ctx, GGML_TYPE_F32, 1, 1, n_subsampling_channels, 1));
    ggml_set_name(model.enc_pre_conv_3_b, "enc_pre_conv_3_b");

    model.enc_pre_conv_5_w = create_tensor(PARAKEET_TENSOR_ENC_PRE_CONV_5_WEIGHT, ggml_new_tensor_4d(ctx, GGML_TYPE_F32, 3, 3, 1, n_subsampling_channels));
    ggml_set_name(model.enc_pre_conv_5_w, "enc_pre_conv_5_w");
    model.enc_pre_conv_5_b = create_tensor(PARAKEET_TENSOR_ENC_PRE_CONV_5_BIAS, ggml_new_tensor_4d(ctx, GGML_TYPE_F32, 1, 1, n_subsampling_channels, 1));
    ggml_set_name(model.enc_pre_conv_5_b, "enc_pre_conv_5_b");

    model.enc_pre_conv_6_w = create_tensor(PARAKEET_TENSOR_ENC_PRE_CONV_6_WEIGHT, ggml_new_tensor_4d(ctx, GGML_TYPE_F32, 1, 1, n_subsampling_channels, n_subsampling_channels));
    ggml_set_name(model.enc_pre_conv_6_w, "enc_pre_conv_6_w");
    model.enc_pre_conv_6_b = create_tensor(PARAKEET_TENSOR_ENC_PRE_CONV_6_BIAS, ggml_new_tensor_4d(ctx, GGML_TYPE_F32, 1, 1, n_subsampling_channels, 1));
    ggml_set_name(model.enc_pre_conv_6_b, "enc_pre_conv_6_b");

    // Encoder layers
    for (int i = 0; i < n_audio_layer; ++i) {
        auto & layer = model.layers[i];

        // Feed forward 1
        layer.norm_ff1_w    = create_tensor(PARAKEET_TENSOR_ENC_NORM_FF1_WEIGHT, ggml_new_tensor_1d(ctx, GGML_TYPE_F32, n_audio_state), i);
        layer.norm_ff1_b    = create_tensor(PARAKEET_TENSOR_ENC_NORM_FF1_BIAS, ggml_new_tensor_1d(ctx, GGML_TYPE_F32, n_audio_state), i);
        layer.ff1_linear1_w = create_tensor(PARAKEET_TENSOR_ENC_FF1_LINEAR1_WEIGHT, ggml_new_tensor_2d(ctx, wtype, n_audio_state, 4*n_audio_state), i);
        ggml_format_name(layer.ff1_linear1_w, "enc_%d_ff1_linear1_w", i);
        layer.ff1_linear2_w = create_tensor(PARAKEET_TENSOR_ENC_FF1_LINEAR2_WEIGHT, ggml_new_tensor_2d(ctx, wtype, 4*n_audio_state, n_audio_state), i);
        ggml_format_name(layer.ff1_linear2_w, "enc_%d_ff1_linear2_w", i);

        // Convolution module
        layer.norm_conv_w         = create_tensor(PARAKEET_TENSOR_ENC_NORM_CONV_WEIGHT, ggml_new_tensor_1d(ctx, GGML_TYPE_F32, n_audio_state), i);
        ggml_format_name(layer.norm_conv_w, "enc_%d_norm_conv_w", i);
        layer.norm_conv_b         = create_tensor(PARAKEET_TENSOR_ENC_NORM_CONV_BIAS, ggml_new_tensor_1d(ctx, GGML_TYPE_F32, n_audio_state), i);
        ggml_format_name(layer.norm_conv_b, "enc_%d_norm_conv_b", i);
        layer.conv_pw1_w          = create_tensor(PARAKEET_TENSOR_ENC_CONV_PW1_WEIGHT, ggml_new_tensor_2d(ctx, wtype, n_audio_state, 2*n_audio_state), i);
        ggml_format_name(layer.conv_pw1_w, "enc_%d_conv_pw1_w", i);
        layer.conv_dw_w           = create_tensor(PARAKEET_TENSOR_ENC_CONV_DW_WEIGHT, ggml_new_tensor_2d(ctx, GGML_TYPE_F32, hparams.n_conv_kernel, n_audio_state), i);
        ggml_format_name(layer.conv_dw_w, "enc_%d_conv_dw_w", i);
        layer.conv_bn_w           = create_tensor(PARAKEET_TENSOR_ENC_CONV_BN_WEIGHT, ggml_new_tensor_1d(ctx, GGML_TYPE_F32, n_audio_state), i);
        ggml_format_name(layer.conv_bn_w, "enc_%d_conv_bn_w", i);
        layer.conv_bn_b           = create_tensor(PARAKEET_TENSOR_ENC_CONV_BN_BIAS, ggml_new_tensor_1d(ctx, GGML_TYPE_F32, n_audio_state), i);
        ggml_format_name(layer.conv_bn_b, "enc_%d_conv_bn_b", i);
        layer.conv_bn_mean        = create_tensor(PARAKEET_TENSOR_ENC_CONV_BN_MEAN, ggml_new_tensor_1d(ctx, GGML_TYPE_F32, n_audio_state), i);
        layer.conv_bn_var         = create_tensor(PARAKEET_TENSOR_ENC_CONV_BN_VAR, ggml_new_tensor_1d(ctx, GGML_TYPE_F32, n_audio_state), i);
        ggml_format_name(layer.conv_bn_var, "enc_%d_conv_bn_var", i);
        layer.conv_bn_num_batches = create_tensor(PARAKEET_TENSOR_ENC_CONV_BN_NUM_BATCHES, ggml_new_tensor_1d(ctx, GGML_TYPE_I32, 1), i);
        layer.conv_pw2_w          = create_tensor(PARAKEET_TENSOR_ENC_CONV_PW2_WEIGHT, ggml_new_tensor_2d(ctx, wtype, n_audio_state, n_audio_state), i);
        ggml_format_name(layer.conv_pw2_w, "enc_%d_conv_pw2_w", i);

        // Self attention
        layer.norm_attn_w      = create_tensor(PARAKEET_TENSOR_ENC_NORM_ATTN_WEIGHT, ggml_new_tensor_1d(ctx, GGML_TYPE_F32, n_audio_state), i);
        layer.norm_attn_b      = create_tensor(PARAKEET_TENSOR_ENC_NORM_ATTN_BIAS, ggml_new_tensor_1d(ctx, GGML_TYPE_F32, n_audio_state), i);
        layer.attn_pos_bias_u  = create_tensor(PARAKEET_TENSOR_ENC_ATTN_POS_BIAS_U, ggml_new_tensor_2d(ctx, GGML_TYPE_F32, hparams.n_audio_state / hparams.n_audio_head, hparams.n_audio_head), i);
        layer.attn_pos_bias_v  = create_tensor(PARAKEET_TENSOR_ENC_ATTN_POS_BIAS_V, ggml_new_tensor_2d(ctx, GGML_TYPE_F32, hparams.n_audio_state / hparams.n_audio_head, hparams.n_audio_head), i);
        layer.attn_q_w         = create_tensor(PARAKEET_TENSOR_ENC_ATTN_Q_WEIGHT, ggml_new_tensor_2d(ctx, wtype, n_audio_state, n_audio_state), i);
        layer.attn_k_w         = create_tensor(PARAKEET_TENSOR_ENC_ATTN_K_WEIGHT, ggml_new_tensor_2d(ctx, wtype, n_audio_state, n_audio_state), i);
        layer.attn_v_w         = create_tensor(PARAKEET_TENSOR_ENC_ATTN_V_WEIGHT, ggml_new_tensor_2d(ctx, wtype, n_audio_state, n_audio_state), i);
        layer.attn_out_w       = create_tensor(PARAKEET_TENSOR_ENC_ATTN_OUT_WEIGHT, ggml_new_tensor_2d(ctx, wtype, n_audio_state, n_audio_state), i);
        layer.attn_pos_w       = create_tensor(PARAKEET_TENSOR_ENC_ATTN_POS_WEIGHT, ggml_new_tensor_2d(ctx, wtype, n_audio_state, n_audio_state), i);
        ggml_format_name(layer.attn_pos_w, "enc_%d_attn_pos_w", i);

        // Feed forward 2
        layer.norm_ff2_w    = create_tensor(PARAKEET_TENSOR_ENC_NORM_FF2_WEIGHT, ggml_new_tensor_1d(ctx, GGML_TYPE_F32, n_audio_state), i);
        layer.norm_ff2_b    = create_tensor(PARAKEET_TENSOR_ENC_NORM_FF2_BIAS, ggml_new_tensor_1d(ctx, GGML_TYPE_F32, n_audio_state), i);
        layer.ff2_linear1_w = create_tensor(PARAKEET_TENSOR_ENC_FF2_LINEAR1_WEIGHT, ggml_new_tensor_2d(ctx, wtype, n_audio_state, 4*n_audio_state), i);
        layer.ff2_linear2_w = create_tensor(PARAKEET_TENSOR_ENC_FF2_LINEAR2_WEIGHT, ggml_new_tensor_2d(ctx, wtype, 4*n_audio_state, n_audio_state), i);

        // Output norm
        layer.norm_out_w = create_tensor(PARAKEET_TENSOR_ENC_NORM_OUT_WEIGHT, ggml_new_tensor_1d(ctx, GGML_TYPE_F32, n_audio_state), i);
        layer.norm_out_b = create_tensor(PARAKEET_TENSOR_ENC_NORM_OUT_BIAS, ggml_new_tensor_1d(ctx, GGML_TYPE_F32, n_audio_state), i);
    }

    // Prediction network (decoder)
    const int dec_hidden   = hparams.n_pred_dim;
    const int n_pred_embed = hparams.n_vocab + 1;                            // vocab + blank token
    const int n_lstm_gates = 4 * dec_hidden;                                 // 4 LSTM gates
    const int n_joint_out  = hparams.n_vocab + hparams.n_tdt_durations + 1;  // vocab + durations + blank

    // The prediction/joint hidden dimension is 640, which is not a multiple of the
    // K-quant block size (256). For K-quant models, we keep these tensors at F32.
    const int blck         = ggml_blck_size(wtype);
    const ggml_type pred_wtype = (blck > 1 && dec_hidden % blck != 0) ? GGML_TYPE_F32 : wtype;
    const ggml_type join_wtype = pred_wtype;

    model.prediction.embed_w = create_tensor(PARAKEET_TENSOR_PRED_EMBED_WEIGHT, ggml_new_tensor_2d(ctx, pred_wtype, dec_hidden, n_pred_embed));
    model.prediction.lstm_layer.resize(hparams.n_pred_layers);
    for (int i = 0; i < hparams.n_pred_layers; ++i) {
        auto & layer = model.prediction.lstm_layer[i];
        layer.ih_w = create_tensor(PARAKEET_TENSOR_PRED_LSTM_WEIGHT_IH, ggml_new_tensor_2d(ctx, pred_wtype, dec_hidden, n_lstm_gates), i);
        ggml_format_name(layer.ih_w, "pred_%d_ih_w", i);

        layer.hh_w = create_tensor(PARAKEET_TENSOR_PRED_LSTM_WEIGHT_HH, ggml_new_tensor_2d(ctx, pred_wtype, dec_hidden, n_lstm_gates), i);
        ggml_format_name(layer.hh_w, "pred_%d_hh_w", i);

        layer.b_h = create_tensor(PARAKEET_TENSOR_PRED_LSTM_BIAS_H, ggml_new_tensor_1d(ctx, GGML_TYPE_F32, n_lstm_gates), i);
        ggml_format_name(layer.b_h, "pred_%d_b_h", i);
    }

    // Joint network
    model.joint.pred_w = create_tensor(PARAKEET_TENSOR_JOINT_PRED_WEIGHT, ggml_new_tensor_2d(ctx, join_wtype, dec_hidden, dec_hidden));
    ggml_set_name(model.joint.pred_w, "pred_w");
    model.joint.pred_b = create_tensor(PARAKEET_TENSOR_JOINT_PRED_BIAS, ggml_new_tensor_1d(ctx, GGML_TYPE_F32, dec_hidden));
    ggml_set_name(model.joint.pred_b, "pred_b");
    model.joint.enc_w  = create_tensor(PARAKEET_TENSOR_JOINT_ENC_WEIGHT, ggml_new_tensor_2d(ctx, wtype, n_audio_state, dec_hidden));
    ggml_set_name(model.joint.enc_w, "enc_w");
    model.joint.enc_b  = create_tensor(PARAKEET_TENSOR_JOINT_ENC_BIAS, ggml_new_tensor_1d(ctx, GGML_TYPE_F32, dec_hidden));
    ggml_set_name(model.joint.enc_b, "enc_b");
    model.joint.net_w  = create_tensor(PARAKEET_TENSOR_JOINT_NET_WEIGHT, ggml_new_tensor_2d(ctx, join_wtype, dec_hidden, n_joint_out));
    ggml_set_name(model.joint.net_w, "net_w");
    model.joint.net_b  = create_tensor(PARAKEET_TENSOR_JOINT_NET_BIAS, ggml_new_tensor_1d(ctx, GGML_TYPE_F32, n_joint_out));
    ggml_set_name(model.joint.net_b, "net_b");

    ggml_free(ctx);

    // allocate tensors in the backend buffers
    for (auto & p : ctx_map) {
        ggml_backend_buffer_type_t buft = p.first;
        ggml_context * ctx = p.second;
        ggml_backend_buffer_t buf = ggml_backend_alloc_ctx_tensors_from_buft(ctx, buft);
        if (buf) {
            wctx.model.buffers.emplace_back(buf);

            size_t size_main = ggml_backend_buffer_get_size(buf);
            PARAKEET_LOG_INFO("%s: %12s total size = %8.2f MB\n", __func__, ggml_backend_buffer_name(buf), size_main / 1e6);
        }
    }

    // load weights
    {
        size_t total_size = 0;

        auto & tensors_map = wctx.model.tensors;
        int & n_loaded = wctx.model.n_loaded;

        n_loaded = 0;

        std::vector<char> read_buf;

        while (true) {
            int32_t n_dims;
            int32_t length;
            int32_t ttype;

            read_safe(loader, n_dims);
            read_safe(loader, length);
            read_safe(loader, ttype);

            if (loader->eof(loader->context)) {
                break;
            }

            int32_t nelements = 1;
            int32_t ne[4] = { 1, 1, 1, 1 };
            for (int i = 0; i < n_dims; ++i) {
                read_safe(loader, ne[i]);
                nelements *= ne[i];
            }

            std::string name;
            std::vector<char> tmp(length); // create a buffer
            loader->read(loader->context, &tmp[0], tmp.size()); // read to buffer
            name.assign(&tmp[0], tmp.size());

            if (tensors_map.find(name) == tensors_map.end()) {
                PARAKEET_LOG_ERROR("%s: unknown tensor '%s' in model file\n", __func__, name.data());
                return false;
            }

            auto tensor = tensors_map[name.data()];

            if (ggml_nelements(tensor) != nelements) {
                PARAKEET_LOG_ERROR("%s: tensor '%s' has wrong size in model file\n", __func__, name.data());
                PARAKEET_LOG_ERROR("%s: shape: [%d, %d, %d], expected: [%d, %d, %d]\n",
                        __func__, ne[0], ne[1], ne[2], (int) tensor->ne[0], (int) tensor->ne[1], (int) tensor->ne[2]);
                return false;
            }

            if (tensor->ne[0] != ne[0] || tensor->ne[1] != ne[1] || tensor->ne[2] != ne[2] || tensor->ne[3] != ne[3]) {
                PARAKEET_LOG_ERROR("%s: tensor '%s' has wrong shape in model file: got [%d, %d, %d, %d], expected [%d, %d, %d, %d]\n",
                        __func__, name.data(), (int) tensor->ne[0], (int) tensor->ne[1], (int) tensor->ne[2], (int) tensor->ne[3], ne[0], ne[1], ne[2], ne[3]);
                return false;
            }

            const size_t bpe = ggml_type_size(ggml_type(ttype));

            if ((nelements*bpe)/ggml_blck_size(tensor->type) != ggml_nbytes(tensor)) {
                PARAKEET_LOG_ERROR("%s: tensor '%s' has wrong size in model file: got %zu, expected %zu\n",
                        __func__, name.data(), ggml_nbytes(tensor), nelements*bpe);
                return false;
            }

            if (ggml_backend_buffer_is_host(tensor->buffer)) {
                // for the CPU and Metal backend, we can read directly into the tensor
                loader->read(loader->context, tensor->data, ggml_nbytes(tensor));
                BYTESWAP_TENSOR(tensor);
            } else {
                // read into a temporary buffer first, then copy to device memory
                read_buf.resize(ggml_nbytes(tensor));

                loader->read(loader->context, read_buf.data(), read_buf.size());

                ggml_backend_tensor_set(tensor, read_buf.data(), 0, ggml_nbytes(tensor));
            }

            total_size += ggml_nbytes(tensor);
            n_loaded++;
        }

        PARAKEET_LOG_INFO("%s: model size    = %7.2f MB\n", __func__, total_size/1e6);

        if (n_loaded == 0) {
            PARAKEET_LOG_WARN("%s: WARN no tensors loaded from model file - assuming empty model for testing\n", __func__);
        } else if (n_loaded != (int) tensors_map.size()) {
            PARAKEET_LOG_ERROR("%s: ERROR not all tensors loaded from model file - expected %zu, got %d\n", __func__, tensors_map.size(), n_loaded);
            return false;
        }
    }

    auto & buffers = wctx.model.buffers;
    for (auto & buf : buffers) {
        ggml_backend_buffer_set_usage(buf, GGML_BACKEND_BUFFER_USAGE_WEIGHTS);
    }

    wctx.t_load_us = ggml_time_us() - t_start_us;

    return true;
}

// conv subsampling + conformer encoder
static struct ggml_cgraph * parakeet_build_graph_encode(parakeet_context & pctx, parakeet_state & pstate) {
    const auto & model    = pctx.model;
    const auto & hparams  = model.hparams;
    const int n_mel_time  = pstate.n_audio_ctx > 0 ? pstate.n_audio_ctx : hparams.n_audio_ctx;
    const int n_mels      = hparams.n_mels;
    const int n_layer     = hparams.n_audio_layer;
    const int n_state     = hparams.n_audio_state;
    const float fc_factor = 0.5f;

    struct ggml_init_params params = {
        /*.mem_size   =*/ pstate.sched_encode.meta.size(),
        /*.mem_buffer =*/ pstate.sched_encode.meta.data(),
        /*.no_alloc   =*/ true,
    };

    struct ggml_context * ctx0 = ggml_init(params);
    ggml_cgraph * gf = ggml_new_graph_custom(ctx0, PARAKEET_MAX_NODES, false);

    // Conv subsampling

    // [freq, time]
    struct ggml_tensor * mel = ggml_new_tensor_4d(ctx0, GGML_TYPE_F32, n_mels, n_mel_time, 1, 1);
    ggml_set_name(mel, "mel");
    ggml_set_input(mel);

    // [freq, time, channels, batch]
    struct ggml_tensor * cur = ggml_conv_2d(ctx0, model.enc_pre_conv_0_w, mel, 2, 2, 1, 1, 1, 1);
    cur = ggml_add(ctx0, cur, model.enc_pre_conv_0_b);
    ggml_set_name(cur, "pre_conv_0");

    cur = ggml_relu(ctx0, cur);
    ggml_set_name(cur, "pre_conv_0_relu");

    // [freq, time, channels, batch]
    cur = ggml_conv_2d_dw_direct(ctx0, model.enc_pre_conv_2_w, cur, 2, 2, 1, 1, 1, 1);
    cur = ggml_add(ctx0, cur, model.enc_pre_conv_2_b);
    ggml_set_name(cur, "pre_conv_2");

    // [freq, time, channels, batch]
    cur = ggml_conv_2d(ctx0, model.enc_pre_conv_3_w, cur, 1, 1, 0, 0, 1, 1);
    cur = ggml_add(ctx0, cur, model.enc_pre_conv_3_b);
    ggml_set_name(cur, "pre_conv_3");

    cur = ggml_relu(ctx0, cur);
    ggml_set_name(cur, "pre_conv_3_relu");

    // [freq, time, channels, batch]
    cur = ggml_conv_2d_dw_direct(ctx0, model.enc_pre_conv_5_w, cur, 2, 2, 1, 1, 1, 1);
    ggml_set_name(cur, "pre_conv_5_direct");
    cur = ggml_add(ctx0, cur, model.enc_pre_conv_5_b);
    ggml_set_name(cur, "pre_conv_5");

    // [freq, time, channels, batch]
    cur = ggml_conv_2d(ctx0, model.enc_pre_conv_6_w, cur, 1, 1, 0, 0, 1, 1);
    cur = ggml_add(ctx0, cur, model.enc_pre_conv_6_b);
    ggml_set_name(cur, "pre_conv_6");

    cur = ggml_relu(ctx0, cur);
    ggml_set_name(cur, "pre_conv_6_relu");

    // [freq, time, chan]
    cur = ggml_permute(ctx0, cur, 0, 2, 1, 3);
    // [freq, chan, time]
    cur = ggml_cont(ctx0, cur);

    const int n_freq   = cur->ne[0]; // 16
    const int n_chan   = cur->ne[1]; // 256
    const int n_frames = cur->ne[2]; // time

    // [freq, time, chan, batch] -> [(freq * chan), time]
    cur = ggml_reshape_2d(ctx0, cur, n_freq * n_chan, n_frames);

    cur = ggml_mul_mat(ctx0, model.enc_pre_out_w, cur);
    cur = ggml_add(ctx0, cur, model.enc_pre_out_b);

    ggml_set_name(cur, "pre_enc_out");

    // Encoder
    // cur: [n_state, n_enc_time]

    const int  n_time      = cur->ne[1];
    const bool local_attn  = n_time > PARAKEET_LOCAL_ATTN_THRESHOLD;
    const int  att_left    = local_attn ? PARAKEET_LOCAL_ATTN_WINDOW : n_time - 1;
    const int  att_right   = local_attn ? PARAKEET_LOCAL_ATTN_WINDOW : n_time - 1;
    const int  window_size = local_attn ? att_left + att_right + 1 : 2 * n_time - 1;
    const int  d_half      = n_state / 2;
    const int  mask_dim    = local_attn ? window_size : n_time;

    // mask [key, n_time]
    struct ggml_tensor * attn_mask = ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, mask_dim, n_time);
    ggml_set_name(attn_mask, "attn_mask");
    ggml_set_input(attn_mask);

    struct ggml_tensor * local_mask = nullptr;
    if (local_attn) {
        const int chunk = att_left + att_right;
        local_mask = ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, chunk + window_size - 1, chunk);
        ggml_set_name(local_mask, "local_mask");
        ggml_set_input(local_mask);
    }

    struct ggml_tensor * pos_freqs = ggml_new_tensor_1d(ctx0, GGML_TYPE_F32, d_half);
    ggml_set_name(pos_freqs, "pos_freqs");
    ggml_set_input(pos_freqs);

    struct ggml_tensor * rel_positions = ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, 1, window_size);
    ggml_set_name(rel_positions, "rel_positions");
    ggml_set_input(rel_positions);

    struct ggml_tensor * freqs = ggml_repeat_4d(ctx0, pos_freqs, d_half, window_size, 1, 1);
    struct ggml_tensor * theta = ggml_mul(ctx0, freqs, rel_positions);

    struct ggml_tensor * sin_t = ggml_reshape_3d(ctx0, ggml_sin(ctx0, theta), 1, d_half, window_size);
    struct ggml_tensor * cos_t = ggml_reshape_3d(ctx0, ggml_cos(ctx0, theta), 1, d_half, window_size);
    // [n_state, window_size]
    struct ggml_tensor * pos_emb = ggml_reshape_2d(ctx0, ggml_cont(ctx0, ggml_concat(ctx0, sin_t, cos_t, 0)), n_state, window_size);
    ggml_set_name(pos_emb, "pos_emb");

    for (int il = 0; il < n_layer; ++il) {
        const auto & layer = model.layers[il];

        // FFN1
        {
            struct ggml_tensor * residual = cur;
            ggml_format_name(cur, "enc_%d_res", il);

            // norm
            cur = ggml_norm(ctx0, cur, hparams.eps);
            cur = ggml_add(ctx0, ggml_mul(ctx0, cur, layer.norm_ff1_w), layer.norm_ff1_b);
            ggml_format_name(cur, "enc_%d_ffn_norm_1", il);

            // ffn_1
            cur = ggml_mul_mat(ctx0, layer.ff1_linear1_w, cur);
            cur = ggml_silu(ctx0, cur);
            ggml_format_name(cur, "enc_%d_silu", il);

            cur = ggml_mul_mat(ctx0, layer.ff1_linear2_w, cur);
            ggml_format_name(cur, "enc_%d_ffn_1", il);

            cur = ggml_add(ctx0, residual, ggml_scale(ctx0, cur, fc_factor));
            ggml_format_name(cur, "enc_%d_res_ffn", il);
        }

        // self attention block using relative positional encoding computed in graph.
        {
            // [feat, time_frames, 1, 1]
            struct ggml_tensor * residual = cur;

            cur = ggml_norm(ctx0, cur, hparams.eps);
            cur = ggml_add(ctx0, ggml_mul(ctx0, cur, layer.norm_attn_w), layer.norm_attn_b);
            ggml_format_name(cur, "enc_%d_attn_norm", il);

            const int n_head = hparams.n_audio_head;
            const int d_head = n_state / n_head;

            // [feat, time_frames, 1, 1]
            struct ggml_tensor * Q_cur = ggml_mul_mat(ctx0, layer.attn_q_w, cur);
            struct ggml_tensor * K_cur = ggml_mul_mat(ctx0, layer.attn_k_w, cur);
            struct ggml_tensor * V_cur = ggml_mul_mat(ctx0, layer.attn_v_w, cur);

            Q_cur = ggml_reshape_3d(ctx0, Q_cur, d_head, n_head, n_time);
            K_cur = ggml_reshape_3d(ctx0, K_cur, d_head, n_head, n_time);
            V_cur = ggml_reshape_3d(ctx0, V_cur, d_head, n_head, n_time);

            struct ggml_tensor * pos = ggml_mul_mat(ctx0, layer.attn_pos_w, pos_emb);
            pos = ggml_reshape_3d(ctx0, pos, d_head, n_head, window_size);
            pos = ggml_cont(ctx0, ggml_permute(ctx0, pos, 0, 2, 1, 3));

            if (local_attn) {
                const int  chunk         = att_left + att_right;
                const int  n_group       = (n_time + chunk - 1) / chunk;
                const int  n_time_padded = n_group * chunk;
                const int  n_kv_chunk    = chunk + window_size - 1;
                const int  n_kv_dense    = n_kv_chunk * n_group;
                const bool need_padding  = n_time_padded > n_time;

                Q_cur = ggml_cont(ctx0, ggml_permute(ctx0, Q_cur, 0, 2, 1, 3));
                K_cur = ggml_cont(ctx0, ggml_permute(ctx0, K_cur, 0, 2, 1, 3));
                V_cur = ggml_cont(ctx0, ggml_permute(ctx0, V_cur, 0, 2, 1, 3));

                // content bias
                struct ggml_tensor * bias_u = ggml_reshape_3d(ctx0, layer.attn_pos_bias_u, d_head, 1, n_head);
                struct ggml_tensor * Q_u = ggml_add(ctx0, Q_cur, bias_u);

                // position bias
                struct ggml_tensor * bias_v = ggml_reshape_3d(ctx0, layer.attn_pos_bias_v, d_head, 1, n_head);
                struct ggml_tensor * Q_v = ggml_add(ctx0, Q_cur, bias_v);

                // right pad the time_frame.
                struct ggml_tensor * Q_u_padded = need_padding ?
                    ggml_pad_ext(ctx0, Q_u, 0, 0, 0, n_time_padded - n_time, 0, 0, 0, 0) : Q_u;
                Q_u_padded = ggml_reshape_4d(ctx0, Q_u_padded, d_head, chunk, n_group, n_head);

                // Add padding to front and back (for the first timeframe and the last timeframe).
                struct ggml_tensor * K_padded = ggml_pad_ext(ctx0, K_cur, 0, 0, att_left, att_right, 0, 0, 0, 0);

                // pad time axis to match n_kv_dense if needed.
                if (n_kv_dense > K_padded->ne[1]) {
                    K_padded = ggml_pad_ext(ctx0, K_padded, 0, 0, 0, n_kv_dense - K_padded->ne[1], 0, 0, 0, 0);
                }

                // Create a 4d tensor where each group spans a wide window of
                // 512 keys (n_kv_chunk), but moving to the next group (nb[2])
                // only jumps forward by 256 frames (chunk * nb[1]). This creates
                // a 256 frame overlap, shared keys in RAM without copies.
                struct ggml_tensor * K_chunk = ggml_view_4d(ctx0, K_padded,
                        d_head, n_kv_chunk, n_group, n_head,
                        K_padded->nb[1],
                        (size_t) chunk * K_padded->nb[1],
                        K_padded->nb[2],
                        0);
                K_chunk = ggml_cont(ctx0, K_chunk);

                struct ggml_tensor * content_scores = ggml_mul_mat(ctx0, K_chunk, Q_u_padded);

                // The above mul_mat operation, combined with K_chunk's overlapping
                // frames, produces a dense matrix. But some of the results in
                // this matrix were computed for keys that aren't part of that
                // query's window. So we shift each row to keep only the results
                // that we want.
                content_scores = ggml_view_4d(ctx0, content_scores,
                        window_size, chunk, n_group, n_head,
                        (size_t) (chunk + window_size) * content_scores->nb[0],
                        content_scores->nb[2],
                        content_scores->nb[3],
                        0);
                content_scores = ggml_cont(ctx0, content_scores);

                // ungrouping.
                content_scores = ggml_reshape_3d(ctx0, content_scores, window_size, n_time_padded, n_head);

                // remove padding if padding was applied (truncating to n_time).
                if (need_padding) {
                    content_scores = ggml_view_3d(ctx0, content_scores,
                            window_size, n_time, n_head,
                            content_scores->nb[1],
                            content_scores->nb[2],
                            0);
                }

                struct ggml_tensor * rel_pos_scores = ggml_mul_mat(ctx0, pos, Q_v);

                // attention_score = content similarity + relative position scores
                struct ggml_tensor * attn_scores = ggml_add(ctx0, content_scores, rel_pos_scores);

                attn_scores = ggml_soft_max_ext(ctx0, attn_scores, attn_mask, 1.0f / std::sqrt(d_head), 0.0f);

                // right pad the probabilites.
                struct ggml_tensor * probs_padded = need_padding ?
                    ggml_pad_ext(ctx0, attn_scores, 0, 0, 0, n_time_padded - n_time, 0, 0, 0, 0) : attn_scores;

                probs_padded = ggml_reshape_4d(ctx0, probs_padded, window_size, chunk, n_group, n_head);
                probs_padded = ggml_pad_ext(ctx0, probs_padded, 0, chunk, 0, 0, 0, 0, 0, 0);
                probs_padded = ggml_view_4d(ctx0, probs_padded,
                        n_kv_chunk, chunk, n_group, n_head,
                        (size_t) n_kv_chunk * probs_padded->nb[0],
                        probs_padded->nb[2],
                        probs_padded->nb[3],
                        0);
                probs_padded = ggml_cont(ctx0, probs_padded);
                probs_padded = ggml_mul(ctx0, probs_padded, local_mask);

                // Add padding to front and back (for the first timeframe and the last timeframe).
                struct ggml_tensor * V_padded = ggml_pad_ext(ctx0, V_cur, 0, 0, att_left, att_right, 0, 0, 0, 0);

                // pad time axis to match n_kv_dense if needed.
                if (n_kv_dense > V_padded->ne[1]) {
                    V_padded = ggml_pad_ext(ctx0, V_padded, 0, 0, 0, n_kv_dense - V_padded->ne[1], 0, 0, 0, 0);
                }

                V_padded = ggml_cont(ctx0, ggml_transpose(ctx0, V_padded));

                struct ggml_tensor * V_chunk = ggml_view_4d(ctx0, V_padded,
                        n_kv_chunk, d_head, n_group, n_head,
                        V_padded->nb[1],
                        (size_t) chunk * V_padded->nb[0],
                        V_padded->nb[2],
                        0);
                V_chunk = ggml_cont(ctx0, V_chunk);

                cur = ggml_mul_mat(ctx0, V_chunk, probs_padded);
                // ungroup.
                cur = ggml_reshape_3d(ctx0, cur, d_head, n_time_padded, n_head);
                // unpad
                if (need_padding) {
                    cur = ggml_view_3d(ctx0, cur, d_head, n_time, n_head, cur->nb[1], cur->nb[2], 0);
                }

                cur = ggml_cont(ctx0, ggml_permute(ctx0, cur, 0, 2, 1, 3));
                cur = ggml_reshape_2d(ctx0, cur, n_state, n_time);
                cur = ggml_mul_mat(ctx0, layer.attn_out_w, cur);
            } else {
                struct ggml_tensor * Q_u = ggml_add(ctx0, Q_cur, layer.attn_pos_bias_u);
                ggml_format_name(Q_u, "enc_%d_attn_q_u", il);

                struct ggml_tensor * K_prep = ggml_permute(ctx0, K_cur, 0, 2, 1, 3);
                struct ggml_tensor * Q_prep = ggml_permute(ctx0, Q_u,   0, 2, 1, 3);
                struct ggml_tensor * content_scores = ggml_mul_mat(ctx0, K_prep, Q_prep);
                ggml_format_name(content_scores, "enc_%d_attn_content_scores", il);

                struct ggml_tensor * Q_v = ggml_add(ctx0, Q_cur, layer.attn_pos_bias_v);
                ggml_format_name(Q_v, "enc_%d_attn_q_v", il);

                Q_v = ggml_permute(ctx0, Q_v, 0, 2, 1, 3);
                Q_v = ggml_cont(ctx0, Q_v);
                ggml_format_name(Q_v, "enc_%d_attn_q_v_perm", il);

                struct ggml_tensor * rel_pos_scores = ggml_mul_mat(ctx0, pos, Q_v);
                ggml_format_name(rel_pos_scores, "enc_%d_attn_rel_pos", il);

                // Relative position shifting is performed in the following block.
                // Some more details on the operations performed below can be found here:
                // https://github.com/danbev/learning-ai/blob/main/notes/whisper/parakeet.md#relative-position-shift
                {
                    const auto pos_window = rel_pos_scores->ne[0];
                    const auto n_frame    = rel_pos_scores->ne[1];
                    const auto n_head_cur = rel_pos_scores->ne[2];

                    rel_pos_scores = ggml_pad(ctx0, rel_pos_scores, 1, 0, 0, 0);
                    rel_pos_scores = ggml_roll(ctx0, rel_pos_scores, 1, 0, 0, 0);

                    rel_pos_scores = ggml_reshape_3d(ctx0, rel_pos_scores, n_frame, pos_window + 1, n_head_cur);
                    ggml_format_name(rel_pos_scores, "enc_%d_attn_rel_pos_reshaped", il);

                    int center = pos_window / 2;
                    size_t offset = rel_pos_scores->nb[0] * (center+1);

                    rel_pos_scores = ggml_view_3d(ctx0, rel_pos_scores,
                                                  n_frame, pos_window, n_head_cur,
                                                  (pos_window) * 4,
                                                  rel_pos_scores->nb[2],
                                                  offset);

                    ggml_format_name(rel_pos_scores, "enc_%d_attn_rel_pos_shifted", il);

                    rel_pos_scores = ggml_view_3d(ctx0, rel_pos_scores,
                                                  content_scores->ne[0],
                                                  content_scores->ne[1],
                                                  rel_pos_scores->ne[2],
                                                  rel_pos_scores->nb[1],
                                                  rel_pos_scores->nb[2],
                                                  0);
                    rel_pos_scores = ggml_cont(ctx0, rel_pos_scores);
                    ggml_format_name(rel_pos_scores, "enc_%d_attn_rel_pos_shifted_view", il);
                }

                struct ggml_tensor * attn_scores = ggml_add(ctx0, content_scores, rel_pos_scores);
                ggml_format_name(attn_scores, "enc_%d_attn_scores", il);
                attn_scores = ggml_scale(ctx0, attn_scores, 1.0f / std::sqrt(d_head));
                attn_scores = ggml_add(ctx0, attn_scores, attn_mask);
                ggml_format_name(attn_scores, "enc_%d_attn_scores_scaled", il);

                struct ggml_tensor * probs = ggml_soft_max(ctx0, attn_scores);
                ggml_format_name(probs, "enc_%d_attn_probs", il);

                V_cur = ggml_cont(ctx0, ggml_permute(ctx0, V_cur, 1, 2, 0, 3));
                ggml_format_name(V_cur, "enc_%d_attn_v_cur", il);
                cur = ggml_mul_mat(ctx0, probs, V_cur);
                ggml_format_name(cur, "enc_%d_attn_inp", il);

                cur = ggml_permute(ctx0, cur, 2, 0, 1, 3);
                cur = ggml_cont_2d(ctx0, cur, n_state, n_time);
                cur = ggml_mul_mat(ctx0, layer.attn_out_w, cur);
            }
            ggml_format_name(cur, "enc_%d_attn_out", il);

            cur = ggml_add(ctx0, residual, cur);
            ggml_format_name(cur, "enc_%d_attn_res", il);
        }

        // Convolution
        {
            struct ggml_tensor * residual = cur;
            ggml_format_name(cur, "enc_%d_residual_conv", il);

            cur = ggml_norm(ctx0, cur, hparams.eps);
            cur = ggml_add(ctx0, ggml_mul(ctx0, cur, layer.norm_conv_w), layer.norm_conv_b);
            ggml_format_name(cur, "enc_%d_norm_conv", il);

            // pointwise 1d convolution: [1024, 138] -> [2048, 138]
            cur = ggml_mul_mat(ctx0, layer.conv_pw1_w, cur);
            ggml_format_name(cur, "enc_%d_conv_pw1", il);

            {
                int64_t d = cur->ne[0] / 2;
                struct ggml_tensor * signal = ggml_view_2d(ctx0, cur, d, cur->ne[1], cur->nb[1], 0);
                struct ggml_tensor * gate   = ggml_view_2d(ctx0, cur, d, cur->ne[1], cur->nb[1], d * cur->nb[0]);

                cur = ggml_mul(ctx0, signal, ggml_sigmoid(ctx0, gate));
                ggml_format_name(cur, "enc_%d_conv_glu", il);
            }

            cur = ggml_cont(ctx0, ggml_transpose(ctx0, cur));

            // use ggml_ssm_conv for f32 precision
            const int dw_pad = (hparams.n_conv_kernel - 1) / 2;
            cur = ggml_pad(ctx0, cur, dw_pad, 0, 0, 0);
            cur = ggml_roll(ctx0, cur, dw_pad, 0, 0, 0);
            cur = ggml_pad(ctx0, cur, dw_pad, 0, 0, 0);
            ggml_format_name(cur, "enc_%d_conv_dw_pad", il);

            cur = ggml_ssm_conv(ctx0, cur, layer.conv_dw_w);
            ggml_format_name(cur, "enc_%d_conv_1d_dw", il);

            cur = ggml_sub(ctx0, cur, layer.conv_bn_mean);
            struct ggml_tensor * std = ggml_sqrt(ctx0, layer.conv_bn_var);
            cur = ggml_div(ctx0, cur, std);
            cur = ggml_add(ctx0, ggml_mul(ctx0, cur, layer.conv_bn_w), layer.conv_bn_b);
            ggml_format_name(cur, "enc_%d_conv_bn", il);

            cur = ggml_silu(ctx0, cur);
            ggml_format_name(cur, "enc_%d_conv_silu", il);

            cur = ggml_mul_mat(ctx0, layer.conv_pw2_w, cur);
            ggml_format_name(cur, "enc_%d_conv_pw2", il);

            cur = ggml_add(ctx0, residual, cur);
            ggml_format_name(cur, "enc_%d_conv_res", il);
        }

        // FFN2
        {
            struct ggml_tensor * residual = cur;
            cur = ggml_norm(ctx0, cur, hparams.eps);
            cur = ggml_add(ctx0, ggml_mul(ctx0, cur, layer.norm_ff2_w), layer.norm_ff2_b);
            ggml_format_name(cur, "enc_%d_ffn_norm_2", il);

            cur = ggml_mul_mat(ctx0, layer.ff2_linear1_w, cur);
            cur = ggml_silu(ctx0, cur);
            cur = ggml_mul_mat(ctx0, layer.ff2_linear2_w, cur);
            cur = ggml_add(ctx0, residual, ggml_scale(ctx0, cur, 0.5));
            ggml_format_name(cur, "enc_%d_ffn_res", il);
        }

        cur = ggml_norm(ctx0, cur, hparams.eps);
        cur = ggml_add(ctx0, ggml_mul(ctx0, cur, layer.norm_out_w), layer.norm_out_b);
    }

    ggml_set_name(cur, "encoder_out");
    pstate.n_frames = cur->ne[1];

    struct ggml_tensor * enc_out_view = ggml_view_2d(ctx0, pstate.enc_out, n_state, pstate.n_frames, pstate.enc_out->nb[1], 0);
    ggml_build_forward_expand(gf, ggml_cpy(ctx0, cur, enc_out_view));

    ggml_free(ctx0);

    return gf;
}

static bool parakeet_encode_internal(
        parakeet_context & pctx,
          parakeet_state & pstate,
              const int   mel_offset,
              const int   n_threads,
    ggml_abort_callback   abort_callback,
                   void * abort_callback_data) {
    const int64_t t_start_us = ggml_time_us();

    auto & sched = pstate.sched_encode.sched;

    ggml_cgraph * gf = parakeet_build_graph_encode(pctx, pstate);

    if (!ggml_backend_sched_alloc_graph(sched, gf)) {
        // should never happen as we pre-allocate the memory
        return false;
    }

    // set mel input
    {
        struct ggml_tensor * mel = ggml_graph_get_tensor(gf, "mel");

        const auto & mel_inp = pstate.mel;
        const int n_ctx      = pstate.n_audio_ctx > 0 ? pstate.n_audio_ctx : pctx.model.hparams.n_audio_ctx;

        assert(mel->type == GGML_TYPE_F32);
        assert(mel_inp.n_mel == pctx.model.hparams.n_mels);

        pstate.inp_mel.resize(ggml_nelements(mel));

        float * dst = pstate.inp_mel.data();
        memset(dst, 0, ggml_nbytes(mel));

        const int i0 = std::min(mel_offset,         mel_inp.n_len);
        const int i1 = std::min(mel_offset + n_ctx, mel_inp.n_len);

        memcpy(dst, mel_inp.data.data() + i0 * mel_inp.n_mel, (i1 - i0) * mel_inp.n_mel * sizeof(float));

        ggml_backend_tensor_set(mel, pstate.inp_mel.data(), 0, ggml_nelements(mel)*sizeof(float));
    }

    // set attention mask
    {
        struct ggml_tensor * attn_mask = ggml_graph_get_tensor(gf, "attn_mask");
        const int n_q = attn_mask->ne[1];
        const int n_k = attn_mask->ne[0];

        const int32_t subsampl_factor = pctx.model.hparams.subsampling_factor;
        const int n_tokens_real = (pstate.mel.n_len_org + subsampl_factor - 1) / subsampl_factor;

        std::vector<float> mask_data(n_q * n_k);
        const float mask_value = -1e30f;

        if (n_k == n_q) {   // full attention
            for (int q = 0; q < n_q; ++q) {
                for (int k = 0; k < n_k; ++k) {
                    mask_data[q * n_k + k] = (k >= n_tokens_real) ? mask_value : 0.0f;
                }
            }
        } else {            // local attention
            const int att_left = n_k / 2;
            for (int q = 0; q < n_q; ++q) {
                for (int k = 0; k < n_k; ++k) {
                    const int key = q - att_left + k;
                    mask_data[q * n_k + k] = (key >= 0 && key < n_tokens_real) ? 0.0f : mask_value;
                }
            }
        }
        ggml_backend_tensor_set(attn_mask, mask_data.data(), 0, mask_data.size() * sizeof(float));
    }

    // set local attention skew mask
    if (struct ggml_tensor * local_mask = ggml_graph_get_tensor(gf, "local_mask")) {
        const int n_k = local_mask->ne[0];
        const int n_q = local_mask->ne[1];

        std::vector<float> mask_data(n_q * n_k);
        const int window_size = n_k - n_q + 1;
        for (int q = 0; q < n_q; ++q) {
            for (int k = 0; k < n_k; ++k) {
                const int rel = k - q;
                mask_data[q * n_k + k] = (rel >= 0 && rel < window_size) ? 1.0f : 0.0f;
            }
        }
        ggml_backend_tensor_set(local_mask, mask_data.data(), 0, mask_data.size() * sizeof(float));
    }

    // set positional frequency
    {
        struct ggml_tensor * pos_freqs_t = ggml_graph_get_tensor(gf, "pos_freqs");
        const int d_half      = pos_freqs_t->ne[0];
        const int n_state     = pctx.model.hparams.n_audio_state;
        const float log_10000 = logf(10000.0f);
        std::vector<float> freqs(d_half);
        for (int k = 0; k < d_half; ++k) {
            freqs[k] = expf(-(float(k * 2) * log_10000 / float(n_state)));
        }
        ggml_backend_tensor_set(pos_freqs_t, freqs.data(), 0, freqs.size() * sizeof(float));
    }

    // set relative position offsets
    {
        struct ggml_tensor * rel_pos_t = ggml_graph_get_tensor(gf, "rel_positions");
        const int window_size = rel_pos_t->ne[1];
        std::vector<float> pos(window_size);
        if (window_size == PARAKEET_LOCAL_ATTN_WINDOW * 2 + 1) {
            for (int t = 0; t < window_size; ++t) {
                pos[t] = float(PARAKEET_LOCAL_ATTN_WINDOW - t);
            }
        } else {
            const int n_time = (window_size + 1) / 2;
            for (int t = 0; t < window_size; ++t) {
                pos[t] = float(n_time - 1 - t);
            }
        }
        ggml_backend_tensor_set(rel_pos_t, pos.data(), 0, pos.size() * sizeof(float));
    }

    if (!ggml_graph_compute_helper(sched, gf, n_threads)) {
        return false;
    }

    pstate.t_encode_us += ggml_time_us() - t_start_us;
    pstate.n_encode++;

    return !(abort_callback && abort_callback(abort_callback_data));
}

static bool parakeet_ensure_encode_sched(
        parakeet_context & pctx,
          parakeet_state & pstate,
                    int    n_audio_ctx) {
    if (pstate.sched_encode.sched && pstate.sched_encode_n_audio_ctx == n_audio_ctx) {
        return true;
    }

    parakeet_sched_free(pstate.sched_encode);

    const int32_t prev_n_audio_ctx = pstate.n_audio_ctx;
    pstate.n_audio_ctx = n_audio_ctx;

    const int subsampl_factor = pctx.model.hparams.subsampling_factor;
    const int n_frames_max = (n_audio_ctx + subsampl_factor - 1) / subsampl_factor;
    if (n_frames_max > pstate.enc_out->ne[1]) {
        ggml_backend_buffer_free(pstate.enc_out_buffer);
        pstate.enc_out_buffer = nullptr;
        pstate.enc_out = nullptr;

        if (!parakeet_enc_state_init(pstate, pstate.backends[0], pctx.model.hparams.n_audio_state, n_frames_max)) {
            pstate.sched_encode_n_audio_ctx = 0;
            pstate.n_audio_ctx = prev_n_audio_ctx;
            return false;
        }
    }

    const bool ok = parakeet_sched_graph_init(pstate.sched_encode, pstate.backends,
            [&]() {
                return parakeet_build_graph_encode(pctx, pstate);
            });

    if (!ok) {
        pstate.sched_encode_n_audio_ctx = 0;
        pstate.n_audio_ctx = prev_n_audio_ctx;
        return false;
    }

    pstate.sched_encode_n_audio_ctx = n_audio_ctx;
    return true;
}

static struct ggml_tensor * parakeet_build_graph_lstm_layer(
        struct ggml_context * ctx0,
         struct ggml_cgraph * gf,
         struct ggml_tensor * x_t,       // the current input token embedding
         struct ggml_tensor * w_ih,      // input to hidden weights (4 weight tensors packed)
         struct ggml_tensor * w_hh,      // hidden to hidden weights (4 weight tensors packed)
         struct ggml_tensor * b_h,       // folded ih+hh bias (4 bias tensors packed)
         struct ggml_tensor * h_state,   // this layers hidden state
         struct ggml_tensor * c_state,   // this layers cell state
                        int   li) {      // layer index (for tensor naming)

    ggml_format_name(x_t, "lstm_layer_%d_x_t", li);
    ggml_format_name(h_state, "lstm_layer_%d_h_state", li);
    ggml_format_name(c_state, "lstm_layer_%d_c_state", li);

    // The 4 gates (i, f, o, c) are packed in the same weight tensor.
    struct ggml_tensor * inp_gates = ggml_mul_mat(ctx0, w_ih, x_t);

    // Hidden-to-Hidden Projections are also packed in the same weight tensor.
    // b_h holds the folded ih+hh bias (see parakeet_model_load), so it is
    // the only bias that needs to be added here.
    struct ggml_tensor * hid_gates = ggml_mul_mat(ctx0, w_hh, h_state);
    hid_gates = ggml_add(ctx0, hid_gates, b_h);

    // Combine the input and hidden contributions of the gates.
    struct ggml_tensor * gates = ggml_add(ctx0, inp_gates, hid_gates);
    ggml_format_name(gates, "lstm_layer_%d_gates", li);

    const int h_dim = h_state->ne[0];
    const size_t row_size = ggml_row_size(gates->type, h_dim);

    // The gates are packed as [i, f, o, c] (reordered at convert time, see
    // parakeet_model_load), so the three sigmoid-gated outputs (i, f, o) are
    // contiguous and can be computed with a single ggml_sigmoid call.
    struct ggml_tensor * ifo = ggml_sigmoid(ctx0, ggml_view_1d(ctx0, gates, 3 * h_dim, 0));
    ggml_format_name(ifo, "lstm_layer_%d_ifo", li);

    // 1. Input Gate at time t.
    struct ggml_tensor * i_t = ggml_view_1d(ctx0, ifo, h_dim, 0 * row_size);
    ggml_format_name(i_t, "lstm_layer_%d_i_t", li);

    // Forget gate.
    struct ggml_tensor * f_t = ggml_view_1d(ctx0, ifo, h_dim, 1 * row_size);
    ggml_format_name(f_t, "lstm_layer_%d_f_t", li);

    // Output gate.
    struct ggml_tensor * o_t = ggml_view_1d(ctx0, ifo, h_dim, 2 * row_size);
    ggml_format_name(o_t, "lstm_layer_%d_o_t", li);

    // Cell gate.
    struct ggml_tensor * c_t = ggml_tanh(ctx0, ggml_view_1d(ctx0, gates, h_dim, 3 * row_size));
    ggml_format_name(c_t, "lstm_layer_%d_c_t", li);

    // Calculate the new cell state.
    struct ggml_tensor * c_new = ggml_add(ctx0,
        ggml_mul(ctx0, f_t, c_state), // apply forget gate to cell state.
        ggml_mul(ctx0, i_t, c_t));    // apply input gate to cell gate.
    ggml_build_forward_expand(gf, ggml_cpy(ctx0, c_new, c_state));

    // Calculate the new hidden state.
    struct ggml_tensor * h_new = ggml_mul(ctx0, o_t, ggml_tanh(ctx0, c_new));
    ggml_set_output(h_new);
    ggml_format_name(h_new, "lstm_layer_%d_h_new", li);
    ggml_build_forward_expand(gf, ggml_cpy(ctx0, h_new, h_state));

    return h_new;
}

static struct ggml_cgraph * parakeet_build_graph_prediction(
         parakeet_context & pctx,
           parakeet_state & pstate,
     const parakeet_batch & batch,
                    bool   worst_case) {
    GGML_UNUSED(worst_case);
    const auto & model   = pctx.model;
    const auto & hparams = model.hparams;
    const int n_tokens   = batch.n_tokens;

    struct ggml_init_params params = {
        /*.mem_size   =*/ pstate.sched_decode.meta.size(),
        /*.mem_buffer =*/ pstate.sched_decode.meta.data(),
        /*.no_alloc   =*/ true,
    };

    struct ggml_context * ctx0 = ggml_init(params);
    ggml_cgraph * gf = ggml_new_graph_custom(ctx0, PARAKEET_MAX_NODES, false);

    // Prediction Network
    struct ggml_tensor * token = ggml_new_tensor_1d(ctx0, GGML_TYPE_I32, n_tokens);
    ggml_set_name(token, "token_inp");
    ggml_set_input(token);

    struct ggml_tensor * token_embd = ggml_get_rows(ctx0, model.prediction.embed_w, token);

    struct ggml_tensor * inpL = token_embd;

    for (int il = 0; il < hparams.n_pred_layers; ++il) {
        inpL = parakeet_build_graph_lstm_layer(ctx0, gf, inpL,
                model.prediction.lstm_layer[il].ih_w,
                model.prediction.lstm_layer[il].hh_w,
                model.prediction.lstm_layer[il].b_h,
                pstate.lstm_state.layer[il].h_state,
                pstate.lstm_state.layer[il].c_state,
                il);
    }

    struct ggml_tensor * pred_out = inpL;
    ggml_format_name(pred_out, "lstm_pred_out");

    // Project the prediction network output to the joint network hidden dimension.
    struct ggml_tensor * pred = ggml_mul_mat(ctx0, model.joint.pred_w, pred_out);
    pred = ggml_add(ctx0, pred, model.joint.pred_b);
    ggml_set_name(pred, "h_pred");

    ggml_build_forward_expand(gf, ggml_cpy(ctx0, pred, pstate.pred_out));

    ggml_free(ctx0);

    return gf;
}

static struct ggml_cgraph * parakeet_build_graph_joint(
         parakeet_context & pctx,
           parakeet_state & pstate,
     const parakeet_batch & batch,
                     bool   worst_case) {
    GGML_UNUSED(worst_case);
    const auto & model   = pctx.model;
    const auto & hparams = model.hparams;

    struct ggml_init_params params = {
        /*.mem_size   =*/ pstate.sched_decode.meta.size(),
        /*.mem_buffer =*/ pstate.sched_decode.meta.data(),
        /*.no_alloc   =*/ true,
    };

    struct ggml_context * ctx0 = ggml_init(params);
    ggml_cgraph * gf = ggml_new_graph_custom(ctx0, PARAKEET_MAX_NODES, false);

    struct ggml_tensor * pred = pstate.pred_out;
    ggml_format_name(pred, "pred");

    const int t_idx = batch.i_time[0];
    struct ggml_tensor * enc_out = ggml_view_1d(ctx0, pstate.enc_out, hparams.n_audio_state,
            (size_t) t_idx * pstate.enc_out->nb[1]);
    ggml_format_name(enc_out, "enc_out_view");

    // Project the encoder output to the joint network hidden dimension.
    struct ggml_tensor * enc  = ggml_mul_mat(ctx0, model.joint.enc_w, enc_out);
    enc = ggml_add(ctx0, enc, model.joint.enc_b);
    ggml_set_name(enc, "enc");

    struct ggml_tensor * joint = ggml_add(ctx0, enc, pred);
    ggml_set_name(joint, "joint");
    joint = ggml_relu(ctx0, joint);

    struct ggml_tensor * logits = ggml_mul_mat(ctx0, model.joint.net_w, joint);
    logits = ggml_add(ctx0, logits, model.joint.net_b);
    ggml_set_output(logits);
    ggml_set_name(logits, "logits");

    struct ggml_tensor * probs = ggml_soft_max(ctx0, logits);
    struct ggml_tensor * log_probs = ggml_log(ctx0, probs);
    ggml_set_output(log_probs);
    ggml_format_name(log_probs, "log_probs");

    ggml_build_forward_expand(gf, log_probs);

    ggml_free(ctx0);

    return gf;
}

static bool parakeet_predict(
        parakeet_context & pctx,
          parakeet_state & pstate,
    const parakeet_batch & batch,
               const int   n_threads,
     ggml_abort_callback   abort_callback,
                   void  * abort_callback_data) {

    const int n_tokens   = batch.n_tokens;

    const int64_t t_start_us = ggml_time_us();

    {
        auto & sched = pstate.sched_decode.sched;

        const int64_t t_build_start_us = ggml_time_us();
        ggml_cgraph * gf = parakeet_build_graph_prediction(pctx, pstate, batch, false);
        pstate.t_predict_build_us += ggml_time_us() - t_build_start_us;

        const int64_t t_alloc_start_us = ggml_time_us();
        if (!ggml_backend_sched_alloc_graph(sched, gf)) {
            // should never happen as we pre-allocate the memory
            return false;
        }
        pstate.t_predict_alloc_us += ggml_time_us() - t_alloc_start_us;

        // set the inputs
        {
            struct ggml_tensor * token_inp = ggml_graph_get_tensor(gf, "token_inp");
            ggml_backend_tensor_set(token_inp, batch.token, 0, n_tokens * ggml_element_size(token_inp));
        }

        const int64_t t_compute_start_us = ggml_time_us();
        if (!ggml_graph_compute_helper(sched, gf, n_threads)) {
            return false;
        }
        pstate.t_predict_compute_us += ggml_time_us() - t_compute_start_us;
    }

    pstate.t_predict_us += ggml_time_us() - t_start_us;
    pstate.n_predict++;

    return !(abort_callback && abort_callback(abort_callback_data));
}

static bool parakeet_joint(
         parakeet_context & pctx,
           parakeet_state & pstate,
     const parakeet_batch & batch,
                const int   n_threads,
      ggml_abort_callback   abort_callback,
                     void * abort_callback_data) {
    const int64_t t_start_us = ggml_time_us();

    const auto & model   = pctx.model;
    const auto & hparams = model.hparams;
    const int n_tokens   = batch.n_tokens;

    auto & logits_out = pstate.logits;

    struct ggml_tensor * logits;

    {
        auto & sched = pstate.sched_decode.sched;

        ggml_cgraph * gf = parakeet_build_graph_joint(pctx, pstate, batch, false);

        if (!ggml_backend_sched_alloc_graph(sched, gf)) {
            // should never happen as we pre-allocate the memory
            return false;
        }

        logits = ggml_graph_node(gf, -1);

        if (!ggml_graph_compute_helper(sched, gf, n_threads)) {
            return false;
        }

    }

    const int n_logits = hparams.n_vocab + hparams.n_tdt_durations + 1; // one for the blank token
    logits_out.resize(n_tokens * n_logits);
    for (int i = 0; i < n_tokens; i++) {
        if (batch.logits[i] == 0) {
            continue;
        }
        ggml_backend_tensor_get(logits, logits_out.data() + (n_logits*i), sizeof(float)*(n_logits*i), sizeof(float)*n_logits);
    }

    if (batch.n_tokens == 1) {
        pstate.t_decode_us += ggml_time_us() - t_start_us;
        pstate.n_decode++;
    }

    return !(abort_callback && abort_callback(abort_callback_data));
}

static bool is_word_start_token(parakeet_vocab & vocab, parakeet_token token_id) {
    const std::string & token_str = vocab.id_to_token[token_id];
    // check if it starts with the SentencePiece meta-space "▁" (U+2581) or 3-byte UTF-8 character: 0xE2 0x96 0x81
    if (!token_str.empty()) {
        if (token_str.find("\xE2\x96\x81") == 0 || token_str[0] == '_') {
            return true;
        }
    }
    return false;
}

static bool is_punctuation_token(parakeet_vocab & vocab, parakeet_token token_id) {
    const std::string & token_str = vocab.id_to_token[token_id];
    static const std::string punct_chars = ".,!?;:'\"-()[]{}";

    if (token_str.empty()) {
        return false;
    }

    std::string clean_token = token_str;
    if (clean_token.find("\xE2\x96\x81") == 0) {
        clean_token = clean_token.substr(3); // Remove the 3-byte UTF-8 character
    } else if (clean_token[0] == '_') {
        clean_token = clean_token.substr(1);
    }

    return clean_token.length() == 1 && punct_chars.find(clean_token[0]) != std::string::npos;
}

// Collapse punctuation timestamps to match the original Parakeet model.
// Punctuations symbols like ',', '.' and others are not spoken words but the
// model will still produce a duration for these tokens. But since these are
// non-spoken we collapse the timestamps so that they don't have an time duration.
static void refine_timestamps_tdt(parakeet_vocab & vocab, std::vector<parakeet_token_data> & tokens) {
    if (tokens.empty()) {
        return;
    }

    int64_t last_non_punct_t1 = -1;

    for (size_t i = 0; i < tokens.size(); ++i) {
        if (is_punctuation_token(vocab, tokens[i].id)) {
            if (last_non_punct_t1 >= 0) {
                tokens[i].t0 = last_non_punct_t1;
                tokens[i].t1 = last_non_punct_t1;
            }
        } else {
            last_non_punct_t1 = tokens[i].t1;
        }
    }
}

static parakeet_token_data create_token_data(
            parakeet_context & pctx,
              parakeet_state & pstate,
               parakeet_token   token_id,
                          int   duration_idx,
                          int   duration_value,
                          int   frame_index,
                        float   token_logit,
                          int   n_vocab_logits) {

    float token_sum = 0.0f;
    for (int i = 0; i < n_vocab_logits; ++i) {
        token_sum += expf(pstate.logits[i]);
    }
    float token_p = expf(token_logit) / token_sum;

    parakeet_token_data token_data;
    token_data.id = token_id;
    token_data.duration_idx = duration_idx;
    token_data.duration_value = duration_value;
    token_data.frame_index = frame_index;
    token_data.p = token_p;
    token_data.plog = token_logit;
    token_data.t0 = frame_index * pctx.model.hparams.subsampling_factor;
    token_data.t1 = (frame_index + duration_value) * pctx.model.hparams.subsampling_factor;
    token_data.is_word_start = is_word_start_token(pctx.vocab, token_id);

    return token_data;
}

static bool parakeet_decode(
              parakeet_context & pctx,
                parakeet_state & pstate,
                parakeet_batch & batch,
                     const int   n_threads,
    const parakeet_full_params * params = nullptr) {
    const auto & hparams       = pctx.model.hparams;
    const auto & tdt_durations = pctx.model.tdt_durations;

    const int  n_tdt_durations          = hparams.n_tdt_durations;
    const int  n_frames                 = pstate.n_frames;
    const int  blank_id                 = pctx.vocab.token_blank;
    const int  n_vocab_logits           = blank_id + 1;
    const int  max_tokens_per_timestep = hparams.n_max_tokens;

    // time index into the encoder frame (current time frame)
    int t = 0;
    // number of symbols emitted for the current time frame
    int tokens_emitted = 0;

    // Start with the blank token (8192)
    parakeet_token last_token = blank_id;

    PARAKEET_LOG_DEBUG("parakeet_decode: starting decode with n_frames=%d\n", n_frames);

    batch.n_tokens  = 1;
    batch.token[0]  = last_token;
    batch.logits[0] = 1;
    batch.i_time[0] = 0;

    // run the prediction network for the initial blank token. This will
    // initialize the LSTM state and produce an initial hidden state that can
    // be used in the joint network below.
    if (!parakeet_predict(pctx, pstate, batch, n_threads,
            params ? params->abort_callback           : nullptr,
            params ? params->abort_callback_user_data : nullptr)) {
        return false;
    }

    // process all time frames of the encoder output
    while (t < n_frames) {
        batch.n_tokens  = 1;
        batch.i_time[0] = t;
        batch.logits[0] = 1;

        // Use the current encoder frame (t) and the output of the prediction to
        // generate probabilities for the next token and duration. batch.i_time
        // is used in to select the correct frame from the encoder output.
        // The joint network outputs logits for all the tokens in the vocabulary
        // plus the blank token, and also n_duration logits for the duration
        // tokens which contain information about how many frames to skip/advance forward.
        if (!parakeet_joint(pctx, pstate, batch, n_threads,
                params ? params->abort_callback           : nullptr,
                params ? params->abort_callback_user_data : nullptr)) {
            return false;
        }

        const int64_t t_start_sample_us = ggml_time_us();

        // find the best token (greedy).
        // TODO: implement beam search?
        int best_token = 0;
        float max_logit = -1e10f;
        for (int i = 0; i < n_vocab_logits; ++i) {
            if (pstate.logits[i] > max_logit) {
                max_logit = pstate.logits[i];
                best_token = i;
            }
        }

        // find the max index of the duration logits, and look up that index
        // value in the tdt_durations array to get the actual duration value.
        int best_duration_idx = 0;
        float best_duration_logit = -1e10f;
        for (int i = 0; i < n_tdt_durations; ++i) {
            if (pstate.logits[n_vocab_logits + i] > best_duration_logit) {
                best_duration_logit = pstate.logits[n_vocab_logits + i];
                best_duration_idx = i;
            }
        }
        // look up that max duration index value in the tdt_durations array to
        // get the actual duration value.
        int duration = tdt_durations[best_duration_idx];

        if (best_token == blank_id) {
            if (duration == 0) {
                duration = 1;
            }
            // skip forward by duration time frames.
            t += duration;
            // reset symbols emitted counter
            tokens_emitted = 0;
            // continue without predicting.
            continue;
        }

        // Emit non-blank token at current frame t.
        pstate.decoded_tokens.push_back(best_token);
        pstate.t_sample_us += ggml_time_us() - t_start_sample_us;
        pstate.n_sample++;

        parakeet_token_data token_data = create_token_data(
            pctx, pstate, best_token, best_duration_idx, duration, t,
            max_logit, n_vocab_logits);

        pstate.decoded_token_data.push_back(token_data);

        // Call token callback if registered (for real-time streaming)
        if (params && params->new_token_callback) {
            params->new_token_callback(&pctx, &pstate, &token_data, params->new_token_callback_user_data);
        }

        last_token = best_token;

        // advance predictor for the non-blank token.
        batch.token[0] = last_token;
        if (!parakeet_predict(pctx, pstate, batch, n_threads,
                params ? params->abort_callback           : nullptr,
                params ? params->abort_callback_user_data : nullptr)) {
            return false;
        }

        // if duration greater than 0, continue looping over the encoder frames
        // and skip to the updated time frame (t).
        if (duration > 0) {
            t += duration;
            tokens_emitted = 0;
            continue;
        }

        // if duration is zero we stay on the current time frame.
        tokens_emitted++;
        if (tokens_emitted >= max_tokens_per_timestep) {
            t += 1; // forced blank/time advance behavior
            tokens_emitted = 0;
        }
    }

    return true;
}

//  500 -> 00:05.000
// 6000 -> 01:00.000
// naive Discrete Fourier Transform
// input is real-valued
// output is complex-valued
static void dft(const float* in, int N, float* out, const parakeet_mel_cache & cache) {
    const int sin_cos_step = cache.n_fft / N;

    for (int k = 0; k < N; k++) {
        float re = 0;
        float im = 0;

        for (int n = 0; n < N; n++) {
            int idx = (k * n * sin_cos_step) % cache.n_fft; // t = 2*M_PI*k*n/N
            re += in[n]*cache.cos_vals[idx]; // cos(t)
            im -= in[n]*cache.sin_vals[idx]; // sin(t)
        }

        out[k*2 + 0] = re;
        out[k*2 + 1] = im;
    }
}

// Cooley-Tukey FFT
// poor man's implementation - use something better
// input is real-valued
// output is complex-valued
static void fft(float* in, int N, float* out, const parakeet_mel_cache & cache) {
    if (N == 1) {
        out[0] = in[0];
        out[1] = 0;
        return;
    }

    const int half_N = N / 2;
    if (N - half_N*2 == 1) {
        dft(in, N, out, cache);
        return;
    }

    float* even = in + N;
    for (int i = 0; i < half_N; ++i) {
        even[i]= in[2*i];
    }
    float* even_fft = out + 2 * N;
    fft(even, half_N, even_fft, cache);

    float* odd = even;
    for (int i = 0; i < half_N; ++i) {
        odd[i] = in[2*i + 1];
    }
    float* odd_fft = even_fft + N;
    fft(odd, half_N, odd_fft, cache);

    const int sin_cos_step = cache.n_fft / N;
    for (int k = 0; k < half_N; k++) {
        int idx = k * sin_cos_step; // t = 2*M_PI*k/N
        float re = cache.cos_vals[idx]; // cos(t)
        float im = -cache.sin_vals[idx]; // sin(t)

        float re_odd = odd_fft[2*k + 0];
        float im_odd = odd_fft[2*k + 1];

        out[2*k + 0] = even_fft[2*k + 0] + re*re_odd - im*im_odd;
        out[2*k + 1] = even_fft[2*k + 1] + re*im_odd + im*re_odd;

        out[2*(k + half_N) + 0] = even_fft[2*k + 0] - re*re_odd + im*im_odd;
        out[2*(k + half_N) + 1] = even_fft[2*k + 1] - re*im_odd - im*re_odd;
    }
}

struct mel_worker_params {
    int ith;
    int window_size;
    int n_samples;
    int frame_size;
    int frame_step;
    int n_threads;
};

static void log_mel_spectrogram_worker_thread(
             mel_worker_params   params,
                   const float * window_func,
      const std::vector<float> & samples,
        const parakeet_filters & filters,
                  parakeet_mel & mel,
      const parakeet_mel_cache & cache) {
    std::vector<float> fft_in(params.frame_size * 2, 0.0);
    std::vector<float> fft_out(params.frame_size * 2 * 2 * 2);

    int n_fb = filters.n_fb;  // number of frequency bins
    int i = params.ith;

    // make sure n_fb == 1 + (frame_size / 2), bin_0 to bin_nyquist
    assert(n_fb == 1 + (params.frame_size / 2));

    const double eps = 5.960464477539063e-08;

    // calculate FFT only when fft_in are not all zero
    for (; i < std::min(params.n_samples / params.frame_step + 1, mel.n_len); i += params.n_threads) {
        const int offset = i * params.frame_step;

        const int window_pad_left = (params.frame_size - params.window_size) / 2;

        // Zero-pad left
        std::fill(fft_in.begin(), fft_in.begin() + window_pad_left, 0.0f);

        // Apply windowed samples in the center
        const int n_to_process = std::min({params.window_size, params.n_samples - offset});
        for (int j = 0; j < n_to_process; j++) {
            fft_in[window_pad_left + j] = window_func[j] * samples[offset + window_pad_left + j];
        }

        // Zero-pad right (and any samples we didn't have)
        std::fill(fft_in.begin() + window_pad_left + n_to_process, fft_in.begin() + params.frame_size, 0.0f);

        // FFT
        fft(fft_in.data(), params.frame_size, fft_out.data(), cache);

        // Calculate modulus^2 of complex numbers
        // Use pow(fft_out[2 * j + 0], 2) + pow(fft_out[2 * j + 1], 2) causes inference quality problem? Interesting.
        for (int j = 0; j < n_fb; j++) {
            fft_out[j] = (fft_out[2 * j + 0] * fft_out[2 * j + 0] + fft_out[2 * j + 1] * fft_out[2 * j + 1]);
        }

        // mel spectrogram
        for (int j = 0; j < mel.n_mel; j++) {
            double sum = 0.0;
            // unroll loop (suggested by GH user @lunixbochs)
            int k = 0;
            for (k = 0; k < n_fb - 3; k += 4) {
                sum +=
                        fft_out[k + 0] * filters.data[j * n_fb + k + 0] +
                        fft_out[k + 1] * filters.data[j * n_fb + k + 1] +
                        fft_out[k + 2] * filters.data[j * n_fb + k + 2] +
                        fft_out[k + 3] * filters.data[j * n_fb + k + 3];
            }
            // handle n_fb remainder
            for (; k < n_fb; k++) {
                sum += fft_out[k] * filters.data[j * n_fb + k];
            }

            mel.data[i * mel.n_mel + j] = std::log(sum + eps);
        }
    }

    // Otherwise fft_out are all zero - use log(eps) for consistency
    const double empty_sum = std::log(eps);
    for (; i < mel.n_len; i += params.n_threads) {
        for (int j = 0; j < mel.n_mel; j++) {
            mel.data[i * mel.n_mel + j] = empty_sum;
        }
    }
}

static bool log_mel_spectrogram(
                  parakeet_state & wstate,
                     const float * samples,
                       const int   n_samples,
                       const int   /*sample_rate*/,
                       const int   frame_size,
                       const int   frame_step,
                       const int   n_mel,
                       const int   n_threads,
          const parakeet_filters & filters,
                      const bool   debug,
                    parakeet_mel & mel,
        const parakeet_mel_cache & cache) {
    const int64_t t_start_us = ggml_time_us();

    const float * window_func = cache.window.empty() ? cache.hann_window.data() : cache.window.data();
    const int window_size = cache.window.empty() ? cache.n_fft : cache.window.size();

    std::vector<float> samples_preprocessed(samples, samples + n_samples);

    // Apply preemphasis filter (high-pass): x[i] = x[i] - 0.97 * x[i-1]
    {
        const float preemph = 0.97f;
        for (int i = n_samples - 1; i > 0; i--) {
            samples_preprocessed[i] = samples_preprocessed[i] - preemph * samples_preprocessed[i - 1];
        }
    }

    // Parakeet Pytorch implementation uses centered contant padding.
    const size_t pad = (size_t)(frame_size / 2);
    std::vector<float> samples_padded(n_samples + 2 * pad, 0.0f);
    std::copy(samples_preprocessed.begin(), samples_preprocessed.end(), samples_padded.begin() + pad);

    mel.n_mel = n_mel;
    mel.n_len = (samples_padded.size() - frame_size) / frame_step + 1;
    mel.n_len_org = mel.n_len;
    mel.data.resize(mel.n_mel * mel.n_len);

    // Worker Threads (STFT + Mel + Natural Log)
    {
        std::vector<std::thread> workers(n_threads - 1);
        const mel_worker_params mel_params { 0, window_size, (int)samples_padded.size(), frame_size, frame_step, n_threads };

        for (int iw = 0; iw < n_threads - 1; ++iw) {
            mel_worker_params params = mel_params;
            params.ith = iw + 1;
            workers[iw] = std::thread(log_mel_spectrogram_worker_thread,
                    params,
                    window_func,
                    std::cref(samples_padded),
                    std::cref(filters),
                    std::ref(mel),
                    std::cref(cache));
        }

        log_mel_spectrogram_worker_thread(
                mel_params,
                window_func,
                samples_padded,
                filters,
                mel,
                cache);

        for (int iw = 0; iw < n_threads - 1; ++iw) {
            workers[iw].join();
        }
    }

    {
        const double eps = 1e-5;
        int valid_frames = n_samples / frame_step;

        for (int j = 0; j < mel.n_mel; j++) {
            double sum = 0.0;
            double sq_diff_sum = 0.0;

            // Calculate Mean ONLY on valid audio frames
            for (int i = 0; i < valid_frames; i++) {
                sum += (double)mel.data[i * mel.n_mel + j];
            }
            double mean = sum / valid_frames;

            // Calculate Variance ONLY on valid audio frames
            for (int i = 0; i < valid_frames; i++) {
                double diff = (double)mel.data[i * mel.n_mel + j] - mean;
                sq_diff_sum += diff * diff;
            }

            double std_dev = std::sqrt(sq_diff_sum / (valid_frames - 1.0));
            double denominator = std_dev + eps;

            // Apply to ALL frames (including the padded ones)
            for (int i = 0; i < mel.n_len; i++) {
                mel.data[i * mel.n_mel + j] = (float)((mel.data[i * mel.n_mel + j] - mean) / denominator);
            }
        }
    }

    wstate.t_mel_us += ggml_time_us() - t_start_us;

    if (debug) {
        std::ofstream outFile("log_mel_spectrogram.json");
        outFile << "[";
        for (uint64_t i = 0; i < mel.data.size() - 1; i++) {
            outFile << mel.data[i] << ", ";
        }
        outFile << mel.data[mel.data.size() - 1] << "]";
        outFile.close();
    }

    return true;
}

static std::vector<parakeet_vocab::id> tokenize(const parakeet_vocab & vocab, const std::string & text) {
    std::vector<parakeet_vocab::id> tokens;
    const std::string normalized = sentencepiece_normalize(text);

    size_t i = 0;
    while (i < normalized.size()) {
        const size_t remaining = normalized.size() - i;
        const size_t max_len = std::min(vocab.max_token_length, remaining);

        bool found = false;
        for (size_t len = max_len; len > 0; --len) {
            const auto it = vocab.token_to_id.find(normalized.substr(i, len));
            if (it != vocab.token_to_id.end() && !is_sentencepiece_control(it->first)) {
                tokens.push_back(it->second);
                i += len;
                found = true;
                break;
            }
        }

        if (!found) {
            if (vocab.token_unk >= 0) {
                tokens.push_back(vocab.token_unk);
            }

            const unsigned char c = static_cast<unsigned char>(normalized[i]);
            i += utf8_codepoint_len(c);
        }
    }

    return tokens;
}


//
// interface implementation
//

struct parakeet_state * parakeet_init_state(parakeet_context * ctx) {
    parakeet_state * state = new parakeet_state;

    state->backends = parakeet_backend_init(ctx->params);
    if (state->backends.empty()) {
        PARAKEET_LOG_ERROR("%s: parakeet_backend_init() failed\n", __func__);
        parakeet_free_state(state);
        return nullptr;
    }

    const int batch_size = ctx->model.hparams.n_audio_ctx;

    state->logits.reserve(ctx->vocab.n_vocab * batch_size);

    state->batch = parakeet_batch_init(batch_size);

    {
        const int n_audio_state    = ctx->model.hparams.n_audio_state;
        const int subsampl_factor  = ctx->model.hparams.subsampling_factor;
        const int n_frames_max     = (batch_size + subsampl_factor - 1) / subsampl_factor;

        if (!parakeet_enc_state_init(*state, state->backends[0], n_audio_state, n_frames_max)) {
            PARAKEET_LOG_ERROR("%s: parakeet_enc_state_init() failed\n", __func__);
            parakeet_free_state(state);
            return nullptr;
        }

        const size_t mem_enc_ctx = state->enc_out_buf.size();
        const size_t mem_enc_out_buf = ggml_backend_buffer_get_size(state->enc_out_buffer);
        PARAKEET_LOG_INFO("%s: enc_out state: %7.2f MB (meta) + %7.2f MB (data)\n", __func__,
                mem_enc_ctx / 1024.0 / 1024.0, mem_enc_out_buf / 1024.0 / 1024.0);
    }

    // conv/encoder allocator
    bool ok = parakeet_sched_graph_init(state->sched_encode, state->backends,
            [&]() {
                return parakeet_build_graph_encode(*ctx, *state);
            });

    if (!ok) {
        PARAKEET_LOG_ERROR("%s: failed to init encode allocator\n", __func__);
        parakeet_free_state(state);
        return nullptr;
    }
    state->sched_encode_n_audio_ctx = state->n_audio_ctx > 0 ? state->n_audio_ctx : ctx->model.hparams.n_audio_ctx;

    if (!parakeet_lstm_state_init(*state, state->backends[0], ctx->model.hparams.n_pred_layers, ctx->model.hparams.n_pred_dim)) {
        PARAKEET_LOG_ERROR("%s: parakeet_lstm_states_init () failed\n", __func__);
        parakeet_free_state(state);
        return nullptr;
    }

    {
        const size_t mem_lstm_ctx = state->lstm_state.ctx_buf.size();
        const size_t mem_lstm_buf = ggml_backend_buffer_get_size(state->lstm_state.buffer);
        PARAKEET_LOG_INFO("%s: lstm state: %7.2f MB (meta) + %7.2f MB (data)\n", __func__,
                mem_lstm_ctx / 1024.0 / 1024.0, mem_lstm_buf / 1024.0 / 1024.0);
    }

    if (!parakeet_pred_state_init(*state, state->backends[0], ctx->model.hparams.n_pred_dim)) {
        PARAKEET_LOG_ERROR("%s: parakeet_pred_state_init() failed\n", __func__);
        parakeet_free_state(state);
        return nullptr;
    }

    {
        const size_t mem_pred_ctx = state->pred_out_buf.size();
        const size_t mem_pred_out_buf = ggml_backend_buffer_get_size(state->pred_out_buffer);
        PARAKEET_LOG_INFO("%s: pred state: %7.2f MB (meta) + %7.2f MB (data)\n", __func__,
                mem_pred_ctx / 1024.0 / 1024.0, mem_pred_out_buf / 1024.0 / 1024.0);
    }

    PARAKEET_LOG_INFO("%s: compute buffer (encode) = %7.2f MB\n", __func__, parakeet_sched_size(state->sched_encode) / 1e6);

    {
        bool ok = parakeet_sched_graph_init(state->sched_decode, state->backends,
                [&]() {
                    const auto & hparams = ctx->model.hparams;
                    const int n_tokens = hparams.n_audio_ctx; // Use audio ctx for Parakeet

                    parakeet_batch_prep_legacy(state->batch, nullptr, n_tokens, 0, 0);

                    return parakeet_build_graph_prediction(*ctx, *state, state->batch, true);
                });

        if (!ok) {
            PARAKEET_LOG_ERROR("%s: failed to init decoder allocator\n", __func__);
            parakeet_free_state(state);
            return nullptr;
        }

        PARAKEET_LOG_INFO("%s: compute buffer (decode) = %7.2f MB\n", __func__, parakeet_sched_size(state->sched_decode) / 1e6);
    }

    return state;
}

struct parakeet_context_params parakeet_context_default_params() {
    struct parakeet_context_params result = {
        /*.use_gpu              =*/ true,
        /*.gpu_device           =*/ 0,
    };
    return result;
}

struct parakeet_context * parakeet_init_from_file_with_params_no_state(const char * path_model, struct parakeet_context_params params) {
    PARAKEET_LOG_INFO("%s: loading model from '%s'\n", __func__, path_model);
#ifdef _MSC_VER
    // Convert UTF-8 path to wide string (UTF-16) for Windows, resolving character encoding issues.
    std::wstring_convert<std::codecvt_utf8<wchar_t>> converter;
    std::wstring path_model_wide = converter.from_bytes(path_model);
    auto fin = std::ifstream(path_model_wide, std::ios::binary);
#else
    auto fin = std::ifstream(path_model, std::ios::binary);
#endif
    if (!fin) {
        PARAKEET_LOG_ERROR("%s: failed to open '%s'\n", __func__, path_model);
        return nullptr;
    }

    parakeet_model_loader loader = {};

    loader.context = &fin;

    loader.read = [](void * ctx, void * output, size_t read_size) {
        std::ifstream * fin = (std::ifstream*)ctx;
        fin->read((char *)output, read_size);
        return read_size;
    };

    loader.eof = [](void * ctx) {
        std::ifstream * fin = (std::ifstream*)ctx;
        return fin->eof();
    };

    loader.close = [](void * ctx) {
        std::ifstream * fin = (std::ifstream*)ctx;
        fin->close();
    };

    auto ctx = parakeet_init_with_params_no_state(&loader, params);

    if (ctx) {
        ctx->path_model = path_model;
    }

    return ctx;
}

struct parakeet_context * parakeet_init_from_buffer_with_params_no_state(void * buffer, size_t buffer_size, struct parakeet_context_params params) {
    struct buf_context {
        uint8_t* buffer;
        size_t size;
        size_t current_offset;
    };

    buf_context ctx = { reinterpret_cast<uint8_t*>(buffer), buffer_size, 0 };

    PARAKEET_LOG_INFO("%s: loading model from buffer\n", __func__);

    parakeet_model_loader loader = {};

    loader.context = &ctx;

    loader.read = [](void * ctx, void * output, size_t read_size) {
        buf_context * buf = reinterpret_cast<buf_context *>(ctx);

        size_t size_to_copy = buf->current_offset + read_size < buf->size ? read_size : buf->size - buf->current_offset;

        memcpy(output, buf->buffer + buf->current_offset, size_to_copy);
        buf->current_offset += size_to_copy;

        return size_to_copy;
    };

    loader.eof = [](void * ctx) {
        buf_context * buf = reinterpret_cast<buf_context *>(ctx);

        return buf->current_offset >= buf->size;
    };

    loader.close = [](void * /*ctx*/) { };

    return parakeet_init_with_params_no_state(&loader, params);
}

struct parakeet_context * parakeet_init_with_params_no_state(struct parakeet_model_loader * loader, struct parakeet_context_params params) {
    ggml_time_init();

    PARAKEET_LOG_INFO("%s: use gpu    = %d\n", __func__, params.use_gpu);
    PARAKEET_LOG_INFO("%s: gpu_device = %d\n", __func__, params.gpu_device);
    PARAKEET_LOG_INFO("%s: devices    = %zu\n", __func__, ggml_backend_dev_count());
    PARAKEET_LOG_INFO("%s: backends   = %zu\n", __func__, ggml_backend_reg_count());

    parakeet_context * ctx = new parakeet_context;
    ctx->params = params;

    bool model_loaded = false;
    try {
        model_loaded = parakeet_model_load(loader, *ctx);
    } catch (const std::exception & e) {
        PARAKEET_LOG_ERROR("%s: exception during model load: %s\n", __func__, e.what());
    } catch (...) {
        PARAKEET_LOG_ERROR("%s: unknown exception during model load\n", __func__);
    }

    if (!model_loaded) {
        loader->close(loader->context);
        PARAKEET_LOG_ERROR("%s: failed to load model\n", __func__);
        delete ctx;
        return nullptr;
    }

    loader->close(loader->context);

    // Initialize mel cache with model's FFT size
    ctx->mel_cache.init(ctx->model.hparams.n_fft);
    PARAKEET_LOG_INFO("%s: initialized mel cache with n_fft = %d\n", __func__, ctx->model.hparams.n_fft);

    return ctx;
}

struct parakeet_context * parakeet_init_from_file_with_params(const char * path_model, struct parakeet_context_params params) {
    parakeet_context * ctx = parakeet_init_from_file_with_params_no_state(path_model, params);
    if (!ctx) {
        return nullptr;
    }

    ctx->state = parakeet_init_state(ctx);
    if (!ctx->state) {
        parakeet_free(ctx);
        return nullptr;
    }

    return ctx;
}

struct parakeet_context * parakeet_init_from_buffer_with_params(void * buffer, size_t buffer_size, struct parakeet_context_params params) {
    parakeet_context * ctx = parakeet_init_from_buffer_with_params_no_state(buffer, buffer_size, params);
    if (!ctx) {
        return nullptr;
    }

    ctx->state = parakeet_init_state(ctx);
    if (!ctx->state) {
        parakeet_free(ctx);
        return nullptr;
    }

    return ctx;
}

struct parakeet_context * parakeet_init_with_params(struct parakeet_model_loader * loader, struct parakeet_context_params params) {
    parakeet_context * ctx = parakeet_init_with_params_no_state(loader, params);
    if (!ctx) {
        return nullptr;
    }

    ctx->state = parakeet_init_state(ctx);
    if (!ctx->state) {
        parakeet_free(ctx);
        return nullptr;
    }

    return ctx;
}

void parakeet_free_state(struct parakeet_state * state) {
    if (state) {
        ggml_backend_buffer_free(state->lstm_state.buffer);
        ggml_backend_buffer_free(state->pred_out_buffer);
        ggml_backend_buffer_free(state->enc_out_buffer);

        parakeet_batch_free(state->batch);

        parakeet_sched_free(state->sched_encode);
        parakeet_sched_free(state->sched_decode);

        for (auto & backend : state->backends) {
            ggml_backend_free(backend);
        }

        delete state;
    }
}

void parakeet_free(struct parakeet_context * ctx) {
    if (ctx) {
        for (ggml_context * context : ctx->model.ctxs) {
            ggml_free(context);
        }

        for (ggml_backend_buffer_t buf : ctx->model.buffers) {
            ggml_backend_buffer_free(buf);
        }

        parakeet_free_state(ctx->state);

        delete ctx;
    }
}

void parakeet_free_context_params(struct parakeet_context_params * params) {
    if (params) {
        delete params;
    }
}

void parakeet_free_params(struct parakeet_full_params * params) {
    if (params) {
        delete params;
    }
}

int parakeet_pcm_to_mel_with_state(struct parakeet_context * ctx, struct parakeet_state * state, const float * samples, int n_samples, int n_threads) {
    if (!log_mel_spectrogram(*state,
                samples,
                n_samples,
                PARAKEET_SAMPLE_RATE,
                ctx->model.hparams.n_fft,
                PARAKEET_HOP_LENGTH,
                ctx->model.filters.n_mel,
                n_threads,
                ctx->model.filters,
                false,                        // debug
                state->mel,
                ctx->mel_cache)) {
        PARAKEET_LOG_ERROR("%s: failed to compute mel spectrogram\n", __func__);
        return -1;
    }

    return 0;
}

int parakeet_pcm_to_mel(struct parakeet_context * ctx, const float * samples, int n_samples, int n_threads) {
    return parakeet_pcm_to_mel_with_state(ctx, ctx->state, samples, n_samples, n_threads);
}

int parakeet_set_mel_with_state(
        struct parakeet_context * ctx,
          struct parakeet_state * state,
                   const float * data,
                           int   n_len,
                           int   n_mel) {
    if (n_mel != ctx->model.filters.n_mel) {
        PARAKEET_LOG_ERROR("%s: invalid number of mel bands: %d (expected %d)\n", __func__, n_mel, ctx->model.filters.n_mel);
        return -1;
    }

    state->mel.n_len     = n_len;
    state->mel.n_len_org = n_len;
    state->mel.n_mel     = n_mel;

    state->mel.data.resize(n_len*n_mel);
    memcpy(state->mel.data.data(), data, n_len*n_mel*sizeof(float));

    return 0;
}

int parakeet_set_mel(
        struct parakeet_context * ctx,
        const float * data,
        int n_len,
        int n_mel) {
    return parakeet_set_mel_with_state(ctx, ctx->state, data, n_len, n_mel);
}

int parakeet_encode_with_state(struct parakeet_context * ctx, struct parakeet_state * state, int offset, int n_threads) {
    if (!parakeet_encode_internal(*ctx, *state, offset, n_threads, nullptr, nullptr)) {
        PARAKEET_LOG_ERROR("%s: failed to eval\n", __func__);
        return -1;
    }

    return 0;
}

int parakeet_encode(struct parakeet_context * ctx, int offset, int n_threads) {
    if (!parakeet_encode_internal(*ctx, *ctx->state, offset, n_threads, nullptr, nullptr)) {
        PARAKEET_LOG_ERROR("%s: failed to eval\n", __func__);
        return -1;
    }

    return 0;
}

int parakeet_tokenize(struct parakeet_context * ctx, const char * text, parakeet_token * tokens, int n_max_tokens) {
    const auto res = tokenize(ctx->vocab, text);

    if (n_max_tokens < (int) res.size()) {
        PARAKEET_LOG_ERROR("%s: too many resulting tokens: %d (max %d)\n", __func__, (int) res.size(), n_max_tokens);
        return -(int) res.size();
    }

    for (int i = 0; i < (int) res.size(); i++) {
        tokens[i] = res[i];
    }

    return res.size();
}

int parakeet_token_count(struct parakeet_context * ctx, const char * text) {
    return -parakeet_tokenize(ctx, text, NULL, 0);
}

int parakeet_model_n_vocab(struct parakeet_context * ctx) {
    return ctx->model.hparams.n_vocab;
}

int parakeet_model_n_audio_ctx(struct parakeet_context * ctx) {
    return ctx->model.hparams.n_audio_ctx;
}

int parakeet_model_n_audio_state(struct parakeet_context * ctx) {
    return ctx->model.hparams.n_audio_state;
}

int parakeet_model_n_audio_head(struct parakeet_context * ctx) {
    return ctx->model.hparams.n_audio_head;
}

int parakeet_model_n_audio_layer(struct parakeet_context * ctx) {
    return ctx->model.hparams.n_audio_layer;
}

int parakeet_model_n_mels(struct parakeet_context * ctx) {
    return ctx->model.hparams.n_mels;
}

int parakeet_model_ftype(struct parakeet_context * ctx) {
    return ctx->model.hparams.ftype;
}

int parakeet_n_len_from_state(struct parakeet_state * state) {
    return state->mel.n_len_org;
}

int parakeet_n_len(struct parakeet_context * ctx) {
    return ctx->state->mel.n_len_org;
}

int parakeet_n_vocab(struct parakeet_context * ctx) {
    return ctx->vocab.n_vocab;
}

int parakeet_n_audio_ctx(struct parakeet_context * ctx) {
    return ctx->model.hparams.n_audio_ctx;
}

float * parakeet_get_logits(struct parakeet_context * ctx) {
    return ctx->state->logits.data();
}

float * parakeet_get_logits_from_state(struct parakeet_state * state) {
    return state->logits.data();
}

const char * parakeet_token_to_str(struct parakeet_context * ctx, parakeet_token token) {
    return ctx->vocab.id_to_token.at(token).c_str();
}

int parakeet_token_to_text(const char * token_str, bool is_first, char * output, int max_len) {
    std::string text = sentencepiece_piece_to_text(token_str, is_first);

    if (output == nullptr) {
        return text.size();
    }

    int bytes_to_copy = std::min((int)text.size(), max_len - 1);
    if (bytes_to_copy > 0) {
        memcpy(output, text.c_str(), bytes_to_copy);
        output[bytes_to_copy] = '\0';
    } else if (max_len > 0) {
        output[0] = '\0';
    }

    return text.size();
}

parakeet_token parakeet_token_bos(struct parakeet_context * ctx) {
    return ctx->vocab.token_bos;
}

parakeet_token parakeet_token_unk(struct parakeet_context * ctx) {
    return ctx->vocab.token_unk;
}

parakeet_token parakeet_token_blank(struct parakeet_context * ctx) {
    return ctx->vocab.token_blank;
}

struct parakeet_timings * parakeet_get_timings(struct parakeet_context * ctx) {
    if (ctx->state == nullptr) {
        return nullptr;
    }
    parakeet_timings * timings = new parakeet_timings;
    timings->sample_ms = 1e-3f * ctx->state->t_sample_us / std::max(1, ctx->state->n_sample);
    timings->encode_ms = 1e-3f * ctx->state->t_encode_us / std::max(1, ctx->state->n_encode);
    timings->decode_ms = 1e-3f * ctx->state->t_decode_us / std::max(1, ctx->state->n_decode);
    return timings;
}

void parakeet_print_timings(struct parakeet_context * ctx) {
    const int64_t t_end_us = ggml_time_us();

    PARAKEET_LOG_INFO("\n");
    PARAKEET_LOG_INFO("%s:     load time = %8.2f ms\n", __func__, ctx->t_load_us / 1000.0f);
    if (ctx->state != nullptr) {

        const int32_t n_sample  = std::max(1, ctx->state->n_sample);
        const int32_t n_encode  = std::max(1, ctx->state->n_encode);
        const int32_t n_decode  = std::max(1, ctx->state->n_decode);
        const int32_t n_predict = std::max(1, ctx->state->n_predict);

        PARAKEET_LOG_INFO("%s:     fallbacks = %3d p / %3d h\n", __func__, ctx->state->n_fail_p, ctx->state->n_fail_h);
        PARAKEET_LOG_INFO("%s:      mel time = %8.2f ms\n", __func__, ctx->state->t_mel_us / 1000.0f);
        PARAKEET_LOG_INFO("%s:   sample time = %8.2f ms / %5d runs ( %8.2f ms per run)\n", __func__, 1e-3f * ctx->state->t_sample_us, n_sample, 1e-3f * ctx->state->t_sample_us / n_sample);
        PARAKEET_LOG_INFO("%s:   encode time = %8.2f ms / %5d runs ( %8.2f ms per run)\n", __func__, 1e-3f * ctx->state->t_encode_us, n_encode, 1e-3f * ctx->state->t_encode_us / n_encode);
        PARAKEET_LOG_INFO("%s:   decode time = %8.2f ms / %5d runs ( %8.2f ms per run)\n", __func__, 1e-3f * ctx->state->t_decode_us, n_decode, 1e-3f * ctx->state->t_decode_us / n_decode);
        PARAKEET_LOG_INFO("%s:  predict time = %8.2f ms / %5d runs ( %8.2f ms per run)\n", __func__, 1e-3f * ctx->state->t_predict_us, n_predict, 1e-3f * ctx->state->t_predict_us / n_predict);
        PARAKEET_LOG_INFO("%s:    - build     = %8.2f ms / %5d runs ( %8.2f ms per run)\n", __func__, 1e-3f * ctx->state->t_predict_build_us, n_predict, 1e-3f * ctx->state->t_predict_build_us / n_predict);
        PARAKEET_LOG_INFO("%s:    - alloc     = %8.2f ms / %5d runs ( %8.2f ms per run)\n", __func__, 1e-3f * ctx->state->t_predict_alloc_us, n_predict, 1e-3f * ctx->state->t_predict_alloc_us / n_predict);
        PARAKEET_LOG_INFO("%s:    - compute   = %8.2f ms / %5d runs ( %8.2f ms per run)\n", __func__, 1e-3f * ctx->state->t_predict_compute_us, n_predict, 1e-3f * ctx->state->t_predict_compute_us / n_predict);

    }
    PARAKEET_LOG_INFO("%s:    total time = %8.2f ms\n", __func__, (t_end_us - ctx->t_start_us)/1000.0f);
}

void parakeet_reset_timings(struct parakeet_context * ctx) {
    ctx->t_start_us = ggml_time_us();
    if (ctx->state != nullptr) {
        ctx->state->t_mel_us = 0;
        ctx->state->t_sample_us = 0;
        ctx->state->t_encode_us = 0;
        ctx->state->t_decode_us = 0;
        ctx->state->t_predict_us = 0;
        ctx->state->t_predict_build_us = 0;
        ctx->state->t_predict_alloc_us = 0;
        ctx->state->t_predict_compute_us = 0;

        ctx->state->n_sample = 0;
        ctx->state->n_encode = 0;
        ctx->state->n_decode = 0;
        ctx->state->n_predict = 0;
    }
}

const char * parakeet_print_system_info(void) {
    static std::string s;

    s  = "";
    s += "PARAKEET : ";

    for (size_t i = 0; i < ggml_backend_reg_count(); i++) {
        auto * reg = ggml_backend_reg_get(i);
        auto * get_features_fn = (ggml_backend_get_features_t) ggml_backend_reg_get_proc_address(reg, "ggml_backend_get_features");
        if (get_features_fn) {
            ggml_backend_feature * features = get_features_fn(reg);
            s += ggml_backend_reg_name(reg);
            s += " : ";
            for (; features->name; features++) {
                s += features->name;
                s += " = ";
                s += features->value;
                s += " | ";
            }
        }
    }
    return s.c_str();
}

struct parakeet_context_params * parakeet_context_default_params_by_ref(void) {
    struct parakeet_context_params params = parakeet_context_default_params();

    struct parakeet_context_params* result = new parakeet_context_params();
    *result = params;
    return result;
}

struct parakeet_full_params * parakeet_full_default_params_by_ref(enum parakeet_sampling_strategy strategy) {
    struct parakeet_full_params params = parakeet_full_default_params(strategy);

    struct parakeet_full_params* result = new parakeet_full_params();
    *result = params;
    return result;
}

struct parakeet_full_params parakeet_full_default_params(enum parakeet_sampling_strategy strategy) {
    struct parakeet_full_params result = {
        /*.strategy                         =*/ strategy,
        /*.n_threads                        =*/ std::min(4, (int32_t) std::thread::hardware_concurrency()),
        /*.offset_ms                        =*/ 0,
        /*.duration_ms                      =*/ 0,
        /*.no_context                       =*/ true,
        /*.audio_ctx                        =*/ 0,
        /*.new_token_callback               =*/ nullptr,
        /*.new_token_callback_user_data     =*/ nullptr,
        /*.new_segment_callback             =*/ nullptr,
        /*.new_segment_callback_user_data   =*/ nullptr,
        /*.progress_callback                =*/ nullptr,
        /*.progress_callback_user_data      =*/ nullptr,
        /*.encoder_begin_callback           =*/ nullptr,
        /*.encoder_begin_callback_user_data =*/ nullptr,
        /*.abort_callback                   =*/ nullptr,
        /*.abort_callback_user_data         =*/ nullptr,
    };

    return result;
}

static void parakeet_reset_state(struct parakeet_state * state) {
    state->decoded_tokens.clear();
    state->decoded_token_data.clear();

    if (state->lstm_state.buffer) {
        ggml_backend_buffer_clear(state->lstm_state.buffer, 0);
    }

}

// Encode and decode the mel spectrogram already in state, without recomputing it.
static int parakeet_chunk_with_state(
      struct parakeet_context   * ctx,
        struct parakeet_state   * state,
    struct parakeet_full_params   params) {
    return parakeet_chunk(ctx, state, params, nullptr, 0);
}

int parakeet_full_with_state(
        struct parakeet_context * ctx,
          struct parakeet_state * state,
    struct parakeet_full_params   params,
                    const float * samples,
                           int    n_samples) {
    state->result_all.clear();

    if (params.no_context) {
        parakeet_reset_state(state);
    }

    if (n_samples > 0) {
        if (parakeet_pcm_to_mel_with_state(ctx, state, samples, n_samples, params.n_threads) != 0) {
            PARAKEET_LOG_ERROR("%s: failed to compute log mel spectrogram\n", __func__);
            return -2;
        }
    }

    const int n_mel_total = state->mel.n_len;
    const int n_audio_ctx = ctx->model.hparams.n_audio_ctx;

    if (n_mel_total <= n_audio_ctx) {
        if (params.progress_callback) {
            params.progress_callback(ctx, state, 0, params.progress_callback_user_data);
        }
        return parakeet_chunk_with_state(ctx, state, params);
    }

    PARAKEET_LOG_DEBUG("%s: audio too long (%d mel > n_audio_ctx=%d), using dynamic encoder graph\n",
                       __func__, n_mel_total, n_audio_ctx);

    if (params.encoder_begin_callback) {
        if (!params.encoder_begin_callback(ctx, state, params.encoder_begin_callback_user_data)) {
            PARAKEET_LOG_ERROR("%s: encoder_begin_callback returned false\n", __func__);
            return -6;
        }
    }

    if (params.progress_callback) {
        params.progress_callback(ctx, state, 0, params.progress_callback_user_data);
    }

    if (!parakeet_ensure_encode_sched(*ctx, *state, n_mel_total)) {
        PARAKEET_LOG_ERROR("%s: failed to allocate dynamic encoder graph for %d mel frames\n",
                __func__, n_mel_total);
        return -6;
    }

    state->n_audio_ctx = n_mel_total;

    if (!parakeet_encode_internal(*ctx, *state, 0, params.n_threads,
                                  params.abort_callback, params.abort_callback_user_data)) {
        PARAKEET_LOG_ERROR("%s: failed to encode\n", __func__);
        return -6;
    }

    if (params.progress_callback) {
        params.progress_callback(ctx, state, 100, params.progress_callback_user_data);
    }

    const size_t tokens_before = state->decoded_tokens.size();

    if (!parakeet_decode(*ctx, *state, state->batch, params.n_threads, &params)) {
        PARAKEET_LOG_ERROR("%s: failed to decode\n", __func__);
        return -7;
    }

    const size_t tokens_after    = state->decoded_tokens.size();
    const size_t new_token_count = tokens_after - tokens_before;

    if (new_token_count > 0) {
        std::string text;
        std::vector<parakeet_token_data> result_tokens;

        for (size_t i = tokens_before; i < tokens_after; i++) {
            const auto token_id  = state->decoded_tokens[i];
            const char * tok_str = parakeet_token_to_str(ctx, token_id);
            if (tok_str) {
                const bool is_first = (tokens_before == 0) && text.empty();
                text += sentencepiece_piece_to_text(tok_str, is_first);
            }
            result_tokens.push_back(state->decoded_token_data[i]);
        }

        refine_timestamps_tdt(ctx->vocab, result_tokens);

        if (!text.empty()) {
            parakeet_segment seg;
            seg.t0     = 0;
            seg.t1     = state->n_frames;
            seg.text   = text;
            seg.tokens = result_tokens;
            state->result_all.push_back(std::move(seg));

            if (params.new_segment_callback) {
                params.new_segment_callback(ctx, state, 1, params.new_segment_callback_user_data);
            }
        }
    }

    return 0;
}

int parakeet_full(
        struct parakeet_context * ctx,
    struct parakeet_full_params   params,
                    const float * samples,
                            int   n_samples) {
    return parakeet_full_with_state(ctx, ctx->state, params, samples, n_samples);
}

int parakeet_chunk(
        struct parakeet_context * ctx,
          struct parakeet_state * state,
    struct parakeet_full_params   params,
                    const float * samples,
                            int   n_samples) {

    if (params.no_context) {
        parakeet_reset_state(state);
    }

    if (n_samples > 0) {
        if (parakeet_pcm_to_mel_with_state(ctx, state, samples, n_samples, params.n_threads) != 0) {
            PARAKEET_LOG_ERROR("%s: failed to compute log mel spectrogram\n", __func__);
            return -2;
        }
    }

    if (params.audio_ctx == 0) {
        const int total_len = parakeet_n_len_from_state(state);
        const int model_max_ctx = parakeet_n_audio_ctx(ctx);
        params.audio_ctx = std::min(total_len, model_max_ctx);
        PARAKEET_LOG_DEBUG("Processing audio: total_frames=%d, chunk_size=%d\n", total_len, params.audio_ctx);
    }
    state->n_audio_ctx = params.audio_ctx;

    const int n_frames = parakeet_n_len_from_state(state);

    if (!parakeet_ensure_encode_sched(*ctx, *state, state->n_audio_ctx)) {
        PARAKEET_LOG_ERROR("%s: failed to allocate encoder graph for %d mel frames\n",
                __func__, state->n_audio_ctx);
        return -6;
    }

    if (params.encoder_begin_callback) {
        if (!params.encoder_begin_callback(ctx, state, params.encoder_begin_callback_user_data)) {
            PARAKEET_LOG_ERROR("%s: encoder_begin_callback returned false - aborting\n", __func__);
            return -6;
        }
    }
    if (!parakeet_encode_internal(*ctx, *state, 0, params.n_threads, params.abort_callback, params.abort_callback_user_data)) {
        PARAKEET_LOG_ERROR("%s: failed to encode\n", __func__);
        return -6;
    }

    const size_t tokens_before = state->decoded_tokens.size();

    if (!parakeet_decode(*ctx, *state, state->batch, params.n_threads, &params)) {
        PARAKEET_LOG_ERROR("%s: failed to decode\n", __func__);
        return -7;
    }

    const size_t tokens_after = state->decoded_tokens.size();
    const size_t new_token_count = tokens_after - tokens_before;

    if (new_token_count > 0) {
        std::string text;
        std::vector<parakeet_token_data> result_tokens;

        for (size_t i = tokens_before; i < tokens_after; i++) {
            const auto token_id = state->decoded_tokens[i];
            const char * token_str = parakeet_token_to_str(ctx, token_id);
            if (token_str) {
                const bool is_first_piece = (tokens_before == 0) && text.empty();
                text += sentencepiece_piece_to_text(token_str, is_first_piece);
            }

            // Use the stored token data from parakeet_decode
            result_tokens.push_back(state->decoded_token_data[i]);
        }

        refine_timestamps_tdt(ctx->vocab, result_tokens);

        if (!text.empty()) {
            parakeet_segment segment;
            segment.t0 = 0; // Caller tracks timing
            segment.t1 = n_frames;
            segment.text = text;
            segment.tokens = result_tokens;

            state->result_all.push_back(std::move(segment));

            if (params.new_segment_callback) {
                params.new_segment_callback(ctx, state, 1, params.new_segment_callback_user_data);
            }
        }
    }

    return 0;
}

int parakeet_full_n_segments_from_state(struct parakeet_state * state) {
    return state->result_all.size();
}

int parakeet_full_n_segments(struct parakeet_context * ctx) {
    return ctx->state->result_all.size();
}

int64_t parakeet_full_get_segment_t0_from_state(struct parakeet_state * state, int i_segment) {
    return state->result_all[i_segment].t0;
}

int64_t parakeet_full_get_segment_t1_from_state(struct parakeet_state * state, int i_segment) {
    return state->result_all[i_segment].t1;
}

int64_t parakeet_full_get_segment_t0(struct parakeet_context * ctx, int i_segment) {
    return parakeet_full_get_segment_t0_from_state(ctx->state, i_segment);
}

int64_t parakeet_full_get_segment_t1(struct parakeet_context * ctx, int i_segment) {
    return parakeet_full_get_segment_t1_from_state(ctx->state, i_segment);
}

const char * parakeet_full_get_segment_text_from_state(struct parakeet_state * state, int i_segment) {
    return state->result_all[i_segment].text.c_str();
}

const char * parakeet_full_get_segment_text(struct parakeet_context * ctx, int i_segment) {
    return ctx->state->result_all[i_segment].text.c_str();
}

int parakeet_full_n_tokens_from_state(struct parakeet_state * state, int i_segment) {
    return state->result_all[i_segment].tokens.size();
}

int parakeet_full_n_tokens(struct parakeet_context * ctx, int i_segment) {
    return ctx->state->result_all[i_segment].tokens.size();
}

const char * parakeet_full_get_token_text_from_state(struct parakeet_context * ctx, struct parakeet_state * state, int i_segment, int i_token) {
    return ctx->vocab.id_to_token[state->result_all[i_segment].tokens[i_token].id].c_str();
}

const char* parakeet_full_get_token_text(struct parakeet_context * ctx, int i_segment, int i_token) {
    return ctx->vocab.id_to_token[ctx->state->result_all[i_segment].tokens[i_token].id].c_str();
}

parakeet_token parakeet_full_get_token_id_from_state(struct parakeet_state * state, int i_segment, int i_token) {
    return state->result_all[i_segment].tokens[i_token].id;
}

parakeet_token parakeet_full_get_token_id(struct parakeet_context * ctx, int i_segment, int i_token) {
    return ctx->state->result_all[i_segment].tokens[i_token].id;
}

struct parakeet_token_data parakeet_full_get_token_data_from_state(struct parakeet_state * state, int i_segment, int i_token) {
    return state->result_all[i_segment].tokens[i_token];
}

struct parakeet_token_data parakeet_full_get_token_data(struct parakeet_context * ctx, int i_segment, int i_token) {
    return ctx->state->result_all[i_segment].tokens[i_token];
}

float parakeet_full_get_token_p_from_state(struct parakeet_state * state, int i_segment, int i_token) {
    return state->result_all[i_segment].tokens[i_token].p;
}

float parakeet_full_get_token_p(struct parakeet_context * ctx, int i_segment, int i_token) {
    return ctx->state->result_all[i_segment].tokens[i_token].p;
}

void parakeet_log_set(ggml_log_callback log_callback, void * user_data) {
    g_state.log_callback = log_callback ? log_callback : parakeet_log_callback_default;
    g_state.log_callback_user_data = user_data;
    ggml_log_set(g_state.log_callback, g_state.log_callback_user_data);
}

const char * parakeet_version(void) {
    return PARAKEET_VERSION;
}

GGML_ATTRIBUTE_FORMAT(2, 3)
static void parakeet_log_internal(ggml_log_level level, const char * format, ...) {
    va_list args;
    va_start(args, format);
    char buffer[1024];
    int len = vsnprintf(buffer, 1024, format, args);
    if (len < 1024) {
        g_state.log_callback(level, buffer, g_state.log_callback_user_data);
    } else {
        char* buffer2 = new char[len+1];
        vsnprintf(buffer2, len+1, format, args);
        buffer2[len] = 0;
        g_state.log_callback(level, buffer2, g_state.log_callback_user_data);
        delete[] buffer2;
    }
    va_end(args);
}

static void parakeet_log_callback_default(ggml_log_level level, const char * text, void * user_data) {
    (void) level;
    (void) user_data;
#ifndef PARAKEET_DEBUG
    if (level == GGML_LOG_LEVEL_DEBUG) {
        return;
    }
#endif
    fputs(text, stderr);
    fflush(stderr);
}