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;
}