23 lines
380 B
C++
23 lines
380 B
C++
#pragma once
|
|
|
|
#include <cstdint>
|
|
#include <unordered_map>
|
|
#include <vector>
|
|
|
|
namespace ngram {
|
|
|
|
struct Result {
|
|
std::vector<int32_t> token;
|
|
std::vector<uint8_t> mask;
|
|
|
|
void truncate(size_t n);
|
|
};
|
|
|
|
struct Node {
|
|
std::unordered_map<int32_t, int32_t> next;
|
|
};
|
|
|
|
Result fillResult(int last_token, int draft_token_num, std::vector<Node>& tree, int root);
|
|
|
|
} // namespace ngram
|