transcription.hpp - mozsearch

Enable keyboard shortcuts

#pragma once

#include "decode_types.hpp"

#include <string>

#include <vector>

namespace pk {

// One decoded word with its time span (seconds) and aggregate confidence.

//

//   text  : the word string (SentencePiece pieces detokenized — `▁`->space,

//           leading space stripped, punctuation attached).

//   start : word start time in seconds (= first token frame * frame_sec).

//   end   : word end time in seconds (= last token end_offset * frame_sec; see

//           group_words for the per-head end_offset convention).

//   conf  : NeMo's 'min' aggregate of the word's per-token confidences.

struct Word {

    std::string text;

    float       start = 0.0f;

    float       end   = 0.0f;

    float       conf  = 0.0f;

};

// A full transcription result: the flat text, the per-word timestamps +

// confidence, and the raw per-token metadata the words were grouped from.

struct Transcription {

    std::string            text;

    std::vector<Word>      words;

    std::vector<TokenInfo> tokens;

};

// Group a per-token decode (TokenInfo sequence, in emission order) into words,

// matching NeMo's word-offset + confidence convention (timestamps=True,

// confidence aggregation='min').

//

// `pieces` is the tokenizer piece table (index by TokenInfo.id). A token whose

// piece begins with `▁` (U+2581) starts a new word; the leading-`▁` of the

// utterance starts the first word too. Consecutive non-`▁` pieces extend the

// current word. Word text = detokenize(the word's piece ids). A supported

// punctuation mark (extracted from `pieces` like NeMo

// extract_punctuation_from_vocab) attaches to the preceding word and has its

// own start/end refined to the previous token's end (NeMo _refine_timestamps).

//

// Per-token offsets (frames): start_offset = TokenInfo.frame; end_offset =

// TokenInfo.frame + TokenInfo.span. For TDT, span is the predicted duration

// (NeMo _compute_offsets_tdt: end = start + duration). For CTC, NeMo's

// end_offset is the NEXT token's start frame (cumulative run lengths); the

// caller therefore sets each CTC token's span to (next_frame - frame) so the

// `frame + span` rule reproduces NeMo exactly (the id-only ctc_greedy keeps

// span == 1; transcribe_with_timestamps does the run-length rewrite).

//

//   word.start = first_token.frame                      * frame_sec

//   word.end   = (last_token.frame + last_token.span)   * frame_sec

//   word.conf  = min over the word's token confidences

std::vector<Word> group_words(const std::vector<TokenInfo>& tokens,

                              const std::vector<std::string>& pieces,

                              float frame_sec);

} // namespace pk