|
|
|
|
|
|
|
|
|
|
|
import os |
|
|
from logging import getLogger |
|
|
from typing import List |
|
|
|
|
|
from sentencepiece import SentencePieceProcessor |
|
|
|
|
|
|
|
|
logger = getLogger() |
|
|
|
|
|
|
|
|
class TextTokenizer: |
|
|
"""Tokenizing and encoding/decoding text using SentencePiece.""" |
|
|
|
|
|
def __init__(self, model_path=None): |
|
|
""" |
|
|
Initializes the Tokenizer with a SentencePiece model. |
|
|
|
|
|
Args: |
|
|
model_path (str): The path to the SentencePiece model file. |
|
|
""" |
|
|
if model_path is None: |
|
|
model_path = os.path.join( |
|
|
os.path.dirname(os.path.abspath(__file__)), "text_tokenizer.model" |
|
|
) |
|
|
|
|
|
assert os.path.isfile(model_path), model_path |
|
|
self.sp_model = SentencePieceProcessor(model_file=model_path) |
|
|
logger.info(f"Reloaded SentencePiece model from {model_path}") |
|
|
|
|
|
self.n_words: int = self.sp_model.vocab_size() |
|
|
self.bos_id: int = self.sp_model.bos_id() |
|
|
self.eos_id: int = self.sp_model.eos_id() |
|
|
self.pad_id: int = self.sp_model.pad_id() |
|
|
self.pad_id += self.n_words if self.pad_id < 0 else 0 |
|
|
logger.info(f"#words: {self.n_words} - BOS ID: {self.bos_id} - EOS ID: {self.eos_id}") |
|
|
assert self.sp_model.vocab_size() == self.sp_model.get_piece_size() |
|
|
|
|
|
def encode(self, s: str, bos: bool, eos: bool) -> List[int]: |
|
|
""" |
|
|
Encodes a string into a list of token IDs. |
|
|
|
|
|
Args: |
|
|
s (str): The input string to be encoded. |
|
|
bos (bool): Whether to prepend the beginning-of-sequence token. |
|
|
eos (bool): Whether to append the end-of-sequence token. |
|
|
|
|
|
Returns: |
|
|
List[int]: A list of token IDs. |
|
|
""" |
|
|
assert type(s) is str |
|
|
t = self.sp_model.encode(s) |
|
|
if bos: |
|
|
t = [self.bos_id] + t |
|
|
if eos: |
|
|
t = t + [self.eos_id] |
|
|
return t |
|
|
|
|
|
def decode(self, t: List[int]) -> str: |
|
|
""" |
|
|
Decodes a list of token IDs into a string. |
|
|
|
|
|
Args: |
|
|
t (List[int]): The list of token IDs to be decoded. |
|
|
|
|
|
Returns: |
|
|
str: The decoded string. |
|
|
""" |
|
|
return self.sp_model.decode(t) |
|
|
|
|
|
def tokenize(self, texts, context_length=None): |
|
|
"""Encode a list of string. |
|
|
|
|
|
Parameters |
|
|
---------- |
|
|
texts : Union[str, List[str]] |
|
|
The input text(s). |
|
|
context_length : int, optional |
|
|
The max token length. |
|
|
|
|
|
Returns |
|
|
------- |
|
|
List[List[int]] |
|
|
The encoded token indices. |
|
|
|
|
|
""" |
|
|
if isinstance(texts, str): |
|
|
texts = [texts] |
|
|
tokens = [self.encode(text, bos=True, eos=True) for text in texts] |
|
|
if context_length is None: |
|
|
return tokens |
|
|
truncated_tokens = [] |
|
|
for k, t in enumerate(tokens): |
|
|
if len(t) > context_length: |
|
|
t = t[:context_length] |
|
|
t[-1] = self.eos_id |
|
|
truncated_tokens.append(t) |
|
|
return truncated_tokens |
|
|
|
|
|
def detokenize(self, tokens): |
|
|
"""Decode a list of string. |
|
|
|
|
|
Parameters |
|
|
---------- |
|
|
tokens : Union[List[List[int]], numpy.ndarray] |
|
|
The input tokens. |
|
|
|
|
|
Returns |
|
|
------- |
|
|
List[str] |
|
|
The decoded text strings. |
|
|
|
|
|
""" |
|
|
if hasattr(tokens, "tolist"): |
|
|
tokens = tokens.tolist() |
|
|
texts = [] |
|
|
for i in range(len(tokens)): |
|
|
t = tokens[i][1:] |
|
|
try: |
|
|
eot_idx = t.index(self.eos_id) |
|
|
t = t[:eot_idx] |
|
|
except ValueError: |
|
|
pass |
|
|
texts.append(self.decode(t)) |
|
|
return texts |
|
|
|
