| File: | root/firefox-clang/third_party/parakeet.cpp/src/tdt.cpp |
| Warning: | line 114, column 24 Value stored to 'skip' is never read |
Press '?' to see keyboard shortcuts
Keyboard shortcuts:
| 1 | #include "tdt.hpp" |
| 2 | #include "decode_common.hpp" |
| 3 | #include <algorithm> |
| 4 | #include <cassert> |
| 5 | #include <cmath> |
| 6 | #include <stdexcept> |
| 7 | #include <utility> |
| 8 | #include "moz-overrides.h" |
| 9 | |
| 10 | namespace pk { |
| 11 | |
| 12 | std::vector<int32_t> tdt_greedy(const PredictionNet& pred, const Joint& joint, |
| 13 | const std::vector<float>& enc, int T, int enc_hidden, |
| 14 | const std::vector<int32_t>& durations, |
| 15 | int blank_id, int max_symbols, |
| 16 | std::vector<TokenInfo>* tokens) { |
| 17 | assert((int)enc.size() == (size_t)T * enc_hidden)(static_cast <bool> ((int)enc.size() == (size_t)T * enc_hidden ) ? void (0) : __assert_fail ("(int)enc.size() == (size_t)T * enc_hidden" , __builtin_FILE (), __builtin_LINE (), __extension__ __PRETTY_FUNCTION__ )); |
| 18 | assert(!durations.empty())(static_cast <bool> (!durations.empty()) ? void (0) : __assert_fail ("!durations.empty()", __builtin_FILE (), __builtin_LINE (), __extension__ __PRETTY_FUNCTION__)); |
| 19 | |
| 20 | const int V_plus = joint.V_plus(); |
| 21 | const int num_dur = (int)durations.size(); |
| 22 | const int token_count = V_plus - num_dur; // vocab + 1 (incl. blank) = 1025 |
| 23 | assert(token_count == joint.vocab_size() + 1)(static_cast <bool> (token_count == joint.vocab_size() + 1) ? void (0) : __assert_fail ("token_count == joint.vocab_size() + 1" , __builtin_FILE (), __builtin_LINE (), __extension__ __PRETTY_FUNCTION__ )); |
| 24 | assert(num_dur == joint.num_durations())(static_cast <bool> (num_dur == joint.num_durations()) ? void (0) : __assert_fail ("num_dur == joint.num_durations()" , __builtin_FILE (), __builtin_LINE (), __extension__ __PRETTY_FUNCTION__ )); |
| 25 | |
| 26 | std::vector<int32_t> hyp; |
| 27 | if (tokens) tokens->clear(); |
| 28 | |
| 29 | // Committed (non-blank) decoding state and last emitted token. |
| 30 | PredState committed = pred.zero_state(); |
| 31 | int32_t last_token = -1; // -1 sentinel: nothing emitted yet -> SOS. |
| 32 | bool emitted_any = false; |
| 33 | |
| 34 | // Precompute the encoder projection over ALL frames ONCE (one matmul on the |
| 35 | // persistent backend), reused for every step. The per-step joint below is a |
| 36 | // tight churn-free graph on the same backend. The old code rebuilt the full |
| 37 | // joint per (t,u) step (a fresh graph per step — the bulk of the |
| 38 | // per-utterance graph dispatches). |
| 39 | std::vector<float> enc_proj; // row-major [T, joint_hidden] |
| 40 | joint.precompute_enc_proj(enc, T, enc_hidden, enc_proj); |
| 41 | const int H = joint.joint_hidden(); |
| 42 | |
| 43 | // Scratch reused across inner steps. |
| 44 | std::vector<float> g; |
| 45 | PredState out_state; |
| 46 | std::vector<float> logits; |
| 47 | |
| 48 | // Prediction-net output cache: `g`/`out_state` depend only on the committed |
| 49 | // (last_token, lstm_state), which change exclusively on an emit (k != blank). |
| 50 | // Steps that don't emit reuse the cached forward pass instead of recomputing |
| 51 | // the LSTM. See the RNN-T loop for the full rationale. |
| 52 | bool g_valid = false; |
| 53 | |
| 54 | int t = 0; |
| 55 | while (t < T) { |
| 56 | int symbols_added = 0; |
| 57 | bool need_loop = true; |
| 58 | int skip = 0; |
| 59 | |
| 60 | while (need_loop && symbols_added < max_symbols) { |
| 61 | // Prediction net step from the committed state — only when the cache |
| 62 | // is stale (first step, or the previous step emitted). SOS until the |
| 63 | // first emit; otherwise feed the last EMITTED token. |
| 64 | if (!g_valid) { |
| 65 | const bool is_sos = !emitted_any; |
| 66 | const int32_t last_label = emitted_any ? last_token : blank_id; |
| 67 | pred.step(last_label, is_sos, committed, g, out_state); |
| 68 | g_valid = true; |
| 69 | } |
| 70 | |
| 71 | // Joint for (t,u): precomputed enc_proj[t] x g -> raw logits [V_plus]. |
| 72 | // `t` is always in [0, T): the outer loop guards the first inner |
| 73 | // iteration, and subsequent inner iterations only run when skip==0 |
| 74 | // (t unchanged). A positive skip exits the inner loop. This matches |
| 75 | // the old per-step `enc[t*enc_hidden]` access exactly. |
| 76 | assert(t < T && "enc_proj row out of range")(static_cast <bool> (t < T && "enc_proj row out of range" ) ? void (0) : __assert_fail ("t < T && \"enc_proj row out of range\"" , __builtin_FILE (), __builtin_LINE (), __extension__ __PRETTY_FUNCTION__ )); |
| 77 | joint.step_logits(enc_proj.data() + (size_t)t * H, |
| 78 | g.data(), (int)g.size(), logits); |
| 79 | |
| 80 | // Split: token logits [0, token_count), duration logits [token_count, V_plus). |
| 81 | const int k = decode_argmax(logits.data(), token_count); |
| 82 | const int d_k = decode_argmax(logits.data() + token_count, num_dur); |
| 83 | skip = durations[d_k]; |
| 84 | |
| 85 | // Commit state + last_token ONLY when k != blank. |
| 86 | if (k != blank_id) { |
| 87 | hyp.push_back((int32_t)k); |
| 88 | if (tokens) { |
| 89 | // NeMo per-token metadata (matches GreedyTDTInfer._greedy_decode |
| 90 | // + max_prob confidence): |
| 91 | // frame = the encoder frame t at emission (hypothesis.timestamp). |
| 92 | // conf = max_prob over the TOKEN slice logits[0:vocab+1] |
| 93 | // (NeMo log_softmaxes that slice; exclude the duration |
| 94 | // logits). N = token_count = vocab + 1. |
| 95 | // span = durations[d_k] (the duration/skip applied to the token). |
| 96 | const float conf = decode_max_prob_conf(logits.data(), token_count, k); |
| 97 | tokens->push_back(TokenInfo{ (int32_t)k, (int32_t)t, conf, |
| 98 | (int32_t)skip }); |
| 99 | } |
| 100 | last_token = (int32_t)k; |
| 101 | committed = out_state; // carry the step's new (h', c') |
| 102 | emitted_any = true; |
| 103 | g_valid = false; // committed state advanced -> recompute g |
| 104 | } |
| 105 | // else: discard out_state; committed/last_token unchanged (g stays valid). |
| 106 | |
| 107 | symbols_added += 1; |
| 108 | t += skip; |
| 109 | need_loop = (skip == 0); |
| 110 | } |
| 111 | |
| 112 | // Infinite-loop guard: if we exited with duration 0 (blank + dur 0), step |
| 113 | // forward by one frame anyway. |
| 114 | if (skip == 0) skip = 1; |
Value stored to 'skip' is never read | |
| 115 | |
| 116 | // If we stopped because max_symbols was hit (not because of a positive |
| 117 | // duration), advance the frame by one to make progress. |
| 118 | if (symbols_added == max_symbols) t += 1; |
| 119 | } |
| 120 | |
| 121 | return hyp; |
| 122 | } |
| 123 | |
| 124 | namespace detail { |
| 125 | |
| 126 | int best_positive_duration_index( |
| 127 | const std::vector<float>& scores, |
| 128 | const std::vector<int32_t>& durations) { |
| 129 | if (scores.size() != durations.size()) |
| 130 | throwabort_with_suppression(); if (false) std::invalid_argument( |
| 131 | "best_positive_duration_index: size mismatch"); |
| 132 | |
| 133 | int best = -1; |
| 134 | for (size_t i = 0; i < durations.size(); ++i) { |
| 135 | if (durations[i] > 0 && |
| 136 | (best < 0 || scores[i] > scores[(size_t)best])) { |
| 137 | best = (int)i; |
| 138 | } |
| 139 | } |
| 140 | return best; |
| 141 | } |
| 142 | |
| 143 | } // namespace detail |
| 144 | |
| 145 | namespace { |
| 146 | |
| 147 | struct BeamState { |
| 148 | TdtBeamHypothesis hyp; |
| 149 | PredState dec_state; |
| 150 | PredState next_state; |
| 151 | std::vector<float> pred_out; |
| 152 | int last_frame = 0; |
| 153 | int32_t last_token = -1; |
| 154 | bool emitted_any = false; |
| 155 | bool pred_valid = false; |
| 156 | }; |
| 157 | |
| 158 | struct RankedIndex { |
| 159 | float score; |
| 160 | int index; |
| 161 | }; |
| 162 | |
| 163 | std::vector<float> log_softmax(const float* logits, int n) { |
| 164 | float max_logit = logits[0]; |
| 165 | for (int i = 1; i < n; ++i) |
| 166 | max_logit = std::max(max_logit, logits[i]); |
| 167 | |
| 168 | double sum = 0.0; |
| 169 | for (int i = 0; i < n; ++i) |
| 170 | sum += std::exp((double)logits[i] - (double)max_logit); |
| 171 | const double log_denom = std::log(sum); |
| 172 | |
| 173 | std::vector<float> out(n); |
| 174 | for (int i = 0; i < n; ++i) |
| 175 | out[i] = (float)((double)logits[i] - (double)max_logit - log_denom); |
| 176 | return out; |
| 177 | } |
| 178 | |
| 179 | std::vector<RankedIndex> top_k(const float* scores, int n, int k) { |
| 180 | std::vector<RankedIndex> ranked; |
| 181 | ranked.reserve(n); |
| 182 | for (int i = 0; i < n; ++i) |
| 183 | ranked.push_back(RankedIndex{scores[i], i}); |
| 184 | |
| 185 | const auto better = [](const RankedIndex& a, const RankedIndex& b) { |
| 186 | if (a.score != b.score) return a.score > b.score; |
| 187 | return a.index < b.index; |
| 188 | }; |
| 189 | const int kept = std::min(n, k); |
| 190 | std::partial_sort(ranked.begin(), ranked.begin() + kept, ranked.end(), better); |
| 191 | ranked.resize(kept); |
| 192 | return ranked; |
| 193 | } |
| 194 | |
| 195 | float log_add_exp(float a, float b) { |
| 196 | const float hi = std::max(a, b); |
| 197 | const float lo = std::min(a, b); |
| 198 | return hi + std::log1p(std::exp(lo - hi)); |
| 199 | } |
| 200 | |
| 201 | bool same_path(const BeamState& a, const BeamState& b) { |
| 202 | if (a.last_frame != b.last_frame || a.hyp.tokens.size() != b.hyp.tokens.size()) |
| 203 | return false; |
| 204 | for (size_t i = 0; i < a.hyp.tokens.size(); ++i) |
| 205 | if (a.hyp.tokens[i].id != b.hyp.tokens[i].id) |
| 206 | return false; |
| 207 | return true; |
| 208 | } |
| 209 | |
| 210 | void merge_duplicate_hypotheses(std::vector<BeamState>& hypotheses) { |
| 211 | std::sort(hypotheses.begin(), hypotheses.end(), |
| 212 | [](const BeamState& a, const BeamState& b) { |
| 213 | return a.hyp.score > b.hyp.score; |
| 214 | }); |
| 215 | |
| 216 | std::vector<BeamState> merged; |
| 217 | merged.reserve(hypotheses.size()); |
| 218 | for (BeamState& candidate : hypotheses) { |
| 219 | auto duplicate = std::find_if( |
| 220 | merged.begin(), merged.end(), |
| 221 | [&](const BeamState& kept) { return same_path(kept, candidate); }); |
| 222 | if (duplicate == merged.end()) { |
| 223 | merged.push_back(std::move(candidate)); |
| 224 | } else { |
| 225 | duplicate->hyp.score = log_add_exp(duplicate->hyp.score, |
| 226 | candidate.hyp.score); |
| 227 | } |
| 228 | } |
| 229 | hypotheses.swap(merged); |
| 230 | } |
| 231 | |
| 232 | float sort_score(const BeamState& hyp, bool score_norm) { |
| 233 | if (!score_norm) return hyp.hyp.score; |
| 234 | return hyp.hyp.score / (float)(hyp.hyp.tokens.size() + 1); |
| 235 | } |
| 236 | |
| 237 | } // namespace |
| 238 | |
| 239 | std::vector<TdtBeamHypothesis> tdt_beam_search( |
| 240 | const PredictionNet& pred, const Joint& joint, |
| 241 | const std::vector<float>& enc, int T, int enc_hidden, |
| 242 | const std::vector<int32_t>& durations, int blank_id, |
| 243 | int beam_size, int nbest, bool score_norm) { |
| 244 | if (T < 0 || enc_hidden <= 0 || |
| 245 | enc.size() != (size_t)T * (size_t)enc_hidden) |
| 246 | throwabort_with_suppression(); if (false) std::invalid_argument("tdt_beam_search: invalid encoder shape"); |
| 247 | if (durations.empty()) |
| 248 | throwabort_with_suppression(); if (false) std::invalid_argument("tdt_beam_search: durations must not be empty"); |
| 249 | if (beam_size < 1 || nbest < 1 || nbest > beam_size) |
| 250 | throwabort_with_suppression(); if (false) std::invalid_argument( |
| 251 | "tdt_beam_search: require beam_size >= nbest >= 1"); |
| 252 | |
| 253 | const int num_durations = (int)durations.size(); |
| 254 | const int token_count = joint.V_plus() - num_durations; |
| 255 | if (token_count != joint.vocab_size() + 1 || |
| 256 | num_durations != joint.num_durations() || |
| 257 | blank_id != token_count - 1 || blank_id < 1) { |
| 258 | throwabort_with_suppression(); if (false) std::invalid_argument("tdt_beam_search: incompatible TDT head"); |
| 259 | } |
| 260 | |
| 261 | int zero_duration_idx = -1; |
| 262 | int zero_duration_count = 0; |
| 263 | bool has_positive_duration = false; |
| 264 | for (int i = 0; i < num_durations; ++i) { |
| 265 | if (durations[i] < 0) |
| 266 | throwabort_with_suppression(); if (false) std::invalid_argument( |
| 267 | "tdt_beam_search: durations must be non-negative"); |
| 268 | if (durations[i] == 0) { |
| 269 | zero_duration_idx = i; |
| 270 | ++zero_duration_count; |
| 271 | } |
| 272 | if (durations[i] > 0) |
| 273 | has_positive_duration = true; |
| 274 | } |
| 275 | if (!has_positive_duration) |
| 276 | throwabort_with_suppression(); if (false) std::invalid_argument( |
| 277 | "tdt_beam_search: at least one positive duration is required"); |
| 278 | if (zero_duration_count > 1) |
| 279 | throwabort_with_suppression(); if (false) std::invalid_argument( |
| 280 | "tdt_beam_search: at most one zero duration is allowed"); |
| 281 | if (T == 0) |
| 282 | return {TdtBeamHypothesis{}}; |
| 283 | |
| 284 | const int beam = std::min(beam_size, blank_id); |
| 285 | const int token_beam = std::min(beam, blank_id - 1); |
| 286 | const int duration_beam = std::min(beam, num_durations); |
| 287 | |
| 288 | std::vector<float> enc_proj; |
| 289 | joint.precompute_enc_proj(enc, T, enc_hidden, enc_proj); |
| 290 | const int joint_hidden = joint.joint_hidden(); |
| 291 | |
| 292 | BeamState start; |
| 293 | start.dec_state = pred.zero_state(); |
| 294 | std::vector<BeamState> kept_hyps; |
| 295 | kept_hyps.push_back(std::move(start)); |
| 296 | |
| 297 | std::vector<float> logits; |
| 298 | for (int time_idx = 0; time_idx < T; ++time_idx) { |
| 299 | std::vector<BeamState> current_hyps; |
| 300 | std::vector<BeamState> future_hyps; |
| 301 | current_hyps.reserve(kept_hyps.size()); |
| 302 | future_hyps.reserve(kept_hyps.size()); |
| 303 | for (BeamState& hyp : kept_hyps) { |
| 304 | if (hyp.last_frame == time_idx) |
| 305 | current_hyps.push_back(std::move(hyp)); |
| 306 | else if (hyp.last_frame > time_idx) |
| 307 | future_hyps.push_back(std::move(hyp)); |
| 308 | } |
| 309 | kept_hyps.swap(future_hyps); |
| 310 | |
| 311 | while (!current_hyps.empty()) { |
| 312 | auto best_it = std::max_element( |
| 313 | current_hyps.begin(), current_hyps.end(), |
| 314 | [](const BeamState& a, const BeamState& b) { |
| 315 | return a.hyp.score < b.hyp.score; |
| 316 | }); |
| 317 | BeamState best = std::move(*best_it); |
| 318 | current_hyps.erase(best_it); |
| 319 | |
| 320 | if (!best.pred_valid) { |
| 321 | const bool is_sos = !best.emitted_any; |
| 322 | const int32_t label = best.emitted_any ? best.last_token : blank_id; |
| 323 | pred.step(label, is_sos, best.dec_state, |
| 324 | best.pred_out, best.next_state); |
| 325 | best.pred_valid = true; |
| 326 | } |
| 327 | |
| 328 | joint.step_logits( |
| 329 | enc_proj.data() + (size_t)time_idx * joint_hidden, |
| 330 | best.pred_out.data(), (int)best.pred_out.size(), logits); |
| 331 | std::vector<float> token_logp = |
| 332 | log_softmax(logits.data(), token_count); |
| 333 | std::vector<float> duration_logp = |
| 334 | log_softmax(logits.data() + token_count, num_durations); |
| 335 | if (!std::all_of(token_logp.begin(), token_logp.end(), |
| 336 | [](float value) { return std::isfinite(value); }) || |
| 337 | !std::all_of(duration_logp.begin(), duration_logp.end(), |
| 338 | [](float value) { return std::isfinite(value); })) { |
| 339 | throwabort_with_suppression(); if (false) std::runtime_error( |
| 340 | "tdt_beam_search: non-finite log probability"); |
| 341 | } |
| 342 | |
| 343 | const std::vector<RankedIndex> best_durations = |
| 344 | top_k(duration_logp.data(), num_durations, duration_beam); |
| 345 | |
| 346 | struct Pair { |
| 347 | float score; |
| 348 | int token; |
| 349 | int duration_idx; |
| 350 | }; |
| 351 | std::vector<Pair> pairs; |
| 352 | const std::vector<RankedIndex> best_tokens = |
| 353 | top_k(token_logp.data(), blank_id, token_beam); |
| 354 | pairs.reserve(best_tokens.size() * best_durations.size()); |
| 355 | for (const RankedIndex& duration : best_durations) |
| 356 | for (const RankedIndex& token : best_tokens) |
| 357 | pairs.push_back(Pair{ |
| 358 | duration.score + token.score, |
| 359 | token.index, |
| 360 | duration.index}); |
| 361 | std::partial_sort( |
| 362 | pairs.begin(), |
| 363 | pairs.begin() + std::min(token_beam, (int)pairs.size()), |
| 364 | pairs.end(), |
| 365 | [](const Pair& a, const Pair& b) { |
| 366 | if (a.score != b.score) return a.score > b.score; |
| 367 | if (a.duration_idx != b.duration_idx) |
| 368 | return a.duration_idx < b.duration_idx; |
| 369 | return a.token < b.token; |
| 370 | }); |
| 371 | pairs.resize(std::min(token_beam, (int)pairs.size())); |
| 372 | |
| 373 | for (const Pair& pair : pairs) { |
| 374 | const int duration = durations[pair.duration_idx]; |
| 375 | BeamState child = best; |
| 376 | child.hyp.score += pair.score; |
| 377 | // A near-certain label's log-prob can round to no change in |
| 378 | // the running score, so only an increase is an error. |
| 379 | if (duration == 0 && child.hyp.score > best.hyp.score) { |
| 380 | throwabort_with_suppression(); if (false) std::runtime_error( |
| 381 | "tdt_beam_search: zero-duration expansion " |
| 382 | "increased score"); |
| 383 | } |
| 384 | child.hyp.tokens.push_back(TdtBeamToken{ |
| 385 | (int32_t)pair.token, (int32_t)time_idx, (int32_t)duration}); |
| 386 | child.dec_state = best.next_state; |
| 387 | child.last_token = (int32_t)pair.token; |
| 388 | child.emitted_any = true; |
| 389 | child.pred_valid = false; |
| 390 | child.last_frame += duration; |
| 391 | if (duration == 0) |
| 392 | current_hyps.push_back(std::move(child)); |
| 393 | else |
| 394 | kept_hyps.push_back(std::move(child)); |
| 395 | } |
| 396 | |
| 397 | for (const RankedIndex& ranked_duration : best_durations) { |
| 398 | int duration_idx = ranked_duration.index; |
| 399 | if (duration_idx == zero_duration_idx) { |
| 400 | if (best_durations.size() != 1) |
| 401 | continue; |
| 402 | duration_idx = detail::best_positive_duration_index( |
| 403 | duration_logp, durations); |
| 404 | } |
| 405 | BeamState child = best; |
| 406 | child.hyp.score += token_logp[blank_id] + |
| 407 | duration_logp[duration_idx]; |
| 408 | child.last_frame += durations[duration_idx]; |
| 409 | kept_hyps.push_back(std::move(child)); |
| 410 | } |
| 411 | |
| 412 | merge_duplicate_hypotheses(kept_hyps); |
| 413 | if (!current_hyps.empty()) { |
| 414 | const float next_score = std::max_element( |
| 415 | current_hyps.begin(), current_hyps.end(), |
| 416 | [](const BeamState& a, const BeamState& b) { |
| 417 | return a.hyp.score < b.hyp.score; |
| 418 | })->hyp.score; |
| 419 | const int more_probable_count = (int)std::count_if( |
| 420 | kept_hyps.begin(), kept_hyps.end(), |
| 421 | [next_score](const BeamState& hyp) { |
| 422 | return hyp.hyp.score > next_score; |
| 423 | }); |
| 424 | if (more_probable_count >= beam) { |
| 425 | std::vector<BeamState> more_probable; |
| 426 | more_probable.reserve(more_probable_count); |
| 427 | for (BeamState& hyp : kept_hyps) |
| 428 | if (hyp.hyp.score > next_score) |
| 429 | more_probable.push_back(std::move(hyp)); |
| 430 | kept_hyps.swap(more_probable); |
| 431 | break; |
| 432 | } |
| 433 | } else { |
| 434 | std::sort(kept_hyps.begin(), kept_hyps.end(), |
| 435 | [](const BeamState& a, const BeamState& b) { |
| 436 | return a.hyp.score > b.hyp.score; |
| 437 | }); |
| 438 | if ((int)kept_hyps.size() > beam) |
| 439 | kept_hyps.resize(beam); |
| 440 | } |
| 441 | } |
| 442 | } |
| 443 | |
| 444 | std::sort(kept_hyps.begin(), kept_hyps.end(), |
| 445 | [score_norm](const BeamState& a, const BeamState& b) { |
| 446 | return sort_score(a, score_norm) > sort_score(b, score_norm); |
| 447 | }); |
| 448 | const int count = std::min(nbest, (int)kept_hyps.size()); |
| 449 | std::vector<TdtBeamHypothesis> result; |
| 450 | result.reserve(count); |
| 451 | for (int i = 0; i < count; ++i) { |
| 452 | kept_hyps[i].hyp.normalized_score = |
| 453 | kept_hyps[i].hyp.score / |
| 454 | (float)(kept_hyps[i].hyp.tokens.size() + 1); |
| 455 | result.push_back(std::move(kept_hyps[i].hyp)); |
| 456 | } |
| 457 | return result; |
| 458 | } |
| 459 | |
| 460 | } // namespace pk |