parakeet : fix TDT decode by outputting raw logits from the joint graph

Author: jhen0409

include/parakeet.h

@@ -176,7 +176,8 @@ extern "C" {
     // Token logits obtained from the last call to parakeet_full/parakeet_chunk
     // The logits for the last token are stored in the last row
     // Rows: n_tokens
-    // Cols: n_vocab
+    // Cols: n_vocab + 1 token logits (the blank token is at index n_vocab),
+    //       followed by n_tdt_durations duration logits
     PARAKEET_API float * parakeet_get_logits           (struct parakeet_context * ctx);
     PARAKEET_API float * parakeet_get_logits_from_state(struct parakeet_state * state);
 

src/parakeet.cpp

@@ -2302,12 +2302,7 @@ static struct ggml_cgraph * parakeet_build_graph_joint(
     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_build_forward_expand(gf, logits);
 
     ggml_free(ctx0);
 
@@ -2473,19 +2468,25 @@ static parakeet_token_data create_token_data(
                         float   token_logit,
                           int   n_vocab_logits) {
 
+    float max_logit = token_logit;
+    for (int i = 0; i < n_vocab_logits; ++i) {
+        max_logit = std::max(max_logit, pstate.logits[i]);
+    }
+
     float token_sum = 0.0f;
     for (int i = 0; i < n_vocab_logits; ++i) {
-        token_sum += expf(pstate.logits[i]);
+        token_sum += expf(pstate.logits[i] - max_logit);
     }
-    float token_p = expf(token_logit) / token_sum;
+
+    const float log_z = max_logit + logf(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.p = expf(token_logit - log_z);
+    token_data.plog = token_logit - log_z;
     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);
@@ -2566,8 +2567,8 @@ static bool parakeet_decode(
         // 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) {
+        float best_duration_logit = pstate.logits[n_vocab_logits];
+        for (int i = 1; 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;