File size: 3,817 Bytes
bf1bec7 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 | #pragma once
#include <cstdint>
#include <array>
#include <list>
#include <mutex>
#include <shared_mutex>
#include <string>
#include <string_view>
#include <tuple>
#include <unordered_map>
#include <vector>
namespace curriculum {
/// Lossless byte-level BPE tokenizer for curriculum-trained vocabularies.
///
/// The Rust trainer and this C++ inference implementation share:
/// - the same deterministic pre-tokenization rules,
/// - numeric merge rules (left ID, right ID, result ID),
/// - byte-exact decoding.
class Tokenizer {
public:
Tokenizer() = default;
bool load(const std::string& vocab_path, const std::string& merges_path);
std::vector<int32_t> encode(const std::string& text) const;
std::string decode(const std::vector<int32_t>& ids) const;
std::vector<std::vector<int32_t>> encode_batch(
const std::vector<std::string>& texts, size_t num_threads = 0) const;
size_t vocab_size() const;
void set_cache_capacity(size_t capacity);
void clear_cache();
std::unordered_map<std::string, size_t> cache_info() const;
/// Safe, printable token representation used by vocab.json.
std::string id_to_token(int32_t id) const;
/// Find an ID from literal token text/bytes. Returns -1 when absent.
int32_t token_to_id(const std::string& token) const;
void set_special_tokens(const std::vector<std::string>& tokens);
void configure_special_tokens(
const std::unordered_map<std::string, int32_t>& token_ids,
const std::unordered_map<int32_t, std::string>& aliases,
size_t expected_vocab_size);
// Optional decode-only aliases explicitly declared by legacy artifacts.
void set_special_token_aliases(const std::unordered_map<int32_t, std::string>& aliases);
private:
struct EncodingScratch;
struct CachedPiece {
std::string bytes;
std::vector<int32_t> ids;
};
struct CacheShard {
mutable std::mutex mutex;
std::list<CachedPiece> lru;
// Views refer to owned, stable strings in the LRU list.
std::unordered_map<std::string_view, std::list<CachedPiece>::iterator> index;
size_t payload_bytes = 0;
size_t hits = 0;
size_t misses = 0;
};
static constexpr size_t kCacheShards = 16;
static constexpr size_t kCachePayloadBytes = 4 * 1024 * 1024;
static constexpr size_t kMaxCachedPieceBytes = 256;
mutable std::shared_mutex model_mutex_;
mutable std::array<CacheShard, kCacheShards> piece_cache_;
size_t cache_capacity_ = 16384;
struct MergeInfo {
size_t rank;
int32_t result;
};
std::vector<std::string> id_to_token_; // raw token bytes
std::unordered_map<std::string, int32_t> token_to_id_; // raw bytes -> ID
// Ordered merge rules: left, right, result.
std::vector<std::tuple<int32_t, int32_t, int32_t>> merges_;
// Earliest rule for each pair. Results are unique in trained artifacts,
// so later duplicate rules can never match after the first rule fires.
std::unordered_map<uint64_t, MergeInfo> merge_lookup_;
bool fast_merge_order_ = true;
std::unordered_map<std::string, int32_t> special_tokens_;
std::unordered_map<int32_t, std::string> special_token_aliases_;
std::vector<int32_t> encode_impl(const std::string& text) const;
void append_regular(std::string_view text, std::vector<int32_t>& output, EncodingScratch& scratch) const;
void append_word(std::string_view bytes, std::vector<int32_t>& output, EncodingScratch& scratch) const;
void encode_word_into(std::string_view bytes, EncodingScratch& scratch) const;
void invalidate_cache(); // Caller holds the exclusive model lock.
bool load_vocab(const std::string& path);
bool load_merges(const std::string& path);
};
} // namespace curriculum
|