Skip to content

Commit d4c8e2c

Browse files
o7siCISC
andauthored
vocab : add tokenizer support for jina-embeddings-v2-base-zh (ggml-org#18756)
* vocab : add jina-embeddings-v2-base-zh (whitespace tokenizer) * lowercase defaults to true * type fix --------- Co-authored-by: Sigbjørn Skjæret <sigbjorn.skjaeret@scala.com>
1 parent 3292da0 commit d4c8e2c

9 files changed

Lines changed: 106 additions & 4 deletions

File tree

conversion/base.py

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1692,6 +1692,16 @@ def _set_vocab_gpt2(self) -> None:
16921692
special_vocab = gguf.SpecialVocab(self.dir_model, load_merges=True)
16931693
special_vocab.add_to_gguf(self.gguf_writer)
16941694

1695+
def _set_vocab_whitespace(self) -> None:
1696+
tokens, toktypes, _ = self.get_vocab_base()
1697+
self.gguf_writer.add_tokenizer_model("whitespace")
1698+
self.gguf_writer.add_tokenizer_pre("whitespace") # pinned, not hash-detected: chktxt hash collides with jina-v1-en
1699+
self.gguf_writer.add_token_list(tokens)
1700+
self.gguf_writer.add_token_types(toktypes)
1701+
1702+
special_vocab = gguf.SpecialVocab(self.dir_model, load_merges=True)
1703+
special_vocab.add_to_gguf(self.gguf_writer)
1704+
16951705
def _set_vocab_hybriddna(self):
16961706
from transformers import AutoTokenizer
16971707
tokenizer = AutoTokenizer.from_pretrained(self.dir_model, trust_remote_code=True)

conversion/bert.py

Lines changed: 10 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -571,7 +571,16 @@ def set_vocab(self):
571571
if tokenizer_class == 'BertTokenizer':
572572
super().set_vocab()
573573
elif tokenizer_class == 'RobertaTokenizer':
574-
self._set_vocab_gpt2()
574+
pre_tokenizer_type = None
575+
tokenizer_json_path = self.dir_model / "tokenizer.json"
576+
if tokenizer_json_path.is_file():
577+
with open(tokenizer_json_path, "r", encoding="utf-8") as f:
578+
pre_tokenizer_type = json.load(f).get("pre_tokenizer", {}).get("type")
579+
580+
if pre_tokenizer_type == "Whitespace":
581+
self._set_vocab_whitespace()
582+
else:
583+
self._set_vocab_gpt2()
575584
self.gguf_writer.add_token_type_count(2)
576585
else:
577586
raise NotImplementedError(f'Tokenizer {tokenizer_class} is not supported for JinaBertModel')

gguf-py/gguf/constants.py

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -268,6 +268,8 @@ class Tokenizer:
268268
CHAT_TEMPLATE = "tokenizer.chat_template"
269269
CHAT_TEMPLATE_N = "tokenizer.chat_template.{name}"
270270
CHAT_TEMPLATES = "tokenizer.chat_templates"
271+
# Normalizer constants
272+
NORMALIZER_LOWERCASE = "tokenizer.ggml.normalizer.lowercase"
271273
# FIM/Infill special tokens constants
272274
FIM_PRE_ID = "tokenizer.ggml.fim_pre_token_id"
273275
FIM_SUF_ID = "tokenizer.ggml.fim_suf_token_id"

gguf-py/gguf/gguf_writer.py

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1110,6 +1110,9 @@ def add_chat_template(self, value: str | Sequence[Mapping[str, str]]) -> None:
11101110

11111111
self.add_string(Keys.Tokenizer.CHAT_TEMPLATE, value)
11121112

1113+
def add_normalizer_lowercase(self, value: bool) -> None:
1114+
self.add_bool(Keys.Tokenizer.NORMALIZER_LOWERCASE, value)
1115+
11131116
def add_eot_token_id(self, id: int) -> None:
11141117
self.add_uint32(Keys.Tokenizer.EOT_ID, id)
11151118

gguf-py/gguf/vocab.py

Lines changed: 27 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -52,6 +52,7 @@ class SpecialVocab:
5252
add_special_token: dict[str, bool]
5353
special_token_ids: dict[str, int]
5454
chat_template: str | Sequence[Mapping[str, str]] | None
55+
normalizer_lowercase: bool | None
5556

5657
def __init__(
5758
self, path: str | os.PathLike[str], load_merges: bool = False,
@@ -64,6 +65,7 @@ def __init__(
6465
self.load_merges = load_merges
6566
self.merges = []
6667
self.chat_template = None
68+
self.normalizer_lowercase = None
6769
if special_token_types is not None:
6870
self.special_token_types = special_token_types
6971
else:
@@ -102,6 +104,10 @@ def add_to_gguf(self, gw: GGUFWriter, quiet: bool = False) -> None:
102104
if not quiet:
103105
logger.info(f'Setting chat_template to {self.chat_template}')
104106
gw.add_chat_template(self.chat_template)
107+
if self.normalizer_lowercase is not None:
108+
if not quiet:
109+
logger.info(f'Setting normalizer_lowercase to {self.normalizer_lowercase}')
110+
gw.add_normalizer_lowercase(self.normalizer_lowercase)
105111

106112
def _load(self, path: Path) -> None:
107113
self._try_load_from_tokenizer_json(path)
@@ -146,6 +152,24 @@ def _set_special_token(self, typ: str, tid: Any) -> None:
146152
return
147153
logger.warning(f'Special token type {typ}, id {tid} out of range, must be under {self.n_vocab} - skipping')
148154

155+
def _parse_normalizer(self, normalizer: dict) -> None:
156+
# ref: https://huggingface.co/docs/tokenizers/api/normalizers
157+
#
158+
# Detects lowercase normalization in three possible formats:
159+
# 1. Standalone: {"type": "Lowercase"}
160+
# 2. BertNormalizer attribute: {"type": "BertNormalizer", "lowercase": true, ...}
161+
# 3. Nested in Sequence: {"type": "Sequence", "normalizers": [...]}
162+
163+
normalizer_type = normalizer.get('type')
164+
if normalizer_type == 'Lowercase':
165+
self.normalizer_lowercase = True
166+
elif normalizer_type == 'BertNormalizer':
167+
if 'lowercase' in normalizer:
168+
self.normalizer_lowercase = normalizer['lowercase']
169+
elif normalizer_type == 'Sequence':
170+
for norm in normalizer.get('normalizers', []):
171+
self._parse_normalizer(norm)
172+
149173
def _try_load_from_tokenizer_json(self, path: Path) -> bool:
150174
tokenizer = None
151175
tokenizer_file = path / 'tokenizer.json'
@@ -178,6 +202,9 @@ def _try_load_from_tokenizer_json(self, path: Path) -> bool:
178202
]
179203
else:
180204
raise ValueError("Unknown tokenizer merges format")
205+
# Parse normalizer configuration (e.g. Lowercase) into metadata
206+
if normalizer := tokenizer.get('normalizer'):
207+
self._parse_normalizer(normalizer)
181208
added_tokens = tokenizer.get('added_tokens', {})
182209
else:
183210
added_tokens = {}

src/llama-arch.cpp

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -319,6 +319,7 @@ static const std::map<llm_kv, const char *> LLM_KV_NAMES = {
319319
{ LLM_KV_TOKENIZER_HF_JSON, "tokenizer.huggingface.json" },
320320
{ LLM_KV_TOKENIZER_RWKV, "tokenizer.rwkv.world" },
321321
{ LLM_KV_TOKENIZER_CHAT_TEMPLATE, "tokenizer.chat_template" },
322+
{ LLM_KV_TOKENIZER_NORMALIZER_LOWERCASE, "tokenizer.ggml.normalizer.lowercase" },
322323
{ LLM_KV_TOKENIZER_FIM_PRE_ID, "tokenizer.ggml.fim_pre_token_id" },
323324
{ LLM_KV_TOKENIZER_FIM_SUF_ID, "tokenizer.ggml.fim_suf_token_id" },
324325
{ LLM_KV_TOKENIZER_FIM_MID_ID, "tokenizer.ggml.fim_mid_token_id" },

src/llama-arch.h

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -308,6 +308,7 @@ enum llm_kv {
308308
LLM_KV_TOKENIZER_HF_JSON,
309309
LLM_KV_TOKENIZER_RWKV,
310310
LLM_KV_TOKENIZER_CHAT_TEMPLATE,
311+
LLM_KV_TOKENIZER_NORMALIZER_LOWERCASE,
311312
LLM_KV_TOKENIZER_FIM_PRE_ID,
312313
LLM_KV_TOKENIZER_FIM_SUF_ID,
313314
LLM_KV_TOKENIZER_FIM_MID_ID,

src/llama-vocab.cpp

Lines changed: 50 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -519,6 +519,13 @@ struct llm_tokenizer_bpe : llm_tokenizer {
519519
"(?:'[sS]|'[tT]|'[rR][eE]|'[vV][eE]|'[mM]|'[lL][lL]|'[dD])|[^\\r\\n\\p{L}\\p{N}]?\\p{L}+|\\p{N}+| ?[^\\s\\p{L}\\p{N}]+[\\r\\n]*|\\s*[\\r\\n]+|\\s+(?!\\S)|\\s+",
520520
};
521521
break;
522+
case LLAMA_VOCAB_PRE_TYPE_WHITESPACE:
523+
// whitespace pre-tokenizer (jinaai/jina-embeddings-v2-base-zh)
524+
regex_exprs = {
525+
"\\S+",
526+
};
527+
byte_encode = false;
528+
break;
522529
default:
523530
// default regex for BPE tokenization pre-processing
524531
regex_exprs = {
@@ -1671,6 +1678,35 @@ struct llm_tokenizer_hybriddna_session : llm_tokenizer_bpe_session {
16711678
const llama_vocab & vocab;
16721679
};
16731680

1681+
struct llm_tokenizer_whitespace_session : llm_tokenizer_bpe_session {
1682+
llm_tokenizer_whitespace_session(const llama_vocab & vocab, const llm_tokenizer_bpe & tokenizer) : llm_tokenizer_bpe_session{vocab, tokenizer}, vocab{vocab} {}
1683+
1684+
void tokenize(const std::string & text, std::vector<llama_token> & output) override {
1685+
const bool lowercase = vocab.get_normalizer_lowercase();
1686+
1687+
std::string segment;
1688+
auto flush = [&]() {
1689+
if (!segment.empty()) {
1690+
llm_tokenizer_bpe_session::tokenize(segment, output);
1691+
segment.clear();
1692+
}
1693+
};
1694+
1695+
for (uint32_t cpt : unicode_cpts_from_utf8(text)) {
1696+
// drop whitespace
1697+
if (unicode_cpt_flags_from_cpt(cpt).is_whitespace) {
1698+
flush();
1699+
} else {
1700+
segment += unicode_cpt_to_utf8(lowercase ? unicode_tolower(cpt) : cpt);
1701+
}
1702+
}
1703+
flush();
1704+
}
1705+
1706+
private:
1707+
const llama_vocab & vocab;
1708+
};
1709+
16741710
//
16751711
// impl
16761712
//
@@ -1751,6 +1787,7 @@ struct llama_vocab::impl {
17511787
bool remove_extra_whitespaces = false;
17521788
bool escape_whitespaces = true;
17531789
bool treat_whitespace_as_suffix = false;
1790+
bool normalizer_lowercase = true; // Lowercase normalizer (tokenizer.json)
17541791

17551792
std::unordered_map<std::string, llama_token> token_to_id;
17561793
std::vector<token_data> id_to_token;
@@ -1900,7 +1937,7 @@ void llama_vocab::impl::load(llama_model_loader & ml, const LLM_KV & kv) {
19001937
special_mask_id = 103;
19011938

19021939
add_sep = true;
1903-
} else if (tokenizer_model == "gpt2" || tokenizer_model == "hybriddna") {
1940+
} else if (tokenizer_model == "gpt2" || tokenizer_model == "hybriddna" || tokenizer_model == "whitespace") {
19041941
type = LLAMA_VOCAB_TYPE_BPE;
19051942

19061943
// read bpe merges and populate bpe ranks
@@ -2119,6 +2156,9 @@ void llama_vocab::impl::load(llama_model_loader & ml, const LLM_KV & kv) {
21192156
tokenizer_pre == "roberta-bpe") {
21202157
pre_type = LLAMA_VOCAB_PRE_TYPE_GPT2;
21212158
add_sep = true;
2159+
} else if (
2160+
tokenizer_pre == "whitespace") {
2161+
pre_type = LLAMA_VOCAB_PRE_TYPE_WHITESPACE;
21222162
} else if (
21232163
tokenizer_pre == "refact") {
21242164
pre_type = LLAMA_VOCAB_PRE_TYPE_REFACT;
@@ -2299,8 +2339,9 @@ void llama_vocab::impl::load(llama_model_loader & ml, const LLM_KV & kv) {
22992339
pre_type = LLAMA_VOCAB_PRE_TYPE_DEFAULT;
23002340
}
23012341

2302-
ml.get_key(LLM_KV_TOKENIZER_ADD_PREFIX, add_space_prefix, false);
2303-
ml.get_key(LLM_KV_TOKENIZER_REMOVE_EXTRA_WS, remove_extra_whitespaces, false);
2342+
ml.get_key(LLM_KV_TOKENIZER_ADD_PREFIX, add_space_prefix, false);
2343+
ml.get_key(LLM_KV_TOKENIZER_REMOVE_EXTRA_WS, remove_extra_whitespaces, false);
2344+
ml.get_key(LLM_KV_TOKENIZER_NORMALIZER_LOWERCASE, normalizer_lowercase, false);
23042345
}
23052346

23062347
const int token_idx = gguf_find_key(ctx, kv(LLM_KV_TOKENIZER_LIST).c_str());
@@ -3264,6 +3305,8 @@ std::vector<llama_token> llama_vocab::impl::tokenize(
32643305
std::unique_ptr<llm_tokenizer_bpe_session> session;
32653306
if (vocab.get_tokenizer_model() == "hybriddna") {
32663307
session = std::make_unique<llm_tokenizer_hybriddna_session>(vocab, *tok_bpe);
3308+
} else if (vocab.get_tokenizer_model() == "whitespace") {
3309+
session = std::make_unique<llm_tokenizer_whitespace_session>(vocab, *tok_bpe);
32673310
} else {
32683311
session = std::make_unique<llm_tokenizer_bpe_session>(vocab, *tok_bpe);
32693312
}
@@ -3892,6 +3935,10 @@ bool llama_vocab::get_treat_whitespace_as_suffix() const {
38923935
return pimpl->treat_whitespace_as_suffix;
38933936
}
38943937

3938+
bool llama_vocab::get_normalizer_lowercase() const {
3939+
return pimpl->normalizer_lowercase;
3940+
}
3941+
38953942
int llama_vocab::max_token_len() const {
38963943
return pimpl->max_token_len;
38973944
}

src/llama-vocab.h

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -61,6 +61,7 @@ enum llama_vocab_pre_type {
6161
LLAMA_VOCAB_PRE_TYPE_GEMMA4 = 50,
6262
LLAMA_VOCAB_PRE_TYPE_SARVAM_MOE = 51,
6363
LLAMA_VOCAB_PRE_TYPE_MINICPM5 = 52,
64+
LLAMA_VOCAB_PRE_TYPE_WHITESPACE = 53,
6465
};
6566

6667
struct LLM_KV;
@@ -138,6 +139,7 @@ struct llama_vocab {
138139
bool get_remove_extra_whitespaces () const;
139140
bool get_escape_whitespaces () const;
140141
bool get_treat_whitespace_as_suffix() const;
142+
bool get_normalizer_lowercase () const;
141143

142144
int max_token_len() const;
143145

0 commit comments

Comments
 (0)