Candidate SuffixTree::_speculate_path()

in csrc/suffix_decoding/suffix_tree.cc [424:458]


Candidate SuffixTree::_speculate_path(Node* node, int idx,
                                      int max_spec_tokens,
                                      float min_token_prob) {
    Candidate ret;
    float prob = 1.0f;
    while (ret.token_ids.size() < max_spec_tokens && prob >= min_token_prob) {
        if (idx < node->length) {
            // Use previous token index as parent; if none, mark as -1.
            ret.parents.push_back(static_cast<int>(ret.token_ids.size()) - 1);
            int token = _seqs[node->ref_seq][node->ref_idx + idx];
            ret.token_ids.push_back(token);
            ret.probs.push_back(prob);
            ret.score += prob;
            idx++;
        } else {
            Node* child = nullptr;
            int64_t count = 0;
            // Choose the child with the maximum count.
            for (const auto& kv : node->children) {
                Node* ch = kv.second.get();
                if (ch->count > count) {
                    child = ch;
                    count = ch->count;
                }
            }
            if (child == nullptr) {
                break;
            }
            prob *= static_cast<float>(count) / node->count;
            node = child;
            idx = 0;
        }
    }
    return ret;
}