| |
| |
| |
| |
| |
| |
|
|
| import json |
| import os |
| from typing import List, Tuple |
| import numpy as np |
|
|
|
|
| class CharTokenizer: |
| """Character/word tokenizer for the sherpa punct CT Transformer model.""" |
|
|
| def __init__(self, tokens_path: str, unk_symbol: str = "<unk>"): |
| if not os.path.exists(tokens_path): |
| raise FileNotFoundError(f"tokens.json not found: {tokens_path}") |
| with open(tokens_path, "r", encoding="utf-8") as f: |
| id2token = json.load(f) |
| self.id2token = id2token |
| self.token2id = {tok: idx for idx, tok in enumerate(id2token)} |
| self.unk_id = self.token2id.get(unk_symbol, 0) |
|
|
| def tokenize(self, text: str) -> List[int]: |
| """Split text into tokens and return token IDs. |
| |
| Chinese characters are segmented individually. |
| English words are kept as whole tokens. |
| """ |
| |
| word_list = text.split() |
|
|
| words = [] |
| for w in word_list: |
| s = "" |
| for c in w: |
| if len(c.encode()) > 1: |
| |
| if s == "": |
| s = c |
| elif len(s[-1].encode()) > 1: |
| s += c |
| else: |
| words.append(s) |
| s = c |
| else: |
| |
| if s == "": |
| s = c |
| elif len(s[-1].encode()) > 1: |
| words.append(s) |
| s = c |
| else: |
| s += c |
| if s: |
| words.append(s) |
|
|
| ids = [] |
| for w in words: |
| if len(w[0].encode()) > 1: |
| |
| for c in w: |
| ids.append(self.token2id.get(c, self.unk_id)) |
| else: |
| ids.append(self.token2id.get(w, self.unk_id)) |
| return ids |
|
|
| def tokenize_full(self, text: str) -> List[int]: |
| """Tokenize full text without truncation or padding.""" |
| return self.tokenize(text) |
|
|
| def encode( |
| self, text: str, pad_length: int = 64 |
| ) -> Tuple[np.ndarray, int]: |
| """Tokenize and pad to fixed length. Truncates if > pad_length. |
| |
| Returns: |
| input_array: (1, pad_length) int32 numpy array |
| original_length: actual token count before padding |
| """ |
| ids = self.tokenize(text) |
| original_len = len(ids) |
|
|
| |
| if len(ids) > pad_length: |
| ids = ids[:pad_length] |
| original_len = pad_length |
|
|
| padded = np.zeros((1, pad_length), dtype=np.int32) |
| padded[0, : len(ids)] = ids |
|
|
| return padded, min(original_len, pad_length) |
|
|
| def encode_long( |
| self, text: str, window_size: int = 64 |
| ) -> Tuple[List[np.ndarray], List[int], List[int]]: |
| """Tokenize long text into sliding windows for batched inference. |
| |
| Splits full token sequence into windows of window_size. |
| Each window is padded to window_size if shorter. |
| |
| Returns: |
| windows: list of (1, window_size) int32 arrays |
| window_token_ids: list of token ID lists per window |
| window_lens: original token lengths per window (before padding) |
| """ |
| ids = self.tokenize(text) |
| if not ids: |
| return [], [], [] |
|
|
| windows = [] |
| window_token_ids = [] |
| window_lens = [] |
|
|
| for start in range(0, len(ids), window_size): |
| chunk = ids[start:start + window_size] |
| chunk_len = len(chunk) |
|
|
| padded = np.zeros((1, window_size), dtype=np.int32) |
| padded[0, :chunk_len] = chunk |
|
|
| windows.append(padded) |
| window_token_ids.append(chunk) |
| window_lens.append(chunk_len) |
|
|
| return windows, window_token_ids, window_lens |
|
|
