634 lines
22 KiB
C++
634 lines
22 KiB
C++
#pragma once
|
|
// bpe.h, Qwen3/GPT-2 byte-level BPE tokenizer (CPU-only, no dependencies)
|
|
//
|
|
// Parses the tokenizer.json fields stored in a GGUF model and produces a
|
|
// byte-level BPE encoder/decoder. Arch-specific special tokens (text
|
|
// markers, language tags, audio sentinels) are loaded through
|
|
// bpe_load_specials_from_keys with a caller-provided list of GGUF KV keys.
|
|
// Loads vocab + merges from a GGUF tokenizer payload. Handles byte-level
|
|
// encoding, GPT-2 regex pre-tokenizer, BPE merges, and a registry of
|
|
// verbatim special tokens (endoftext plus any caller-registered tokens
|
|
// such as TTS style markers and language tags).
|
|
|
|
#include "gguf.h"
|
|
|
|
#include <cassert>
|
|
#include <climits>
|
|
#include <cstdio>
|
|
#include <cstring>
|
|
#include <string>
|
|
#include <unordered_map>
|
|
#include <vector>
|
|
|
|
// GPT-2 byte-level encoding table
|
|
// Maps byte [0..255] -> Unicode char for BPE vocab keys.
|
|
// Printable ASCII stays as-is, control/space bytes get remapped.
|
|
static void build_byte_encoder(std::string byte2str[256]) {
|
|
// Standard GPT-2 byte encoder
|
|
int bs[256], cs[256], n = 0, total = 0;
|
|
// Printable ranges that map to themselves
|
|
for (int b = '!'; b <= '~'; b++) {
|
|
bs[total] = b;
|
|
cs[total] = b;
|
|
total++;
|
|
}
|
|
for (int b = 0xA1; b <= 0xAC; b++) {
|
|
bs[total] = b;
|
|
cs[total] = b;
|
|
total++;
|
|
}
|
|
for (int b = 0xAE; b <= 0xFF; b++) {
|
|
bs[total] = b;
|
|
cs[total] = b;
|
|
total++;
|
|
}
|
|
// Remaining bytes get mapped to 256+
|
|
bool used[256] = {};
|
|
for (int i = 0; i < total; i++) {
|
|
used[bs[i]] = true;
|
|
}
|
|
for (int b = 0; b < 256; b++) {
|
|
if (!used[b]) {
|
|
bs[total] = b;
|
|
cs[total] = 256 + n;
|
|
n++;
|
|
total++;
|
|
}
|
|
}
|
|
assert(total == 256);
|
|
// Convert codepoints to UTF-8 strings
|
|
for (int i = 0; i < 256; i++) {
|
|
int cp = cs[i];
|
|
char buf[4];
|
|
int len;
|
|
if (cp < 0x80) {
|
|
buf[0] = (char) cp;
|
|
len = 1;
|
|
} else if (cp < 0x800) {
|
|
buf[0] = (char) (0xC0 | (cp >> 6));
|
|
buf[1] = (char) (0x80 | (cp & 0x3F));
|
|
len = 2;
|
|
} else {
|
|
buf[0] = (char) (0xE0 | (cp >> 12));
|
|
buf[1] = (char) (0x80 | ((cp >> 6) & 0x3F));
|
|
buf[2] = (char) (0x80 | (cp & 0x3F));
|
|
len = 3;
|
|
}
|
|
byte2str[bs[i]] = std::string(buf, len);
|
|
}
|
|
}
|
|
|
|
// UTF-8 helpers
|
|
static int utf8_codepoint(const char * s, int * advance) {
|
|
unsigned char c = s[0];
|
|
if (c < 0x80) {
|
|
*advance = 1;
|
|
return c;
|
|
}
|
|
if ((c & 0xE0) == 0xC0) {
|
|
*advance = 2;
|
|
return ((c & 0x1F) << 6) | (s[1] & 0x3F);
|
|
}
|
|
if ((c & 0xF0) == 0xE0) {
|
|
*advance = 3;
|
|
return ((c & 0x0F) << 12) | ((s[1] & 0x3F) << 6) | (s[2] & 0x3F);
|
|
}
|
|
if ((c & 0xF8) == 0xF0) {
|
|
*advance = 4;
|
|
return ((c & 0x07) << 18) | ((s[1] & 0x3F) << 12) | ((s[2] & 0x3F) << 6) | (s[3] & 0x3F);
|
|
}
|
|
// invalid lead byte: advance one and return the raw byte to avoid an
|
|
// infinite loop. The Python tokenizer never reaches this path since Python
|
|
// str guarantees valid UTF 8; in C++ the std::string input has no such
|
|
// guarantee, so this branch handles malformed input defensively.
|
|
*advance = 1;
|
|
return c;
|
|
}
|
|
|
|
// Unicode category checks (simplified but covers Latin + common scripts)
|
|
static bool is_letter(int cp) {
|
|
if ((cp >= 'A' && cp <= 'Z') || (cp >= 'a' && cp <= 'z')) {
|
|
return true;
|
|
}
|
|
if (cp < 0x80) {
|
|
return false;
|
|
}
|
|
// Latin Extended: U+00C0-U+00D6, U+00D8-U+00F6, U+00F8-U+01BF + Latin Extended-A/B
|
|
if (cp >= 0xC0 && cp <= 0x024F && cp != 0xD7 && cp != 0xF7) {
|
|
return true;
|
|
}
|
|
// Common CJK, Cyrillic, Greek, Arabic, etc., treat as letters
|
|
if (cp >= 0x0370 && cp <= 0x1FFF) {
|
|
return true; // Greek, Cyrillic, Armenian, etc.
|
|
}
|
|
if (cp >= 0x2C00 && cp <= 0x2DFF) {
|
|
return true; // Georgian, etc.
|
|
}
|
|
if (cp >= 0x3040 && cp <= 0x9FFF) {
|
|
return true; // CJK
|
|
}
|
|
if (cp >= 0xAC00 && cp <= 0xD7AF) {
|
|
return true; // Korean
|
|
}
|
|
if (cp >= 0xF900 && cp <= 0xFAFF) {
|
|
return true; // CJK compatibility
|
|
}
|
|
if (cp >= 0x10000) {
|
|
return true; // SMP, mostly letters/symbols
|
|
}
|
|
return false;
|
|
}
|
|
|
|
static bool is_digit(int cp) {
|
|
return cp >= '0' && cp <= '9';
|
|
}
|
|
|
|
static bool is_whitespace(int cp) {
|
|
return cp == ' ' || cp == '\t' || cp == '\n' || cp == '\r' || cp == 0x0B || cp == 0x0C || cp == 0xA0 ||
|
|
cp == 0x2000 || cp == 0x2001 || cp == 0x2002 || cp == 0x200B;
|
|
}
|
|
|
|
static bool is_newline(int cp) {
|
|
return cp == '\n' || cp == '\r';
|
|
}
|
|
|
|
// GPT-2 pre-tokenizer regex (manual implementation)
|
|
// Pattern: (?i:'s|'t|'re|'ve|'m|'ll|'d)|[^\r\n\p{L}\p{N}]?\p{L}+|\p{N}|
|
|
// \s?[^\s\p{L}\p{N}]+[\r\n]*|\s*[\r\n]+|\s+(?!\S)|\s+
|
|
// Splits text into non-overlapping chunks (on original text, not byte-encoded).
|
|
static std::vector<std::string> gpt2_pre_tokenize(const std::string & text) {
|
|
std::vector<std::string> chunks;
|
|
const char * s = text.c_str();
|
|
int len = (int) text.size();
|
|
int i = 0;
|
|
|
|
while (i < len) {
|
|
int adv;
|
|
int cp = utf8_codepoint(s + i, &adv);
|
|
|
|
// Rule 1: Contractions 's 't 're 've 'm 'll 'd
|
|
if ((cp == '\'' || cp == 0x2019) && i + adv < len) {
|
|
const char * rest = s + i + adv;
|
|
int rlen = len - i - adv;
|
|
auto try_match = [&](const char * suffix, int slen) -> bool {
|
|
if (rlen >= slen) {
|
|
// case-insensitive compare
|
|
for (int k = 0; k < slen; k++) {
|
|
char c1 = rest[k], c2 = suffix[k];
|
|
if (c1 >= 'A' && c1 <= 'Z') {
|
|
c1 = (char) (c1 + 32);
|
|
}
|
|
if (c1 != c2) {
|
|
return false;
|
|
}
|
|
}
|
|
// next char should NOT be a letter
|
|
if (rlen > slen) {
|
|
int a2;
|
|
int cp2 = utf8_codepoint(rest + slen, &a2);
|
|
if (is_letter(cp2)) {
|
|
return false;
|
|
}
|
|
}
|
|
chunks.push_back(std::string(s + i, adv + slen));
|
|
i += adv + slen;
|
|
return true;
|
|
}
|
|
return false;
|
|
};
|
|
if (try_match("ll", 2)) {
|
|
continue;
|
|
}
|
|
if (try_match("re", 2)) {
|
|
continue;
|
|
}
|
|
if (try_match("ve", 2)) {
|
|
continue;
|
|
}
|
|
if (try_match("s", 1)) {
|
|
continue;
|
|
}
|
|
if (try_match("t", 1)) {
|
|
continue;
|
|
}
|
|
if (try_match("m", 1)) {
|
|
continue;
|
|
}
|
|
if (try_match("d", 1)) {
|
|
continue;
|
|
}
|
|
}
|
|
|
|
// Rule 2: [^\r\n\p{L}\p{N}]?\p{L}+
|
|
if (is_letter(cp)) {
|
|
int start = i;
|
|
i += adv;
|
|
while (i < len) {
|
|
int a2;
|
|
int cp2 = utf8_codepoint(s + i, &a2);
|
|
if (!is_letter(cp2)) {
|
|
break;
|
|
}
|
|
i += a2;
|
|
}
|
|
chunks.push_back(std::string(s + start, i - start));
|
|
continue;
|
|
}
|
|
if (!is_newline(cp) && !is_letter(cp) && !is_digit(cp) && !is_whitespace(cp)) {
|
|
// Non-letter/number/space, check if followed by letters
|
|
int start = i;
|
|
int after = i + adv;
|
|
if (after < len) {
|
|
int a2;
|
|
int cp2 = utf8_codepoint(s + after, &a2);
|
|
if (is_letter(cp2)) {
|
|
i = after + a2;
|
|
while (i < len) {
|
|
int a3;
|
|
int cp3 = utf8_codepoint(s + i, &a3);
|
|
if (!is_letter(cp3)) {
|
|
break;
|
|
}
|
|
i += a3;
|
|
}
|
|
chunks.push_back(std::string(s + start, i - start));
|
|
continue;
|
|
}
|
|
}
|
|
}
|
|
|
|
// Rule 3: \p{N}+ (digits, consume consecutively)
|
|
if (is_digit(cp)) {
|
|
int start = i;
|
|
while (i < len && is_digit((unsigned char) s[i])) {
|
|
i++;
|
|
}
|
|
// GPT-2 regex matches single \p{N}, let's match one at a time
|
|
// to be safe, but in practice consecutive digits usually merge anyway.
|
|
// The regex is \p{N} (single digit), so split each digit:
|
|
for (int j = start; j < i; j++) {
|
|
chunks.push_back(std::string(s + j, 1));
|
|
}
|
|
continue;
|
|
}
|
|
|
|
// Rule 5: \s*[\r\n]+ (newlines with optional leading whitespace)
|
|
if (is_newline(cp)) {
|
|
int start = i;
|
|
while (i < len && is_newline((unsigned char) s[i])) {
|
|
i++;
|
|
}
|
|
chunks.push_back(std::string(s + start, i - start));
|
|
continue;
|
|
}
|
|
|
|
// Rule 6: whitespace handling
|
|
// Regex order: \s+(?!\S) first (trailing whitespace), then \s+ as fallback
|
|
// \s+(?!\S) backtracks: consumes whitespace NOT followed by non-whitespace
|
|
// This peels off leading spaces, leaving the last space to combine with the next word
|
|
if (is_whitespace(cp)) {
|
|
int start = i;
|
|
// Find end of whitespace run
|
|
int ws_end = i + adv;
|
|
while (ws_end < len && is_whitespace((unsigned char) s[ws_end]) && !is_newline((unsigned char) s[ws_end])) {
|
|
ws_end++;
|
|
}
|
|
// Check what follows the whitespace run
|
|
bool followed_by_non_ws =
|
|
(ws_end < len && !is_whitespace((unsigned char) s[ws_end]) && !is_newline((unsigned char) s[ws_end]));
|
|
if (followed_by_non_ws && ws_end - start > 1) {
|
|
// \s+(?!\S) matches all but the last space
|
|
// Leave one space for the next iteration to combine with word
|
|
int trailing = ws_end - 1;
|
|
chunks.push_back(std::string(s + start, trailing - start));
|
|
i = trailing;
|
|
continue;
|
|
}
|
|
// Single space followed by word: combine space + word as one chunk
|
|
i = start + adv;
|
|
if (i < len) {
|
|
int a2;
|
|
int cp2 = utf8_codepoint(s + i, &a2);
|
|
if (is_letter(cp2)) {
|
|
i += a2;
|
|
while (i < len) {
|
|
int a3;
|
|
int cp3 = utf8_codepoint(s + i, &a3);
|
|
if (!is_letter(cp3)) {
|
|
break;
|
|
}
|
|
i += a3;
|
|
}
|
|
chunks.push_back(std::string(s + start, i - start));
|
|
continue;
|
|
}
|
|
if (is_digit(cp2)) {
|
|
chunks.push_back(std::string(s + start, i - start));
|
|
continue;
|
|
}
|
|
if (!is_whitespace(cp2) && !is_newline(cp2)) {
|
|
int pstart = start;
|
|
while (i < len) {
|
|
int a3;
|
|
int cp3 = utf8_codepoint(s + i, &a3);
|
|
if (is_whitespace(cp3) || is_letter(cp3) || is_digit(cp3)) {
|
|
break;
|
|
}
|
|
i += a3;
|
|
}
|
|
while (i < len && is_newline((unsigned char) s[i])) {
|
|
i++;
|
|
}
|
|
chunks.push_back(std::string(s + pstart, i - pstart));
|
|
continue;
|
|
}
|
|
}
|
|
// Trailing whitespace (end of string or before newline), consume all
|
|
i = ws_end;
|
|
while (i < len) {
|
|
int a2;
|
|
int cp2 = utf8_codepoint(s + i, &a2);
|
|
if (!is_whitespace(cp2)) {
|
|
break;
|
|
}
|
|
i += a2;
|
|
}
|
|
chunks.push_back(std::string(s + start, i - start));
|
|
continue;
|
|
}
|
|
|
|
// Rule 4: [^\s\p{L}\p{N}]+[\r\n]* (punctuation/symbols)
|
|
{
|
|
int start = i;
|
|
i += adv;
|
|
while (i < len) {
|
|
int a2;
|
|
int cp2 = utf8_codepoint(s + i, &a2);
|
|
if (is_whitespace(cp2) || is_letter(cp2) || is_digit(cp2) || is_newline(cp2)) {
|
|
break;
|
|
}
|
|
i += a2;
|
|
}
|
|
// trailing newlines
|
|
while (i < len && is_newline((unsigned char) s[i])) {
|
|
i++;
|
|
}
|
|
chunks.push_back(std::string(s + start, i - start));
|
|
}
|
|
}
|
|
return chunks;
|
|
}
|
|
|
|
// BPE tokenizer struct
|
|
struct BPETokenizer {
|
|
std::unordered_map<std::string, int> vocab; // token_str -> id
|
|
std::unordered_map<std::string, int> merges; // "a b" -> rank
|
|
std::string byte2str[256]; // byte -> GPT-2 UTF-8 string
|
|
int eos_id; // <|endoftext|>
|
|
int n_vocab;
|
|
std::vector<std::string> id_to_str; // id -> token_str (reverse vocab)
|
|
|
|
// Registered special tokens. Each (str, id) pair is matched verbatim in
|
|
// bpe_encode and emitted as a single id, bypassing the BPE merge passes.
|
|
std::vector<std::pair<std::string, int>> specials;
|
|
};
|
|
|
|
// Register a special token. Strings already registered are skipped (no dup).
|
|
static void bpe_add_special(BPETokenizer * tok, const std::string & str, int id) {
|
|
for (const auto & sp : tok->specials) {
|
|
if (sp.first == str) {
|
|
return;
|
|
}
|
|
}
|
|
tok->specials.emplace_back(str, id);
|
|
}
|
|
|
|
// Load tokenizer from GGUF KV (tokenizer.ggml.tokens + tokenizer.ggml.merges)
|
|
static bool load_bpe_from_gguf(BPETokenizer * tok, const char * gguf_path) {
|
|
build_byte_encoder(tok->byte2str);
|
|
|
|
struct gguf_init_params gp = { true, NULL };
|
|
struct gguf_context * ctx = gguf_init_from_file(gguf_path, gp);
|
|
if (!ctx) {
|
|
fprintf(stderr, "[BPE] Failed to open %s\n", gguf_path);
|
|
return false;
|
|
}
|
|
|
|
int64_t tok_key = gguf_find_key(ctx, "tokenizer.ggml.tokens");
|
|
int64_t mrg_key = gguf_find_key(ctx, "tokenizer.ggml.merges");
|
|
if (tok_key < 0 || mrg_key < 0) {
|
|
fprintf(stderr, "[BPE] Tokenizer not found in %s\n", gguf_path);
|
|
gguf_free(ctx);
|
|
return false;
|
|
}
|
|
|
|
int n_tokens = (int) gguf_get_arr_n(ctx, tok_key);
|
|
int n_merges = (int) gguf_get_arr_n(ctx, mrg_key);
|
|
|
|
for (int i = 0; i < n_tokens; i++) {
|
|
const char * s = gguf_get_arr_str(ctx, tok_key, (size_t) i);
|
|
tok->vocab[std::string(s)] = i;
|
|
}
|
|
|
|
for (int i = 0; i < n_merges; i++) {
|
|
const char * s = gguf_get_arr_str(ctx, mrg_key, (size_t) i);
|
|
tok->merges[std::string(s)] = i;
|
|
}
|
|
|
|
gguf_free(ctx);
|
|
|
|
tok->n_vocab = (int) tok->vocab.size();
|
|
|
|
tok->id_to_str.resize(tok->n_vocab);
|
|
for (auto & kv : tok->vocab) {
|
|
if (kv.second >= 0 && kv.second < tok->n_vocab) {
|
|
tok->id_to_str[kv.second] = kv.first;
|
|
}
|
|
}
|
|
|
|
// Resolve eos_id from the vocab itself rather than hard-coding 151643.
|
|
// Falls back to -1 if the standard sentinel is absent.
|
|
auto eos_it = tok->vocab.find("<|endoftext|>");
|
|
tok->eos_id = (eos_it != tok->vocab.end()) ? eos_it->second : -1;
|
|
if (tok->eos_id >= 0) {
|
|
bpe_add_special(tok, "<|endoftext|>", tok->eos_id);
|
|
}
|
|
|
|
fprintf(stderr, "[BPE] Loaded from GGUF: %d vocab, %d merges, eos_id=%d\n", tok->n_vocab, n_merges, tok->eos_id);
|
|
return true;
|
|
}
|
|
|
|
// Read arch-specific special tokens from a caller-provided list of GGUF KV
|
|
// keys. Each key holds a u32 vocab id, mapped back to its vocab string and
|
|
// registered through bpe_add_special. The endoftext sentinel is already
|
|
// registered by load_bpe_from_gguf, so callers should not list it here.
|
|
static bool bpe_load_specials_from_keys(BPETokenizer * tok,
|
|
const char * gguf_path,
|
|
const char * const * keys,
|
|
int n_keys) {
|
|
struct gguf_init_params gp = { true, NULL };
|
|
struct gguf_context * ctx = gguf_init_from_file(gguf_path, gp);
|
|
if (!ctx) {
|
|
fprintf(stderr, "[BPE] Failed to open %s for specials\n", gguf_path);
|
|
return false;
|
|
}
|
|
|
|
int n_added = 0;
|
|
for (int i = 0; i < n_keys; i++) {
|
|
int64_t k = gguf_find_key(ctx, keys[i]);
|
|
if (k < 0) {
|
|
fprintf(stderr, "[BPE] WARNING: missing %s in GGUF\n", keys[i]);
|
|
continue;
|
|
}
|
|
int id = (int) gguf_get_val_u32(ctx, k);
|
|
if (id < 0 || id >= tok->n_vocab) {
|
|
fprintf(stderr, "[BPE] WARNING: %s id=%d out of vocab range\n", keys[i], id);
|
|
continue;
|
|
}
|
|
const std::string & s = tok->id_to_str[id];
|
|
if (s.empty()) {
|
|
fprintf(stderr, "[BPE] WARNING: %s id=%d has empty vocab string\n", keys[i], id);
|
|
continue;
|
|
}
|
|
bpe_add_special(tok, s, id);
|
|
n_added++;
|
|
}
|
|
|
|
gguf_free(ctx);
|
|
fprintf(stderr, "[BPE] Registered %d arch special tokens (total specials=%zu)\n", n_added, tok->specials.size());
|
|
return true;
|
|
}
|
|
|
|
// Byte-level encode: raw text bytes -> GPT-2 BPE string
|
|
static std::string byte_level_encode(const BPETokenizer * tok, const std::string & text) {
|
|
std::string out;
|
|
for (unsigned char c : text) {
|
|
out += tok->byte2str[c];
|
|
}
|
|
return out;
|
|
}
|
|
|
|
// BPE merge algorithm
|
|
// Input: list of symbols (strings). Merges pairs by priority.
|
|
static std::vector<std::string> bpe_merge(const std::unordered_map<std::string, int> & merge_rank,
|
|
const std::vector<std::string> & symbols) {
|
|
if (symbols.size() <= 1) {
|
|
return symbols;
|
|
}
|
|
|
|
std::vector<std::string> work = symbols;
|
|
|
|
while (work.size() > 1) {
|
|
// Find the pair with lowest rank (highest priority)
|
|
int best_rank = INT_MAX;
|
|
int best_pos = -1;
|
|
for (int i = 0; i < (int) work.size() - 1; i++) {
|
|
std::string key = work[i] + " " + work[i + 1];
|
|
auto it = merge_rank.find(key);
|
|
if (it != merge_rank.end() && it->second < best_rank) {
|
|
best_rank = it->second;
|
|
best_pos = i;
|
|
}
|
|
}
|
|
if (best_pos < 0) {
|
|
break; // no more merges
|
|
}
|
|
|
|
// Merge the pair
|
|
std::string merged = work[best_pos] + work[best_pos + 1];
|
|
work[best_pos] = merged;
|
|
work.erase(work.begin() + best_pos + 1);
|
|
}
|
|
return work;
|
|
}
|
|
|
|
// Encode a single pre-tokenized chunk -> token ids
|
|
static void encode_chunk(const BPETokenizer * tok, const std::string & chunk, std::vector<int> & ids) {
|
|
// Byte-level encode
|
|
std::string encoded = byte_level_encode(tok, chunk);
|
|
|
|
// Split into individual UTF-8 characters (each is a BPE symbol)
|
|
std::vector<std::string> symbols;
|
|
const char * s = encoded.c_str();
|
|
int len = (int) encoded.size();
|
|
int i = 0;
|
|
while (i < len) {
|
|
int adv;
|
|
utf8_codepoint(s + i, &adv);
|
|
symbols.push_back(std::string(s + i, adv));
|
|
i += adv;
|
|
}
|
|
|
|
// Apply BPE merges
|
|
std::vector<std::string> merged = bpe_merge(tok->merges, symbols);
|
|
|
|
// Look up in vocab
|
|
for (const auto & piece : merged) {
|
|
auto it = tok->vocab.find(piece);
|
|
if (it != tok->vocab.end()) {
|
|
ids.push_back(it->second);
|
|
} else {
|
|
// Fallback: encode each byte individually (should not happen with byte-level BPE)
|
|
fprintf(stderr, "[BPE] WARNING: unknown token '%s'\n", piece.c_str());
|
|
for (unsigned char c : piece) {
|
|
auto it2 = tok->vocab.find(std::string(1, c));
|
|
if (it2 != tok->vocab.end()) {
|
|
ids.push_back(it2->second);
|
|
}
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
// Full encode: text -> token ids.
|
|
// Walks the text from left to right, matching any registered special token
|
|
// verbatim. For each segment between specials, runs the GPT-2 byte-level
|
|
// pre-tokenizer + BPE merges. The endoftext sentinel is auto-registered as
|
|
// a special by load_bpe_from_gguf, so existing call sites that embed
|
|
// "<|endoftext|>" in the input text keep working.
|
|
// add_eos = true appends the eos_id at the end (post-processor behavior).
|
|
static std::vector<int> bpe_encode(const BPETokenizer * tok, const std::string & text, bool add_eos = true) {
|
|
std::vector<int> ids;
|
|
|
|
auto encode_segment = [&](const std::string & seg) {
|
|
if (seg.empty()) {
|
|
return;
|
|
}
|
|
auto chunks = gpt2_pre_tokenize(seg);
|
|
for (const auto & chunk : chunks) {
|
|
encode_chunk(tok, chunk, ids);
|
|
}
|
|
};
|
|
|
|
size_t pos = 0;
|
|
while (pos < text.size()) {
|
|
// Find the leftmost occurrence of any registered special token.
|
|
size_t best_pos = std::string::npos;
|
|
int best_idx = -1;
|
|
for (size_t i = 0; i < tok->specials.size(); i++) {
|
|
size_t p = text.find(tok->specials[i].first, pos);
|
|
if (p != std::string::npos && p < best_pos) {
|
|
best_pos = p;
|
|
best_idx = (int) i;
|
|
}
|
|
}
|
|
|
|
if (best_idx < 0) {
|
|
encode_segment(text.substr(pos));
|
|
break;
|
|
}
|
|
|
|
if (best_pos > pos) {
|
|
encode_segment(text.substr(pos, best_pos - pos));
|
|
}
|
|
const auto & sp = tok->specials[(size_t) best_idx];
|
|
ids.push_back(sp.second);
|
|
pos = best_pos + sp.first.size();
|
|
}
|
|
|
|
if (add_eos && tok->eos_id >= 0) {
|
|
ids.push_back(tok->eos_id);
|
|
}
|
|
return ids;
|
|
}
|