parakeet : fix TDT decode by outputting raw logits from the joint graph (#4017)

This commit is contained in:
Jhen-Jie Hong
2026-08-29 05:28:05 +02:00
committed by GitHub
parent 978113305b
commit c4ac0012a8
2 changed files with 15 additions and 13 deletions
+2 -1
View File
@@ -176,7 +176,8 @@ extern "C" {
// Token logits obtained from the last call to parakeet_full/parakeet_chunk // Token logits obtained from the last call to parakeet_full/parakeet_chunk
// The logits for the last token are stored in the last row // The logits for the last token are stored in the last row
// Rows: n_tokens // 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 (struct parakeet_context * ctx);
PARAKEET_API float * parakeet_get_logits_from_state(struct parakeet_state * state); PARAKEET_API float * parakeet_get_logits_from_state(struct parakeet_state * state);
+13 -12
View File
@@ -2302,12 +2302,7 @@ static struct ggml_cgraph * parakeet_build_graph_joint(
ggml_set_output(logits); ggml_set_output(logits);
ggml_set_name(logits, "logits"); ggml_set_name(logits, "logits");
struct ggml_tensor * probs = ggml_soft_max(ctx0, logits); ggml_build_forward_expand(gf, 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); ggml_free(ctx0);
@@ -2473,19 +2468,25 @@ static parakeet_token_data create_token_data(
float token_logit, float token_logit,
int n_vocab_logits) { 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; float token_sum = 0.0f;
for (int i = 0; i < n_vocab_logits; ++i) { 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; parakeet_token_data token_data;
token_data.id = token_id; token_data.id = token_id;
token_data.duration_idx = duration_idx; token_data.duration_idx = duration_idx;
token_data.duration_value = duration_value; token_data.duration_value = duration_value;
token_data.frame_index = frame_index; token_data.frame_index = frame_index;
token_data.p = token_p; token_data.p = expf(token_logit - log_z);
token_data.plog = token_logit; token_data.plog = token_logit - log_z;
token_data.t0 = frame_index * pctx.model.hparams.subsampling_factor; token_data.t0 = frame_index * pctx.model.hparams.subsampling_factor;
token_data.t1 = (frame_index + duration_value) * 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); 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 // 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. // value in the tdt_durations array to get the actual duration value.
int best_duration_idx = 0; int best_duration_idx = 0;
float best_duration_logit = -1e10f; float best_duration_logit = pstate.logits[n_vocab_logits];
for (int i = 0; i < n_tdt_durations; ++i) { for (int i = 1; i < n_tdt_durations; ++i) {
if (pstate.logits[n_vocab_logits + i] > best_duration_logit) { if (pstate.logits[n_vocab_logits + i] > best_duration_logit) {
best_duration_logit = pstate.logits[n_vocab_logits + i]; best_duration_logit = pstate.logits[n_vocab_logits + i];
best_duration_idx = i; best_duration_idx = i;