| from pathlib import Path
|
| import json
|
| import heapq
|
|
|
|
|
| class SutraTokenizerFast:
|
|
|
| def __init__(self, directory):
|
|
|
| directory = Path(directory)
|
|
|
| with open(directory / "vocab.json", "r", encoding="utf-8") as f:
|
| self.vocab = json.load(f)
|
|
|
| with open(directory / "merges.json", "r", encoding="utf-8") as f:
|
| self.merges = [tuple(pair) for pair in json.load(f)]
|
|
|
| self.token_to_id = self.vocab
|
|
|
|
|
| self.id_to_token = {
|
| int(idx): token
|
| for token, idx in self.vocab.items()
|
| }
|
|
|
|
|
| self.unk_id = self.vocab.get("<|unk|>")
|
| self.pad_token_id = self.vocab.get("<|pad|>")
|
| self.bos_token_id = self.vocab.get("<|bos|>")
|
| self.eos_token_id = self.vocab.get("<|eos|>")
|
| self.unk_token_id = self.unk_id
|
|
|
|
|
| self.special_token_ids = {
|
| self.pad_token_id,
|
| self.bos_token_id,
|
| self.eos_token_id,
|
| self.unk_token_id,
|
| } - {None}
|
|
|
| self.merge_rank = {
|
| pair: rank
|
| for rank, pair in enumerate(self.merges)
|
| }
|
|
|
| def encode(self, text):
|
|
|
|
|
|
|
|
|
|
|
| text = text.replace(" ", "▁")
|
|
|
| symbols = list(text)
|
| n = len(symbols)
|
|
|
| if n == 0:
|
| return []
|
|
|
| if n == 1:
|
| return [
|
| self.token_to_id.get(
|
| symbols[0],
|
| self.unk_id
|
| )
|
| ]
|
|
|
| merge_rank = self.merge_rank
|
| token_to_id = self.token_to_id
|
| unk_id = self.unk_id
|
|
|
|
|
|
|
|
|
|
|
| heap = []
|
|
|
| for i in range(n - 1):
|
|
|
| pair = (
|
| symbols[i],
|
| symbols[i + 1]
|
| )
|
|
|
| rank = merge_rank.get(pair)
|
|
|
| if rank is not None:
|
| heapq.heappush(
|
| heap,
|
| (rank, i)
|
| )
|
|
|
|
|
|
|
|
|
|
|
| previous = [i - 1 for i in range(n)]
|
| following = [i + 1 for i in range(n)]
|
|
|
| following[-1] = -1
|
|
|
| alive = bytearray(b"\x01") * n
|
|
|
|
|
|
|
|
|
|
|
| while heap:
|
|
|
| rank, left = heapq.heappop(heap)
|
|
|
| if not alive[left]:
|
| continue
|
|
|
| right = following[left]
|
|
|
| if right == -1 or not alive[right]:
|
| continue
|
|
|
| pair = (
|
| symbols[left],
|
| symbols[right]
|
| )
|
|
|
| if merge_rank.get(pair) != rank:
|
| continue
|
|
|
|
|
|
|
|
|
|
|
| symbols[left] += symbols[right]
|
|
|
| alive[right] = 0
|
|
|
| next_index = following[right]
|
|
|
| following[left] = next_index
|
|
|
| if next_index != -1:
|
| previous[next_index] = left
|
|
|
|
|
|
|
|
|
|
|
| prev_index = previous[left]
|
|
|
| if prev_index != -1:
|
|
|
| new_pair = (
|
| symbols[prev_index],
|
| symbols[left]
|
| )
|
|
|
| new_rank = merge_rank.get(new_pair)
|
|
|
| if new_rank is not None:
|
| heapq.heappush(
|
| heap,
|
| (
|
| new_rank,
|
| prev_index
|
| )
|
| )
|
|
|
|
|
|
|
|
|
|
|
| if next_index != -1:
|
|
|
| new_pair = (
|
| symbols[left],
|
| symbols[next_index]
|
| )
|
|
|
| new_rank = merge_rank.get(new_pair)
|
|
|
| if new_rank is not None:
|
| heapq.heappush(
|
| heap,
|
| (
|
| new_rank,
|
| left
|
| )
|
| )
|
|
|
|
|
|
|
|
|
|
|
| result = []
|
|
|
| index = 0
|
|
|
| while index != -1:
|
|
|
| if alive[index]:
|
| result.append(
|
| token_to_id.get(
|
| symbols[index],
|
| unk_id
|
| )
|
| )
|
|
|
| index = following[index]
|
|
|
| return result
|
|
|
| def decode(self, ids, skip_special_tokens=False, **kwargs):
|
|
|
| if skip_special_tokens:
|
| ids = [t for t in ids if t not in self.special_token_ids]
|
|
|
| text = "".join(
|
| self.id_to_token.get(
|
| int(idx),
|
| "<|unk|>"
|
| )
|
| for idx in ids
|
| )
|
|
|
| return text.replace("▁", " ") |
