Source code
Revision control
Copy as Markdown
Other Tools
#include "search.hpp"
#include <cmath>
namespace pk {
namespace {
// NeMo's rescaled `max_prob` confidence (confidence method 'max_prob', alpha==1.0):
// conf = (N * p_max - 1) / (N - 1)
// where p_max is the softmax probability of the argmax class (== exp of the max
// LOG-prob over the slice) and N is the number of classes in that slice
// (num_tokens = vocab_size + 1, blank included). See
// asr_confidence_utils.get_confidence_measure_bank()["max_prob"] (the t==1.0
// branch) and ConfidenceMethodMixin._init_confidence_method (num_tokens = blank+1).
float max_prob_conf(float max_logprob, int num_classes) {
const float p_max = std::exp(max_logprob);
const float N = (float)num_classes;
return (N * p_max - 1.0f) / (N - 1.0f);
}
} // namespace
std::vector<int32_t> ctc_greedy(const std::vector<float>& logits,
int T, int vocab_plus_1, int blank_id,
std::vector<TokenInfo>* tokens) {
std::vector<int32_t> out;
if (tokens) tokens->clear();
if (T <= 0 || vocab_plus_1 <= 0) return out;
out.reserve(T);
// Mirror NeMo's AbstractCTCDecoding.decode_hypothesis fold_consecutive loop:
// previous = blank_id
// emit p iff (p != previous or previous == blank_id) and p != blank_id
// previous = p (updated every frame, including blank frames)
//
// Per-frame argmax + per-frame max_prob confidence (CTC logits are
// log-softmax, so p_max = exp(max log-prob)). For TokenInfo we additionally
// reproduce NeMo's per-token timestamp + confidence:
// * frame = the collapsed token's `start_offset`. NeMo computes
// start_offset[0] = max(0, peak[0] - 1) and
// start_offset[i] = peak[i-1] for i>=1, where peak[i] is the
// argmax frame at which token i is emitted (the first frame of
// its consecutive run). I.e. the run-START frame: token i's
// start_offset is the PREVIOUS token's emit frame.
// * conf = MIN over token i's consecutive argmax run of the per-frame
// max_prob confidence (aggregation='min', exclude_blank=True).
// * span = 1.
const int N = vocab_plus_1; // num classes in the CTC softmax (vocab + blank).
int32_t previous = blank_id;
int prev_peak = 0; // emit frame (peak) of the previously emitted token.
bool have_prev = false; // any token emitted yet?
float cur_run_min = 1.0f; // running min per-frame conf over the current run.
for (int t = 0; t < T; ++t) {
const float* row = logits.data() + (size_t)t * vocab_plus_1;
// argmax over the vocab axis (ties resolve to the lowest index, as in
// PyTorch/NumPy argmax — matches NeMo's torch.argmax / prediction.max).
int32_t p = 0;
float best_val = row[0];
for (int v = 1; v < vocab_plus_1; ++v) {
if (row[v] > best_val) { best_val = row[v]; p = v; }
}
const bool emit = (p != previous || previous == blank_id) && p != blank_id;
if (emit) {
out.push_back(p);
if (tokens) {
// Close out the PREVIOUS token's accumulated run-min conf before
// starting the new token's run.
if (!tokens->empty()) tokens->back().conf = cur_run_min;
// start_offset: token 0 -> max(0, peak0 - 1); else previous peak.
const int32_t frame = have_prev ? (int32_t)prev_peak
: (int32_t)(t > 0 ? t - 1 : 0);
tokens->push_back(TokenInfo{ p, frame, /*conf placeholder*/0.0f, 1 });
cur_run_min = max_prob_conf(best_val, N); // first frame of new run.
prev_peak = t;
have_prev = true;
}
} else if (p == previous && p != blank_id) {
// Continuation of the current token's consecutive argmax run.
if (tokens) {
const float c = max_prob_conf(best_val, N);
if (c < cur_run_min) cur_run_min = c;
}
}
previous = p;
}
// Finalize the last token's run-min.
if (tokens && !tokens->empty()) tokens->back().conf = cur_run_min;
return out;
}
} // namespace pk