forked from mistralai/mistral-inference
-
Notifications
You must be signed in to change notification settings - Fork 0
/
tokenizer.py
36 lines (28 loc) · 954 Bytes
/
tokenizer.py
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
from pathlib import Path
from sentencepiece import SentencePieceProcessor
from typing import List
class Tokenizer:
def __init__(self, model_path: str):
assert Path(model_path).exists(), model_path
self._model = SentencePieceProcessor(model_file=model_path)
assert self._model.vocab_size() == self._model.get_piece_size()
@property
def n_words(self) -> int:
return self._model.vocab_size()
@property
def bos_id(self) -> int:
return self._model.bos_id()
@property
def eos_id(self) -> int:
return self._model.eos_id()
@property
def pad_id(self) -> int:
return self._model.pad_id()
def encode(self, s: str, bos: bool = True) -> List[int]:
assert isinstance(s, str)
t = self._model.encode(s)
if bos:
t = [self.bos_id, *t]
return t
def decode(self, t: List[int]) -> str:
return self._model.decode(t)