|
|
--- |
|
|
library_name: keras-hub |
|
|
--- |
|
|
|
|
|
```py |
|
|
import tensorflow as tf |
|
|
from tokenizers import Tokenizer as HFTokenizer # pip install tokenizers |
|
|
from keras_hub.tokenizers import Tokenizer as KerasTokenizerBase |
|
|
|
|
|
class HFRustTokenizerWrapper(KerasTokenizerBase): |
|
|
def __init__(self, hf_tokenizer): |
|
|
super().__init__() |
|
|
""" |
|
|
hf_tokenizer: either a tokenizers.Tokenizer instance (recommended) |
|
|
or a path to a tokenizer.json that Tokenizer.from_file can load. |
|
|
""" |
|
|
# 如果传入的是路径字符串，就从文件加载 |
|
|
if isinstance(hf_tokenizer, str): |
|
|
self.tk = HFTokenizer.from_file(hf_tokenizer) |
|
|
else: |
|
|
# 假设是已经构造好的 tokenizers.Tokenizer |
|
|
self.tk = hf_tokenizer |
|
|
|
|
|
self._dtype = "int32" |
|
|
|
|
|
def tokenize(self, inputs): |
|
|
""" |
|
|
inputs: tf.Tensor(dtype=string), shape [batch] or scalar |
|
|
return: tf.RaggedTensor(dtype=int32) |
|
|
""" |
|
|
inputs = tf.convert_to_tensor(inputs, dtype=tf.string) |
|
|
|
|
|
def _py_tokenize(x): |
|
|
# x: tf.Tensor[string] in eager context (inside py_function) |
|
|
arr = x.numpy() |
|
|
texts = [ |
|
|
s.decode("utf-8") if isinstance(s, (bytes, bytearray)) else str(s) |
|
|
for s in arr |
|
|
] |
|
|
|
|
|
encs = self.tk.encode_batch(texts, add_special_tokens=False) |
|
|
ids = [enc.ids for enc in encs] |
|
|
|
|
|
# 返回 RaggedTensor 的 components |
|
|
return tf.ragged.constant(ids, dtype=tf.int32) |
|
|
|
|
|
# tf.py_function 只能返回 Tensor / CompositeTensor |
|
|
ragged = tf.py_function( |
|
|
func=_py_tokenize, |
|
|
inp=[inputs], |
|
|
Tout=tf.RaggedTensorSpec( |
|
|
shape=[None, None], |
|
|
dtype=tf.int32, |
|
|
ragged_rank=1, |
|
|
), |
|
|
) |
|
|
|
|
|
# 修正 static shape（否则下游有时会 complain） |
|
|
#ragged.set_shape([None, None]) |
|
|
return ragged |
|
|
|
|
|
def detokenize(self, inputs): |
|
|
""" |
|
|
inputs: RaggedTensor / Tensor / list |
|
|
返回: tf.Tensor(dtype string) — batch of decoded strings, or scalar if single input |
|
|
""" |
|
|
# 规范化为 python list[list[int]] |
|
|
if isinstance(inputs, tf.RaggedTensor): |
|
|
ids_list = inputs.to_list() |
|
|
elif isinstance(inputs, tf.Tensor): |
|
|
# 可能是 [batch, seq] 的定长 tensor |
|
|
ids_list = inputs.numpy().tolist() |
|
|
else: |
|
|
ids_list = inputs |
|
|
|
|
|
# 如果传入的是单条 ids (like [1,2,3]), wrap 成 batch |
|
|
if ids_list and isinstance(ids_list[0], int): |
|
|
ids_list = [ids_list] |
|
|
|
|
|
texts = [] |
|
|
for ids in ids_list: |
|
|
# tokenizers.Tokenizer 提供 decode(ids) |
|
|
# 有些 tokenizer 实现有 decode_batch，但使用循环以兼容更多版本 |
|
|
texts.append(self.tk.decode(ids, skip_special_tokens=True)) |
|
|
|
|
|
# 如果原来是单条输入，返回 scalar string tensor 与原行为更接近 |
|
|
if len(texts) == 1: |
|
|
return tf.convert_to_tensor(texts[0]) |
|
|
return tf.convert_to_tensor(texts) |
|
|
|
|
|
def vocabulary_size(self): |
|
|
# Tokenizers API 提供 get_vocab_size() 或 len(self.tk.get_vocab()) |
|
|
try: |
|
|
return self.tk.get_vocab_size() |
|
|
except Exception: |
|
|
# 兜底 |
|
|
try: |
|
|
return len(self.tk.get_vocab()) |
|
|
except Exception: |
|
|
# 如果都不可用，返回 0 |
|
|
return 0 |
|
|
|
|
|
def id_to_token(self, id_): |
|
|
try: |
|
|
return self.tk.id_to_token(id_) |
|
|
except Exception: |
|
|
# 有些版本的 API 叫 token_to_id 的反向，需要手动查 vocab |
|
|
try: |
|
|
inv = {v: k for k, v in self.tk.get_vocab().items()} |
|
|
return inv.get(int(id_), "") |
|
|
except Exception: |
|
|
return "" |
|
|
|
|
|
def token_to_id(self, token): |
|
|
try: |
|
|
return self.tk.token_to_id(token) |
|
|
except Exception: |
|
|
try: |
|
|
return self.tk.get_vocab().get(token, None) |
|
|
except Exception: |
|
|
return None |
|
|
|
|
|
@property |
|
|
def dtype(self): |
|
|
return tf.int32 |
|
|
|
|
|
from huggingface_hub import hf_hub_download |
|
|
from tokenizers import Tokenizer |
|
|
|
|
|
tokenizer_path = hf_hub_download( |
|
|
repo_id="Qwen/Qwen3-4B-Base", |
|
|
filename="tokenizer.json", |
|
|
) |
|
|
|
|
|
hf_tokenizer = Tokenizer.from_file(tokenizer_path) |
|
|
wrapper = HFRustTokenizerWrapper(hf_tokenizer) |
|
|
wrapper.start_token_id = 151643 # endoftext |
|
|
wrapper.end_token_id = 151643 |
|
|
wrapper.pad_token_id = 151643 |
|
|
gemma_lm.preprocessor.tokenizer = wrapper |
|
|
gemma_lm.preprocessor.add_end_token = True |
|
|
gemma_lm.preprocessor.add_start_token = False |
|
|
``` |
