38 lines
1.2 KiB
Python
38 lines
1.2 KiB
Python
import mlx.core as mx
|
|
import numpy as np
|
|
from transformers import PreTrainedTokenizer
|
|
|
|
|
|
class TokenizerFibo:
|
|
def __init__(self, tokenizer: PreTrainedTokenizer, bot_token_id: int = 128000):
|
|
self.tokenizer = tokenizer
|
|
self.bot_token_id = bot_token_id
|
|
|
|
def tokenize(
|
|
self,
|
|
prompts: list[str],
|
|
max_length: int = 2048,
|
|
padding: str = "longest",
|
|
truncation: bool = True,
|
|
add_special_tokens: bool = True,
|
|
) -> tuple[mx.array, mx.array]:
|
|
prompts = [p if p is not None else "" for p in prompts]
|
|
|
|
if all(p == "" for p in prompts):
|
|
batch_size = len(prompts)
|
|
input_ids_mx = mx.array(np.empty((batch_size, 0), dtype=np.int32))
|
|
attention_mask_mx = mx.array(np.empty((batch_size, 0), dtype=np.int32))
|
|
else:
|
|
tokenized = self.tokenizer(
|
|
prompts,
|
|
padding=padding,
|
|
max_length=max_length,
|
|
truncation=truncation,
|
|
add_special_tokens=add_special_tokens,
|
|
return_tensors="mlx",
|
|
)
|
|
input_ids_mx = tokenized["input_ids"]
|
|
attention_mask_mx = tokenized["attention_mask"]
|
|
|
|
return input_ids_mx, attention_mask_mx
|