Spaces:
Sleeping
Sleeping
File size: 6,310 Bytes
c407ed8 | 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 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 | from __future__ import annotations
from collections import Counter
from dataclasses import dataclass
from typing import Protocol
class Tokenizer(Protocol):
@property
def vocab_size(self) -> int: ...
def encode(self, text: str) -> list[int]: ...
def decode(self, ids: list[int]) -> str: ...
def id_to_token(self, idx: int) -> str: ...
def to_dict(self) -> dict[str, object]: ...
@dataclass(frozen=True)
class CharTokenizer:
stoi: dict[str, int]
itos: dict[int, str]
@classmethod
def train(cls, text: str) -> "CharTokenizer":
if not text:
raise ValueError("Cannot train a tokenizer on empty text")
chars = sorted(set(text))
stoi = {char: idx for idx, char in enumerate(chars)}
itos = {idx: char for char, idx in stoi.items()}
return cls(stoi=stoi, itos=itos)
@property
def vocab_size(self) -> int:
return len(self.stoi)
def encode(self, text: str) -> list[int]:
try:
return [self.stoi[char] for char in text]
except KeyError as exc:
char = exc.args[0]
raise ValueError(f"Character {char!r} is not in the tokenizer vocabulary") from exc
def decode(self, ids: list[int]) -> str:
try:
return "".join(self.itos[idx] for idx in ids)
except KeyError as exc:
idx = exc.args[0]
raise ValueError(f"Token id {idx!r} is not in the tokenizer vocabulary") from exc
def id_to_token(self, idx: int) -> str:
try:
return self.itos[idx]
except KeyError as exc:
raise ValueError(f"Token id {idx!r} is not in the tokenizer vocabulary") from exc
def to_dict(self) -> dict[str, object]:
return {"type": "char", "stoi": self.stoi, "itos": self.itos}
@classmethod
def from_dict(cls, payload: dict[str, dict[str, int] | dict[int | str, str]]) -> "CharTokenizer":
stoi = {str(char): int(idx) for char, idx in payload["stoi"].items()}
itos = {int(idx): str(char) for idx, char in payload["itos"].items()}
return cls(stoi=stoi, itos=itos)
@dataclass(frozen=True)
class BytePairTokenizer:
stoi: dict[str, int]
itos: dict[int, str]
merges: list[tuple[str, str]]
@classmethod
def train(cls, text: str, vocab_size: int = 256) -> "BytePairTokenizer":
if not text:
raise ValueError("Cannot train a tokenizer on empty text")
if vocab_size <= 0:
raise ValueError("vocab_size must be positive")
base_vocab = sorted(set(text))
sequences = [[char for char in text]]
merges: list[tuple[str, str]] = []
vocab = set(base_vocab)
while len(vocab) < vocab_size:
pair_counts = _count_pairs(sequences)
if not pair_counts:
break
pair, count = pair_counts.most_common(1)[0]
merged = "".join(pair)
if count < 2 or merged in vocab:
break
sequences = [_merge_pair(sequence, pair, merged) for sequence in sequences]
merges.append(pair)
vocab.add(merged)
tokens = sorted(vocab, key=lambda token: (len(token), token))
stoi = {token: idx for idx, token in enumerate(tokens)}
itos = {idx: token for token, idx in stoi.items()}
return cls(stoi=stoi, itos=itos, merges=merges)
@property
def vocab_size(self) -> int:
return len(self.stoi)
def encode(self, text: str) -> list[int]:
unknown = sorted(set(text) - {token for token in self.stoi if len(token) == 1})
if unknown:
raise ValueError(f"Characters {unknown!r} are not in the tokenizer vocabulary")
pieces = list(text)
for left, right in self.merges:
pieces = _merge_pair(pieces, (left, right), left + right)
return [self.stoi[piece] for piece in pieces]
def decode(self, ids: list[int]) -> str:
try:
return "".join(self.itos[idx] for idx in ids)
except KeyError as exc:
idx = exc.args[0]
raise ValueError(f"Token id {idx!r} is not in the tokenizer vocabulary") from exc
def id_to_token(self, idx: int) -> str:
try:
return self.itos[idx]
except KeyError as exc:
raise ValueError(f"Token id {idx!r} is not in the tokenizer vocabulary") from exc
def to_dict(self) -> dict[str, object]:
return {
"type": "bpe",
"stoi": self.stoi,
"itos": self.itos,
"merges": self.merges,
}
@classmethod
def from_dict(cls, payload: dict[str, object]) -> "BytePairTokenizer":
raw_stoi = payload["stoi"]
raw_itos = payload["itos"]
raw_merges = payload["merges"]
if not isinstance(raw_stoi, dict) or not isinstance(raw_itos, dict):
raise ValueError("Invalid BPE tokenizer payload")
if not isinstance(raw_merges, list):
raise ValueError("Invalid BPE merge payload")
stoi = {str(token): int(idx) for token, idx in raw_stoi.items()}
itos = {int(idx): str(token) for idx, token in raw_itos.items()}
merges = [(str(left), str(right)) for left, right in raw_merges]
return cls(stoi=stoi, itos=itos, merges=merges)
def tokenizer_from_dict(payload: dict[str, object]) -> Tokenizer:
tokenizer_type = payload.get("type", "char")
if tokenizer_type == "char":
return CharTokenizer.from_dict(payload) # type: ignore[arg-type]
if tokenizer_type == "bpe":
return BytePairTokenizer.from_dict(payload)
raise ValueError(f"Unknown tokenizer type: {tokenizer_type!r}")
def _count_pairs(sequences: list[list[str]]) -> Counter[tuple[str, str]]:
counts: Counter[tuple[str, str]] = Counter()
for sequence in sequences:
counts.update(zip(sequence, sequence[1:]))
return counts
def _merge_pair(sequence: list[str], pair: tuple[str, str], merged: str) -> list[str]:
out: list[str] = []
idx = 0
while idx < len(sequence):
if idx < len(sequence) - 1 and (sequence[idx], sequence[idx + 1]) == pair:
out.append(merged)
idx += 2
else:
out.append(sequence[idx])
idx += 1
return out
|